Cache more information for Closed Classes and Unions on NamedTypeSymbol

src/Compilers/CSharp/Portable/Symbols/NamedTypeSymbol.cs

@@ -46,16 +46,31 @@ private class UnionData
             public StrongBox<NullableFlowState>? _lazyValueDeclaredNullableFlowState;
             public Symbol? _lazyHasValueProperty = ErrorTypeSymbol.UnknownResultType;
             public ImmutableArray<MethodSymbol> _lazyTryGetValueMethods;
+            public ImmutableArray<TypeUnionValueSet.CaseInfo> _lazyTypeUnionValueSetCases;
         }
 
         private sealed class UnionDataForDefinition : UnionData
         {
             public NamedTypeSymbol? _lazyMemberProviderInterface = ErrorTypeSymbol.UnknownResultType;
         }
 
+        private sealed class ClosedClassData
+        {
+            /// <summary>
+            /// This field is safe to read only after checking that <see cref="_lazySubtypes"/> has been initialized.
+            /// Reading of <see cref="_lazySubtypes"/> for this purpose should be done
+            /// using <see cref="RoslynImmutableInterlocked.VolatileRead{T}(ref readonly ImmutableArray{T})"/> API
+            /// to enforce order of read operations between the fields.
+            /// </summary>
+            public bool _lazySubtypesIsComplete;
+            public ImmutableArray<NamedTypeSymbol> _lazySubtypes;
+            public ImmutableArray<TypeUnionValueSet.CaseInfo> _lazyTypeUnionValueSetCases;
+        }
+
         private sealed partial class UncommonProperties
         {
             public UnionData? _lazyUnionData;
+            public ClosedClassData? _lazyClosedClassData;
         }
 #nullable disable
 
@@ -745,22 +760,37 @@ internal bool TryGetClosedSubtypes(out ImmutableArray<NamedTypeSymbol> subtypes,
                 return false;
             }
 
-            var candidateSubtypes = CandidateClosedSubtypeDefinitions;
-            if (!IsGenericType && candidateSubtypes.All(subtype => !subtype.IsGenericType))
+            ClosedClassData lazyClosedClassData = GetClosedClassData();
+            ImmutableArray<NamedTypeSymbol> lazySubtypes = RoslynImmutableInterlocked.VolatileRead(in lazyClosedClassData._lazySubtypes);
+            if (!lazySubtypes.IsDefault)
             {
-                subtypes = candidateSubtypes;
-                return true;
+                subtypes = lazySubtypes;
+                return lazyClosedClassData._lazySubtypesIsComplete;
             }
 
-            var resultBuilder = ArrayBuilder<NamedTypeSymbol>.GetInstance(candidateSubtypes.Length);
-            var baseTypeTypeParameters = PooledHashSet<TypeParameterSymbol>.GetInstance();
-            this.FindTypeParameters(baseTypeTypeParameters);
+            (lazyClosedClassData._lazySubtypesIsComplete, lazySubtypes) = calculateClosedSubtypes(cancellationToken);
+            ImmutableInterlocked.InterlockedInitialize(ref lazyClosedClassData._lazySubtypes, lazySubtypes);
+
+            subtypes = lazyClosedClassData._lazySubtypes;
+            return lazyClosedClassData._lazySubtypesIsComplete;
 
-            var success = tryGetSpeakableSubtypes(this, candidateSubtypes, resultBuilder, baseTypeTypeParameters, cancellationToken);
-            baseTypeTypeParameters.Free();
+            (bool, ImmutableArray<NamedTypeSymbol>) calculateClosedSubtypes(CancellationToken cancellationToken)
+            {
+                var candidateSubtypes = CandidateClosedSubtypeDefinitions;
+                if (!IsGenericType && candidateSubtypes.All(subtype => !subtype.IsGenericType))
+                {
+                    return (true, candidateSubtypes);
+                }
 
-            subtypes = resultBuilder.ToImmutableAndFree();
-            return success;
+                var resultBuilder = ArrayBuilder<NamedTypeSymbol>.GetInstance(candidateSubtypes.Length);
+                var baseTypeTypeParameters = PooledHashSet<TypeParameterSymbol>.GetInstance();
+                this.FindTypeParameters(baseTypeTypeParameters);
+
+                var success = tryGetSpeakableSubtypes(this, candidateSubtypes, resultBuilder, baseTypeTypeParameters, cancellationToken);
+                baseTypeTypeParameters.Free();
+
+                return (success, resultBuilder.ToImmutableAndFree());
+            }
 
             static bool tryGetSpeakableSubtypes(NamedTypeSymbol @this, ImmutableArray<NamedTypeSymbol> candidateSubtypes, ArrayBuilder<NamedTypeSymbol> resultBuilder, HashSet<TypeParameterSymbol> baseTypeTypeParameters, CancellationToken cancellationToken)
             {
@@ -786,6 +816,24 @@ static bool tryGetSpeakableSubtypes(NamedTypeSymbol @this, ImmutableArray<NamedT
             }
         }
 
+        internal ImmutableArray<TypeUnionValueSet.CaseInfo> ClosedClassTypeUnionValueSetCases()
+        {
+            ClosedClassData lazyClosedClassData = GetClosedClassData();
+            ImmutableArray<TypeUnionValueSet.CaseInfo> lazyTypeUnionValueSetCases = lazyClosedClassData._lazyTypeUnionValueSetCases;
+            if (!lazyTypeUnionValueSetCases.IsDefault)
+            {
+                return lazyTypeUnionValueSetCases;
+            }
+
+            var builder = ArrayBuilder<TypeUnionValueSet.CaseInfo>.GetInstance();
+            var setBuilder = AllIgnoreOptionsSetPool.Allocate();
+            ValueSetFactory.ClosedClassTypeUnionValueSetFactory.ExpandClosedSubtypes(this, builder, setBuilder);
+            setBuilder.Free();
+
+            ImmutableInterlocked.InterlockedInitialize(ref lazyClosedClassData._lazyTypeUnionValueSetCases, builder.ToImmutableAndFree());
+            return lazyClosedClassData._lazyTypeUnionValueSetCases;
+        }
+
         /// <summary>
         /// Gets the set of subtype definitions in the same module whose base type has same original definition as 'this'.
         /// </summary>
@@ -1973,6 +2021,28 @@ void addUseSiteInfoForCachedResult(ref CompoundUseSiteInfo<AssemblySymbol> useSi
             }
         }
 
+        internal ImmutableArray<TypeUnionValueSet.CaseInfo> UnionTypeUnionValueSetCases()
+        {
+            UnionData lazyUnionData = GetUnionData();
+            ImmutableArray<TypeUnionValueSet.CaseInfo> lazyTypeUnionValueSetCases = lazyUnionData._lazyTypeUnionValueSetCases;
+            if (!lazyTypeUnionValueSetCases.IsDefault)
+            {
+                return lazyTypeUnionValueSetCases;
+            }
+
+            var builder = ArrayBuilder<TypeUnionValueSet.CaseInfo>.GetInstance();
+            var setBuilder = AllIgnoreOptionsSetPool.Allocate();
+            foreach (var caseType in UnionCaseTypesNoUseSiteDiagnostics)
+            {
+                ValueSetFactory.ClosedClassTypeUnionValueSetFactory.ExpandClosedSubtypes(caseType.StrippedType(), builder, setBuilder);
+            }
+
+            setBuilder.Free();
+
+            ImmutableInterlocked.InterlockedInitialize(ref lazyUnionData._lazyTypeUnionValueSetCases, builder.ToImmutableAndFree());
+            return lazyUnionData._lazyTypeUnionValueSetCases;
+        }
+
         internal ImmutableArray<MethodSymbol> UnionFactoryMethods(ref CompoundUseSiteInfo<AssemblySymbol> membersInterfaceForDefinitionInterfacesUseSiteInfo)
         {
             Debug.Assert(IsUnionType);
@@ -2235,6 +2305,22 @@ private UnionData GetUnionData()
             return lazyUnionData;
         }
 
+        private ClosedClassData GetClosedClassData()
+        {
+            UncommonProperties lazyUncommonProperties = GetUncommonProperties();
+
+            ClosedClassData? lazyClosedClassData = lazyUncommonProperties._lazyClosedClassData;
+            if (lazyClosedClassData is null)
+            {
+                Debug.Assert(IsClosed);
+                Interlocked.CompareExchange(ref lazyUncommonProperties._lazyClosedClassData, new ClosedClassData(), null);
+                Debug.Assert(lazyUncommonProperties._lazyClosedClassData is not null);
+                return lazyUncommonProperties._lazyClosedClassData;
+            }
+
+            return lazyClosedClassData;
+        }
+
         internal static bool IsSuitableUnionConstructor(MethodSymbol ctor)
         {
             Debug.Assert(ctor.MethodKind is MethodKind.Constructor);

src/Compilers/CSharp/Portable/Utilities/TypeUnionValueSet.cs

@@ -188,9 +188,9 @@ public bool IsEmpty(ref CompoundUseSiteInfo<AssemblySymbol> useSiteInfo)
                 && effectiveClosedBase.TryGetClosedSubtypes(out var derivedTypes)
                 && derivedTypes is not [])
             {
-                foreach (var derivedType in derivedTypes)
+                foreach (var caseInfo in effectiveClosedBase.ClosedClassTypeUnionValueSetCases())
                 {
-                    if (closedDerivedTypeMayBePresent(node, derivedType, ref useSiteInfo))
+                    if (evaluateCore(node, caseInfo.CaseType, ref useSiteInfo) != false)
                         return null;
                 }
 
@@ -283,25 +283,6 @@ public bool IsEmpty(ref CompoundUseSiteInfo<AssemblySymbol> useSiteInfo)
 
                 return null;
             }
-
-            bool closedDerivedTypeMayBePresent(Node root, NamedTypeSymbol derivedType, ref CompoundUseSiteInfo<AssemblySymbol> useSiteInfo)
-            {
-                Debug.Assert(derivedType.BaseTypeNoUseSiteDiagnostics.IsClosed);
-                if (derivedType.TryGetClosedSubtypes(out var nestedDerivedTypes) && nestedDerivedTypes is not [])
-                {
-                    foreach (var nestedDerivedType in nestedDerivedTypes)
-                    {
-                        if (closedDerivedTypeMayBePresent(root, nestedDerivedType, ref useSiteInfo))
-                            return true;
-                    }
-                }
-                else if (evaluateCore(root, derivedType, ref useSiteInfo) != false)
-                {
-                    return true;
-                }
-
-                return false;
-            }
         }
 
         public TypeSymbol? SampleType(Binder binder, ref CompoundUseSiteInfo<AssemblySymbol> useSiteInfo)

src/Compilers/CSharp/Portable/Utilities/ValueSetFactory.ClosedClassTypeUnionValueSetFactory.cs

@@ -12,7 +12,7 @@ namespace Microsoft.CodeAnalysis.CSharp
 {
     internal static partial class ValueSetFactory
     {
-        private sealed class ClosedClassTypeUnionValueSetFactory : ITypeUnionValueSetFactory
+        internal sealed class ClosedClassTypeUnionValueSetFactory : ITypeUnionValueSetFactory
         {
             private readonly TypeSymbol _closedClassOrTypeParameter;
 
@@ -61,11 +61,12 @@ private static void AddCaseInfo(ArrayBuilder<TypeUnionValueSet.CaseInfo> builder
 
             private ImmutableArray<TypeUnionValueSet.CaseInfo> ClosedSubtypes()
             {
-                var builder = ArrayBuilder<TypeUnionValueSet.CaseInfo>.GetInstance();
-                var setBuilder = TypeSymbol.AllIgnoreOptionsSetPool.Allocate();
-                ExpandClosedSubtypes(_closedClassOrTypeParameter, builder, setBuilder);
-                setBuilder.Free();
-                return builder.ToImmutableAndFree();
+                if (_closedClassOrTypeParameter is NamedTypeSymbol namedType)
+                {
+                    return namedType.ClosedClassTypeUnionValueSetCases();
+                }
+
+                return [new TypeUnionValueSet.CaseInfo((TypeParameterSymbol)_closedClassOrTypeParameter, originalClosedBase: null)];
             }
 
             public TypeUnionValueSet AllValues(ConversionsBase conversions)

src/Compilers/CSharp/Portable/Utilities/ValueSetFactory.UnionTypeTypeUnionValueSetFactory.cs

@@ -23,15 +23,7 @@ public UnionTypeTypeUnionValueSetFactory(NamedTypeSymbol unionType)
 
             private ImmutableArray<TypeUnionValueSet.CaseInfo> AdjustedTypesInUnion()
             {
-                var builder = ArrayBuilder<TypeUnionValueSet.CaseInfo>.GetInstance();
-                var setBuilder = TypeSymbol.AllIgnoreOptionsSetPool.Allocate();
-                foreach (var caseType in _unionType.UnionCaseTypesNoUseSiteDiagnostics)
-                {
-                    ClosedClassTypeUnionValueSetFactory.ExpandClosedSubtypes(caseType.StrippedType(), builder, setBuilder);
-                }
-
-                setBuilder.Free();
-                return builder.ToImmutableAndFree();
+                return _unionType.UnionTypeUnionValueSetCases();
             }
 
             public TypeUnionValueSet AllValues(ConversionsBase conversions)

src/Compilers/CSharp/Test/CSharp15/ClosedClassesTests.cs

@@ -7732,4 +7732,77 @@ static partial class C3;
         Assert.True(c3.IsClosed);
         Assert.True(c3.IsStatic);
     }
+
+    [Fact]
+    public void ClosedUnion_01()
+    {
+        var src = @"
+[System.Runtime.CompilerServices.Union]
+closed class C1
+{
+    protected readonly object _value;
+    public C1(int x) { _value = x; }
+    public C1(string x) { _value = x; }
+    public object Value => _value;
+}
+
+class C2() : C1(1)
+{
+}
+
+class C3() : C1("""")
+{
+}
+
+class Program
+{
+    static int Test1(C1 u)
+    {
+#line 100
+        return u switch 
+        {
+            int => 2
+        };
+    }   
+
+    static int Test2(C1 u)
+    {
+#line 200
+        return u switch 
+        {
+            int => 2,
+            string => 3, 
+        };
+    }   
+
+    static int Test3(C1 u)
+    {
+#line 300
+        return u switch 
+        {
+            C2 => 1,
+        };
+    }   
+
+    static int Test4(C1 u)
+    {
+#line 400
+        return u switch 
+        {
+            C2 => 1,
+            C3 => 2,
+        };
+    }   
+}
+";
+        var comp = CreateCompilation([src, UnionAttributeSource, IsClosedTypeAttributeDefinition, CompilerFeatureRequiredAttribute]);
+        comp.VerifyDiagnostics(
+            // (100,18): warning CS8509: The switch expression does not handle all possible values of its input type (it is not exhaustive). For example, the pattern 'string' is not covered.
+            //         return u switch 
+            Diagnostic(ErrorCode.WRN_SwitchExpressionNotExhaustive, "switch").WithArguments("string").WithLocation(100, 18),
+            // (300,18): warning CS8509: The switch expression does not handle all possible values of its input type (it is not exhaustive). For example, the pattern 'C3' is not covered.
+            //         return u switch 
+            Diagnostic(ErrorCode.WRN_SwitchExpressionNotExhaustive, "switch").WithArguments("C3").WithLocation(300, 18)
+            );
+    }
 }