Support more collections in STJ source generator

src/libraries/System.Text.Json/Common/ReflectionExtensions.cs

@@ -3,11 +3,227 @@
 
 using System.Diagnostics;
 using System.Reflection;
+#if !BUILDING_SOURCE_GENERATOR
+using System.Diagnostics.CodeAnalysis;
+#endif
 
 namespace System.Text.Json.Reflection
 {
     internal static partial class ReflectionExtensions
     {
+        // Immutable collection types.
+        private const string ImmutableArrayGenericTypeName = "System.Collections.Immutable.ImmutableArray`1";
+        private const string ImmutableListGenericTypeName = "System.Collections.Immutable.ImmutableList`1";
+        private const string ImmutableListGenericInterfaceTypeName = "System.Collections.Immutable.IImmutableList`1";
+        private const string ImmutableStackGenericTypeName = "System.Collections.Immutable.ImmutableStack`1";
+        private const string ImmutableStackGenericInterfaceTypeName = "System.Collections.Immutable.IImmutableStack`1";
+        private const string ImmutableQueueGenericTypeName = "System.Collections.Immutable.ImmutableQueue`1";
+        private const string ImmutableQueueGenericInterfaceTypeName = "System.Collections.Immutable.IImmutableQueue`1";
+        private const string ImmutableSortedSetGenericTypeName = "System.Collections.Immutable.ImmutableSortedSet`1";
+        private const string ImmutableHashSetGenericTypeName = "System.Collections.Immutable.ImmutableHashSet`1";
+        private const string ImmutableSetGenericInterfaceTypeName = "System.Collections.Immutable.IImmutableSet`1";
+        private const string ImmutableDictionaryGenericTypeName = "System.Collections.Immutable.ImmutableDictionary`2";
+        private const string ImmutableDictionaryGenericInterfaceTypeName = "System.Collections.Immutable.IImmutableDictionary`2";
+        private const string ImmutableSortedDictionaryGenericTypeName = "System.Collections.Immutable.ImmutableSortedDictionary`2";
+
+        // Immutable collection builder types.
+        private const string ImmutableArrayTypeName = "System.Collections.Immutable.ImmutableArray";
+        private const string ImmutableListTypeName = "System.Collections.Immutable.ImmutableList";
+        private const string ImmutableStackTypeName = "System.Collections.Immutable.ImmutableStack";
+        private const string ImmutableQueueTypeName = "System.Collections.Immutable.ImmutableQueue";
+        private const string ImmutableSortedSetTypeName = "System.Collections.Immutable.ImmutableSortedSet";
+        private const string ImmutableHashSetTypeName = "System.Collections.Immutable.ImmutableHashSet";
+        private const string ImmutableDictionaryTypeName = "System.Collections.Immutable.ImmutableDictionary";
+        private const string ImmutableSortedDictionaryTypeName = "System.Collections.Immutable.ImmutableSortedDictionary";
+
+        public const string CreateRangeMethodName = "CreateRange";
+
+        public static Type? GetCompatibleGenericBaseClass(
+            this Type type,
+            Type baseType,
+            Type? objectType = null,
+            bool sourceGenType = false)
+        {
+            Debug.Assert(baseType.IsGenericType);
+            Debug.Assert(!baseType.IsInterface);
+            Debug.Assert(baseType == baseType.GetGenericTypeDefinition());
+
+            // Work around not being able to use typeof(object) directly during compile-time src gen type analysis.
+            objectType ??= typeof(object);
+
+            Type? baseTypeToCheck = type;
+
+            while (baseTypeToCheck != null && baseTypeToCheck != typeof(object))
+            {
+                if (baseTypeToCheck.IsGenericType)
+                {
+                    Type genericTypeToCheck = baseTypeToCheck.GetGenericTypeDefinition();
+                    if (genericTypeToCheck == baseType ||
+                        (sourceGenType && (OpenGenericTypesHaveSamePrefix(baseType, genericTypeToCheck))))
+                    {
+                        return baseTypeToCheck;
+                    }
+                }
+
+                baseTypeToCheck = baseTypeToCheck.BaseType;
+            }
+
+            return null;
+        }
+
+#if !BUILDING_SOURCE_GENERATOR
+        [UnconditionalSuppressMessage("ReflectionAnalysis", "IL2070:UnrecognizedReflectionPattern",
+            Justification = "The 'interfaceType' must exist and so trimmer kept it. In which case " +
+                "It also kept it on any type which implements it. The below call to GetInterfaces " +
+                "may return fewer results when trimmed but it will return the 'interfaceType' " +
+                "if the type implemented it, even after trimming.")]
+#endif
+        public static Type? GetCompatibleGenericInterface(this Type type, Type interfaceType)
+        {
+            Debug.Assert(interfaceType.IsGenericType);
+            Debug.Assert(interfaceType.IsInterface);
+            Debug.Assert(interfaceType == interfaceType.GetGenericTypeDefinition());
+
+            Type interfaceToCheck = type;
+
+            if (interfaceToCheck.IsGenericType)
+            {
+                interfaceToCheck = interfaceToCheck.GetGenericTypeDefinition();
+            }
+
+            if (interfaceToCheck == interfaceType)
+            {
+                return type;
+            }
+
+            foreach (Type typeToCheck in type.GetInterfaces())
+            {
+                if (typeToCheck.IsGenericType)
+                {
+                    Type genericInterfaceToCheck = typeToCheck.GetGenericTypeDefinition();
+                    if (genericInterfaceToCheck == interfaceType)
+                    {
+                        return typeToCheck;
+                    }
+                }
+            }
+
+            return null;
+        }
+
+        public static bool IsImmutableDictionaryType(this Type type, bool sourceGenType = false)
+        {
+            if (!type.IsGenericType || !type.Assembly.FullName!.StartsWith("System.Collections.Immutable", StringComparison.Ordinal))
+            {
+                return false;
+            }
+
+            switch (GetBaseNameFromGenericType(type, sourceGenType))
+            {
+                case ImmutableDictionaryGenericTypeName:
+                case ImmutableDictionaryGenericInterfaceTypeName:
+                case ImmutableSortedDictionaryGenericTypeName:
+                    return true;
+                default:
+                    return false;
+            }
+        }
+
+        public static bool IsImmutableEnumerableType(this Type type, bool sourceGenType = false)
+        {
+            if (!type.IsGenericType || !type.Assembly.FullName!.StartsWith("System.Collections.Immutable", StringComparison.Ordinal))
+            {
+                return false;
+            }
+
+            switch (GetBaseNameFromGenericType(type, sourceGenType))
+            {
+                case ImmutableArrayGenericTypeName:
+                case ImmutableListGenericTypeName:
+                case ImmutableListGenericInterfaceTypeName:
+                case ImmutableStackGenericTypeName:
+                case ImmutableStackGenericInterfaceTypeName:
+                case ImmutableQueueGenericTypeName:
+                case ImmutableQueueGenericInterfaceTypeName:
+                case ImmutableSortedSetGenericTypeName:
+                case ImmutableHashSetGenericTypeName:
+                case ImmutableSetGenericInterfaceTypeName:
+                    return true;
+                default:
+                    return false;
+            }
+        }
+
+        public static string? GetImmutableDictionaryConstructingTypeName(this Type type, bool sourceGenType = false)
+        {
+            Debug.Assert(type.IsImmutableDictionaryType(sourceGenType));
+
+            // Use the generic type definition of the immutable collection to determine
+            // an appropriate constructing type, i.e. a type that we can invoke the
+            // `CreateRange<T>` method on, which returns the desired immutable collection.
+            switch (GetBaseNameFromGenericType(type, sourceGenType))
+            {
+                case ImmutableDictionaryGenericTypeName:
+                case ImmutableDictionaryGenericInterfaceTypeName:
+                    return ImmutableDictionaryTypeName;
+                case ImmutableSortedDictionaryGenericTypeName:
+                    return ImmutableSortedDictionaryTypeName;
+                default:
+                    // We verified that the type is an immutable collection, so the
+                    // generic definition is one of the above.
+                    return null;
+            }
+        }
+
+        public static string? GetImmutableEnumerableConstructingTypeName(this Type type, bool sourceGenType = false)
+        {
+            Debug.Assert(type.IsImmutableEnumerableType(sourceGenType));
+
+            // Use the generic type definition of the immutable collection to determine
+            // an appropriate constructing type, i.e. a type that we can invoke the
+            // `CreateRange<T>` method on, which returns the desired immutable collection.
+            switch (GetBaseNameFromGenericType(type, sourceGenType))
+            {
+                case ImmutableArrayGenericTypeName:
+                    return ImmutableArrayTypeName;
+                case ImmutableListGenericTypeName:
+                case ImmutableListGenericInterfaceTypeName:
+                    return ImmutableListTypeName;
+                case ImmutableStackGenericTypeName:
+                case ImmutableStackGenericInterfaceTypeName:
+                    return ImmutableStackTypeName;
+                case ImmutableQueueGenericTypeName:
+                case ImmutableQueueGenericInterfaceTypeName:
+                    return ImmutableQueueTypeName;
+                case ImmutableSortedSetGenericTypeName:
+                    return ImmutableSortedSetTypeName;
+                case ImmutableHashSetGenericTypeName:
+                case ImmutableSetGenericInterfaceTypeName:
+                    return ImmutableHashSetTypeName;
+                default:
+                    // We verified that the type is an immutable collection, so the
+                    // generic definition is one of the above.
+                    return null;
+            }
+        }
+
+        private static bool OpenGenericTypesHaveSamePrefix(Type t1, Type t2)
+            => t1.FullName == GetBaseNameFromGenericTypeDef(t2);
+
+        private static string GetBaseNameFromGenericType(Type genericType, bool sourceGenType)
+        {
+            Type genericTypeDef = genericType.GetGenericTypeDefinition();
+            return sourceGenType ? GetBaseNameFromGenericTypeDef(genericTypeDef) : genericTypeDef.FullName!;
+        }
+
+        private static string GetBaseNameFromGenericTypeDef(Type genericTypeDef)
+        {
+            Debug.Assert(genericTypeDef.IsGenericType);
+            string fullName = genericTypeDef.FullName!;
+            int length = fullName.IndexOf("`") + 2;
+            return fullName.Substring(0, length);
+        }
+
         public static bool IsVirtual(this PropertyInfo? propertyInfo)
         {
             Debug.Assert(propertyInfo != null);

src/libraries/System.Text.Json/gen/CollectionType.cs

@@ -9,11 +9,28 @@ namespace System.Text.Json.SourceGeneration
 {
     internal enum CollectionType
     {
-        NotApplicable = 0,
-        Array = 1,
-        List = 2,
-        IEnumerable = 3,
-        IList = 4,
-        Dictionary = 5
+        NotApplicable,
+        // Dictionary types
+        IDictionary,
+        Dictionary,
+        ImmutableDictionary,
+        IDictionaryOfTKeyTValue,
+        IReadOnlyDictionary,
+        // Non-dictionary types
+        Array,
+        List,
+        IEnumerable,
+        IList,
+        IListOfT,
+        ISet,
+        ICollectionOfT,
+        StackOfT,
+        QueueOfT,
+        ConcurrentStack,
+        ConcurrentQueue,
+        IEnumerableOfT,
+        Stack,
+        Queue,
+        ImmutableEnumerable
     }
 }

src/libraries/System.Text.Json/gen/ContextGenerationSpec.cs

@@ -4,13 +4,15 @@
 using System.Collections.Generic;
 using System.Text.Json.Serialization;
 using System.Text.Json.Reflection;
+using System.Diagnostics;
 
 namespace System.Text.Json.SourceGeneration
 {
     /// <summary>
     /// Represents the set of input types and options needed to provide an
     /// implementation for a user-provided JsonSerializerContext-derived type.
     /// </summary>
+    [DebuggerDisplay("ContextTypeRef={ContextTypeRef}")]
     internal sealed class ContextGenerationSpec
     {
         public JsonSourceGenerationOptionsAttribute GenerationOptions { get; init; }
@@ -34,6 +36,6 @@ internal sealed class ContextGenerationSpec
         /// </summary>
         public HashSet<string> RuntimePropertyNames { get; } = new();
 
-        public string ContextTypeRef => $"global::{ContextType.GetUniqueCompilableTypeName()}";
+        public string ContextTypeRef => ContextType.GetCompilableName();
     }
 }