Generate a global shape provider per compilation. (#57)

eiriktsarpalis · Nov 15, 2024 · eabb24a · eabb24a
1 parent 20dc8b2
commit eabb24a
Show file tree

Hide file tree

Showing 31 changed files with 244 additions and 303 deletions.
diff --git a/src/PolyType.SourceGenerator/Helpers/RoslynHelpers.IncrementalValues.cs b/src/PolyType.SourceGenerator/Helpers/RoslynHelpers.IncrementalValues.cs
@@ -17,26 +17,25 @@ internal static partial class RoslynHelpers
     /// <summary>
     /// Replacement for <see cref="SyntaxValueProvider.ForAttributeWithMetadataName" /> that handles generic attributes correctly.
     /// </summary>
-    public static IncrementalValuesProvider<TypeWithAttributeDeclarationContext> ForTypesWithAttributeDeclaration(
-        this SyntaxValueProvider provider, string attributeFullyQualifiedName,
+    public static IncrementalValuesProvider<TypeWithAttributeDeclarationContext> ForTypesWithAttributeDeclarations(
+        this SyntaxValueProvider provider,
+        string[] attributeFullyQualifiedNames,
         Func<BaseTypeDeclarationSyntax, CancellationToken, bool> predicate)
     {
-        NameSyntax attributeNameSyntax = SyntaxFactory.ParseName(attributeFullyQualifiedName);
-        string attributeName = GetTypeName(attributeNameSyntax, out int attributeArity);
-        ImmutableArray<string> attributeNamespace = GetNamespaceTokens(attributeNameSyntax);
-        string? attributeNameMinusSuffix = attributeName.EndsWith("Attribute", StringComparison.Ordinal) ? attributeName[..^"Attribute".Length] : null;
+        Debug.Assert(attributeFullyQualifiedNames.Length is > 0 and <= 3, "Does not optimize for large lists of attributes.");
+        ParseAttributeFullyQualifiedNames(attributeFullyQualifiedNames, out var attributeData, out var attributeSyntaxNodeCandidates);
 
         return provider.CreateSyntaxProvider(
             predicate: (SyntaxNode node, CancellationToken token) => node is BaseTypeDeclarationSyntax typeDecl && IsAnnotatedTypeDeclaration(typeDecl, token),
             transform: Transform)
             .Where(ctx => ctx.Type != null)
             .GroupBy(
-                keySelector: value => value.Type, 
-                resultSelector: static (key, values) => 
-                    new TypeWithAttributeDeclarationContext 
-                    { 
-                        TypeSymbol = (ITypeSymbol)key!, 
-                        Declarations = values.Select(v => (v.Syntax, v.Model)).ToImmutableArray() 
+                keySelector: value => value.Type,
+                resultSelector: static (key, values) =>
+                    new TypeWithAttributeDeclarationContext
+                    {
+                        TypeSymbol = (ITypeSymbol)key!,
+                        Declarations = values.Select(v => (v.Syntax, v.Model)).ToImmutableArray()
                     },
 
                 keyComparer: SymbolEqualityComparer.Default);
@@ -48,9 +47,12 @@ bool IsAnnotatedTypeDeclaration(BaseTypeDeclarationSyntax typeDecl, Cancellation
                 foreach (AttributeSyntax attribute in attributeList.Attributes)
                 {
                     string name = GetTypeName(attribute.Name, out int arity);
-                    if ((name == attributeName || name == attributeNameMinusSuffix) && arity == attributeArity)
+                    foreach (var candidate in attributeSyntaxNodeCandidates)
                     {
-                        return predicate(typeDecl, token);
+                        if (candidate.name == name && candidate.arity == arity)
+                        {
+                            return predicate(typeDecl, token);
+                        }
                     }
                 }
             }
@@ -65,12 +67,17 @@ bool IsAnnotatedTypeDeclaration(BaseTypeDeclarationSyntax typeDecl, Cancellation
 
             foreach (AttributeData attrData in typeSymbol.GetAttributes())
             {
-                if (attrData.AttributeClass is INamedTypeSymbol attributeType &&
-                    attributeType.Name == attributeName &&
-                    attributeType.Arity == attributeArity &&
-                    attributeType.ContainingNamespace.MatchesNamespace(attributeNamespace))
+                if (attrData.AttributeClass is INamedTypeSymbol attributeType)
                 {
-                    return (typeSymbol, typeDecl, ctx.SemanticModel);
+                    foreach (var (attributeNamespace, attributeName, attributeArity) in attributeData)
+                    {
+                        if (attributeType.Name == attributeName &&
+                            attributeType.Arity == attributeArity &&
+                            attributeType.ContainingNamespace.MatchesNamespace(attributeNamespace))
+                        {
+                            return (typeSymbol, typeDecl, ctx.SemanticModel);
+                        }
+                    }
                 }
             }
 
@@ -130,6 +137,32 @@ void Traverse(NameSyntax current)
                 }
             }
         }
+
+        static void ParseAttributeFullyQualifiedNames(
+            string[] attributeFullyQualifiedNames,
+            out (ImmutableArray<string> Namespace, string Name, int Arity)[] attributeData,
+            out (string name, int arity)[] attributeSyntaxNodeCandidates)
+        {
+            attributeData = new (ImmutableArray<string> Namespace, string Name, int Arity)[attributeFullyQualifiedNames.Length];
+            List<(string name, int arity)> attributeSyntaxNodeCandidateList = new();
+            int i = 0;
+
+            foreach (string attributeFqn in attributeFullyQualifiedNames)
+            {
+                NameSyntax attributeNameSyntax = SyntaxFactory.ParseName(attributeFqn);
+                string attributeName = GetTypeName(attributeNameSyntax, out int attributeArity);
+
+                attributeSyntaxNodeCandidateList.Add((attributeName, attributeArity));
+                if (attributeName.EndsWith("Attribute", StringComparison.Ordinal))
+                {
+                    attributeSyntaxNodeCandidateList.Add((attributeName[..^"Attribute".Length], attributeArity));
+                }
+
+                attributeData[i++] = (GetNamespaceTokens(attributeNameSyntax), attributeName, attributeArity);
+            }
+
+            attributeSyntaxNodeCandidates = attributeSyntaxNodeCandidateList.ToArray();
+        }
     }
 
     // Cf. https://github.com/dotnet/roslyn/issues/72667

diff --git a/src/PolyType.SourceGenerator/Model/TypeDeclarationModel.cs b/src/PolyType.SourceGenerator/Model/TypeDeclarationModel.cs
@@ -10,5 +10,6 @@ public record TypeDeclarationModel
     public required ImmutableEquatableArray<string> ContainingTypes { get; init; }
     public required string SourceFilenamePrefix { get; init; }
     public required string? Namespace { get; init; }
-    public required bool IsValidTypeDeclaration { get; init; }
+    public required bool ImplementsITypeShapeProvider { get; init; }
+    public required ImmutableEquatableSet<TypeId> ShapeableOfTImplementations { get; init; }
 }
diff --git a/src/PolyType.SourceGenerator/Model/TypeId.cs b/src/PolyType.SourceGenerator/Model/TypeId.cs
@@ -8,7 +8,7 @@ namespace PolyType.SourceGenerator.Model;
 public readonly struct TypeId : IEquatable<TypeId>
 {
     public required string FullyQualifiedName { get; init; }
-    public required string GeneratedPropertyName { get; init; }
+    public required string TypeIdentifier { get; init; }
     public required bool IsValueType { get; init; }
     public required SpecialType SpecialType { get; init; }
 

diff --git a/src/PolyType.SourceGenerator/Model/TypeShapeModel.cs b/src/PolyType.SourceGenerator/Model/TypeShapeModel.cs
@@ -5,5 +5,4 @@ namespace PolyType.SourceGenerator.Model;
 public abstract record TypeShapeModel
 {
     public required TypeId Type { get; init; }
-    public required bool EmitGenericTypeShapeProviderImplementation { get; init; }
 }
diff --git a/src/PolyType.SourceGenerator/Model/TypeShapeProviderModel.cs b/src/PolyType.SourceGenerator/Model/TypeShapeProviderModel.cs
@@ -4,8 +4,8 @@ namespace PolyType.SourceGenerator.Model;
 
 public sealed record TypeShapeProviderModel
 {
-    public required TypeDeclarationModel Declaration { get; init; }
+    public required TypeDeclarationModel ProviderDeclaration { get; init; }
     public required ImmutableEquatableDictionary<TypeId, TypeShapeModel> ProvidedTypes { get; init; }
     public required ImmutableEquatableSet<EquatableDiagnostic> Diagnostics { get; init; }
-    public required ImmutableEquatableArray<TypeDeclarationModel> GenerateShapeTypes { get; init; }
+    public required ImmutableEquatableArray<TypeDeclarationModel> AnnotatedTypes { get; init; }
 }
diff --git a/src/PolyType.SourceGenerator/Parser/Parser.ModelMapper.cs b/src/PolyType.SourceGenerator/Parser/Parser.ModelMapper.cs
@@ -7,23 +7,20 @@ namespace PolyType.SourceGenerator;
 
 public sealed partial class Parser
 {
-    private TypeShapeModel MapModel(TypeId typeId, TypeDataModel model, bool isGeneratedViaWitnessType)
+    private TypeShapeModel MapModel(TypeId typeId, TypeDataModel model)
     {
-        bool emitGenericProviderImplementation = isGeneratedViaWitnessType && model.IsRootType;
         return model switch
         {
             EnumDataModel enumModel => new EnumShapeModel
             {
                 Type = typeId,
                 UnderlyingType = CreateTypeId(enumModel.UnderlyingType),
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
             },
 
             NullableDataModel nullableModel => new NullableShapeModel
             {
                 Type = typeId,
                 ElementType = CreateTypeId(nullableModel.ElementType),
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
             },
 
             EnumerableDataModel enumerableModel => new EnumerableShapeModel
@@ -61,7 +58,6 @@ enumerableModel.ConstructionStrategy is CollectionModelConstructionStrategy.List
 
                 Kind = enumerableModel.EnumerableKind,
                 Rank = enumerableModel.Rank,
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
                 ElementTypeContainsNullableAnnotations = enumerableModel.ElementType.ContainsNullabilityAnnotations(),
             },
 
@@ -94,7 +90,6 @@ dictionaryModel.ConstructionStrategy is CollectionModelConstructionStrategy.Muta
                 StaticFactoryMethod = dictionaryModel.FactoryMethod is { IsStatic: true } m ? m.GetFullyQualifiedName() : null,
                 IsTupleEnumerableFactory = dictionaryModel.ConstructionStrategy is CollectionModelConstructionStrategy.TupleEnumerable,
                 Kind = dictionaryModel.DictionaryKind,
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
                 CtorRequiresDictionaryConversion =
                     dictionaryModel.ConstructionStrategy is CollectionModelConstructionStrategy.Dictionary &&
                     !IsFactoryAcceptingIEnumerable(dictionaryModel.FactoryMethod),
@@ -118,7 +113,6 @@ dictionaryModel.ConstructionStrategy is CollectionModelConstructionStrategy.Dict
                 IsValueTupleType = false,
                 IsTupleType = false,
                 IsRecordType = model.Type.IsRecord,
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
             },
 
             TupleDataModel tupleModel => new ObjectShapeModel
@@ -132,7 +126,6 @@ dictionaryModel.ConstructionStrategy is CollectionModelConstructionStrategy.Dict
                 IsValueTupleType = tupleModel.IsValueTuple,
                 IsTupleType = true,
                 IsRecordType = false,
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
             },
 
             _ => new ObjectShapeModel
@@ -143,7 +136,6 @@ dictionaryModel.ConstructionStrategy is CollectionModelConstructionStrategy.Dict
                 IsValueTupleType = false,
                 IsTupleType = false,
                 IsRecordType = false,
-                EmitGenericTypeShapeProviderImplementation = emitGenericProviderImplementation,
             }
         };