Optimized the Type Module Source Generator (#6388)

ChilliCream · Jul 25, 2023 · 6ab57c1 · 6ab57c1
1 parent 16355ea
commit 6ab57c1
Show file tree

Hide file tree

Showing 16 changed files with 280 additions and 174 deletions.
diff --git a/src/HotChocolate/Core/.vscode/settings.json b/src/HotChocolate/Core/.vscode/settings.json
@@ -0,0 +1,3 @@
+{
+    "dotnet.defaultSolution": "HotChocolate.Core.sln"
+}
diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Generators/DataLoaderGenerator.cs b/src/HotChocolate/Core/src/Types.Analyzers/Generators/DataLoaderGenerator.cs
@@ -40,21 +40,10 @@ public bool Consume(ISyntaxInfo syntaxInfo)
     public void Generate(
         SourceProductionContext context,
         Compilation compilation,
-        IReadOnlyCollection<ISyntaxInfo> syntaxInfos)
+        ReadOnlySpan<ISyntaxInfo> syntaxInfos)
     {
-        var module =
-            syntaxInfos.OfType<ModuleInfo>().FirstOrDefault() ??
-            new ModuleInfo(
-                compilation.AssemblyName is null
-                    ? "AssemblyTypes"
-                    : compilation.AssemblyName?.Split('.').Last() + "Types",
-                ModuleOptions.Default);
-
-        var defaults =
-            syntaxInfos.OfType<DataLoaderDefaultsInfo>().FirstOrDefault() ??
-            new DataLoaderDefaultsInfo(null, null, true, true);
-
-        var processed = new HashSet<string>(StringComparer.Ordinal);
+        var (module, defaults) = syntaxInfos.GetDataLoaderDefaults(compilation.AssemblyName);
+
         var dataLoaders = new List<DataLoaderInfo>();
         var sourceText = StringBuilderPool.Get();
 
@@ -66,80 +55,79 @@ compilation.AssemblyName is null
 
         foreach (var syntaxInfo in syntaxInfos)
         {
-            if (syntaxInfo is DataLoaderInfo dataLoader)
+            if (syntaxInfo is not DataLoaderInfo dataLoader)
             {
-                if (dataLoader.MethodSymbol.Parameters.Length == 0)
-                {
-                    context.ReportDiagnostic(
-                        Diagnostic.Create(
-                            _keyParameterMissing,
-                            Location.Create(
-                                dataLoader.MethodSyntax.SyntaxTree,
-                                dataLoader.MethodSyntax.ParameterList.Span)));
-                    continue;
-                }
-
-                if (dataLoader.MethodSymbol.DeclaredAccessibility is not Accessibility.Public
-                    and not Accessibility.Internal and not Accessibility.ProtectedAndInternal)
-                {
-                    context.ReportDiagnostic(
-                        Diagnostic.Create(
-                            _methodAccessModifierInvalid,
-                            Location.Create(
-                                dataLoader.MethodSyntax.SyntaxTree,
-                                dataLoader.MethodSyntax.Modifiers.Span)));
-                    continue;
-                }
+                continue;
+            }
 
-                var keyArg = dataLoader.MethodSymbol.Parameters[0];
-                var keyType = keyArg.Type;
-                var cancellationTokenIndex = -1;
-                var serviceMap = new Dictionary<int, string>();
+            if (dataLoader.MethodSymbol.Parameters.Length == 0)
+            {
+                context.ReportDiagnostic(
+                    Diagnostic.Create(
+                        _keyParameterMissing,
+                        Location.Create(
+                            dataLoader.MethodSyntax.SyntaxTree,
+                            dataLoader.MethodSyntax.ParameterList.Span)));
+                continue;
+            }
 
-                if (IsKeysArgument(keyType))
-                {
-                    keyType = ExtractKeyType(keyType);
-                }
+            if (dataLoader.MethodSymbol.DeclaredAccessibility is not Accessibility.Public
+                and not Accessibility.Internal and not Accessibility.ProtectedAndInternal)
+            {
+                context.ReportDiagnostic(
+                    Diagnostic.Create(
+                        _methodAccessModifierInvalid,
+                        Location.Create(
+                            dataLoader.MethodSyntax.SyntaxTree,
+                            dataLoader.MethodSyntax.Modifiers.Span)));
+                continue;
+            }
 
-                InspectDataLoaderParameters(
-                    dataLoader,
-                    ref cancellationTokenIndex,
-                    serviceMap);
+            var keyArg = dataLoader.MethodSymbol.Parameters[0];
+            var keyType = keyArg.Type;
+            var cancellationTokenIndex = -1;
+            var serviceMap = new Dictionary<int, string>();
 
-                DataLoaderKind kind;
+            if (IsKeysArgument(keyType))
+            {
+                keyType = ExtractKeyType(keyType);
+            }
 
-                if (IsReturnTypeDictionary(dataLoader.MethodSymbol.ReturnType, keyType))
-                {
-                    kind = DataLoaderKind.Batch;
-                }
-                else if (IsReturnTypeLookup(dataLoader.MethodSymbol.ReturnType, keyType))
-                {
-                    kind = DataLoaderKind.Group;
-                }
-                else
-                {
-                    keyType = keyArg.Type;
-                    kind = DataLoaderKind.Cache;
-                }
+            InspectDataLoaderParameters(
+                dataLoader,
+                ref cancellationTokenIndex,
+                serviceMap);
 
-                var valueType = ExtractValueType(dataLoader.MethodSymbol.ReturnType, kind);
+            DataLoaderKind kind;
 
-                if (processed.Add(dataLoader.FullName))
-                {
-                    dataLoaders.Add(dataLoader);
-
-                    GenerateDataLoader(
-                        dataLoader,
-                        defaults,
-                        kind,
-                        keyType,
-                        valueType,
-                        dataLoader.MethodSymbol.Parameters.Length,
-                        cancellationTokenIndex,
-                        serviceMap,
-                        sourceText);
-                }
+            if (IsReturnTypeDictionary(dataLoader.MethodSymbol.ReturnType, keyType))
+            {
+                kind = DataLoaderKind.Batch;
+            }
+            else if (IsReturnTypeLookup(dataLoader.MethodSymbol.ReturnType, keyType))
+            {
+                kind = DataLoaderKind.Group;
+            }
+            else
+            {
+                keyType = keyArg.Type;
+                kind = DataLoaderKind.Cache;
             }
+
+            var valueType = ExtractValueType(dataLoader.MethodSymbol.ReturnType, kind);
+
+            dataLoaders.Add(dataLoader);
+
+            GenerateDataLoader(
+                dataLoader,
+                defaults,
+                kind,
+                keyType,
+                valueType,
+                dataLoader.MethodSymbol.Parameters.Length,
+                cancellationTokenIndex,
+                serviceMap,
+                sourceText);
         }
 
         // if we find no valid DataLoader we will not create any file.
@@ -604,3 +592,59 @@ private static ITypeSymbol ExtractValueType(ITypeSymbol returnType, DataLoaderKi
     private static string ToTypeNameNoGenerics(ITypeSymbol typeSymbol)
         => $"{typeSymbol.ContainingNamespace}.{typeSymbol.Name}";
 }
+
+internal static class GeneratorUtils
+{
+    public static ModuleInfo GetModuleInfo(
+        this ReadOnlySpan<ISyntaxInfo> syntaxInfos,
+        string? assemblyName,
+        out bool defaultModule)
+    {
+        foreach (var syntaxInfo in syntaxInfos)
+        {
+            if (syntaxInfo is ModuleInfo module)
+            {
+                defaultModule = false;
+                return module;
+            }
+        }
+
+        defaultModule = true;
+        return new ModuleInfo(CreateModuleName(assemblyName), ModuleOptions.Default);
+    }
+
+    public static (ModuleInfo, DataLoaderDefaultsInfo) GetDataLoaderDefaults(
+        this ReadOnlySpan<ISyntaxInfo> syntaxInfos,
+        string? assemblyName)
+    {
+        ModuleInfo? moduleInfo = null;
+        DataLoaderDefaultsInfo? dataLoaderDefaultsInfo = null;
+
+        foreach (var syntaxInfo in syntaxInfos)
+        {
+            if (moduleInfo is null && syntaxInfo is ModuleInfo mi)
+            {
+                moduleInfo = mi;
+            }
+            else if (dataLoaderDefaultsInfo is null && syntaxInfo is DataLoaderDefaultsInfo dldi)
+            {
+                dataLoaderDefaultsInfo = dldi;
+            }
+
+            if (moduleInfo is not null && dataLoaderDefaultsInfo is not null)
+            {
+                return (moduleInfo, dataLoaderDefaultsInfo);
+            }
+        }
+
+        moduleInfo ??= new ModuleInfo(CreateModuleName(assemblyName), ModuleOptions.Default);
+        dataLoaderDefaultsInfo ??= new DataLoaderDefaultsInfo(null, null, true, true);
+
+        return (moduleInfo, dataLoaderDefaultsInfo);
+    }
+
+    private static string CreateModuleName(string? assemblyName)
+        => assemblyName is null
+            ? "AssemblyTypes"
+            : assemblyName.Split('.').Last() + "Types";
+}
diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Generators/ISyntaxGenerator.cs b/src/HotChocolate/Core/src/Types.Analyzers/Generators/ISyntaxGenerator.cs
@@ -25,5 +25,5 @@ public interface ISyntaxGenerator
     void Generate(
         SourceProductionContext context,
         Compilation compilation,
-        IReadOnlyCollection<ISyntaxInfo> consumed);
+        ReadOnlySpan<ISyntaxInfo> consumed);
 }
diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Generators/ModuleGenerator.cs b/src/HotChocolate/Core/src/Types.Analyzers/Generators/ModuleGenerator.cs
@@ -22,19 +22,12 @@ public bool Consume(ISyntaxInfo syntaxInfo)
     public void Generate(
         SourceProductionContext context,
         Compilation compilation,
-        IReadOnlyCollection<ISyntaxInfo> syntaxInfos)
+        ReadOnlySpan<ISyntaxInfo> syntaxInfos)
     {
-        var module =
-            syntaxInfos.OfType<ModuleInfo>().FirstOrDefault() ??
-            new ModuleInfo(
-                compilation.AssemblyName is null
-                    ? "AssemblyTypes"
-                    : compilation.AssemblyName?.Split('.').Last() + "Types",
-                ModuleOptions.Default);
+        var module = syntaxInfos.GetModuleInfo(compilation.AssemblyName, out var defaultModule);
 
-        var batch = new List<ISyntaxInfo>(syntaxInfos.Where(static t => t is not ModuleInfo));
-
-        if (batch.Count == 0)
+        // if there is only the module info we do not need to generate a module.
+        if (!defaultModule && syntaxInfos.Length == 1)
         {
             return;
         }
@@ -74,7 +67,7 @@ compilation.AssemblyName is null
 
         var operations = OperationType.No;
 
-        foreach (var syntaxInfo in batch.Distinct())
+        foreach (var syntaxInfo in syntaxInfos)
         {
             switch (syntaxInfo)
             {

diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/AggregateInfo.cs b/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/AggregateInfo.cs
@@ -1,41 +0,0 @@
-namespace HotChocolate.Types.Analyzers.Inspectors;
-
-public sealed class AggregateInfo : ISyntaxInfo, IEquatable<AggregateInfo>
-{
-    private readonly IReadOnlyList<ISyntaxInfo> _syntaxInfos;
-
-    public AggregateInfo(IReadOnlyList<ISyntaxInfo> syntaxInfos)
-    {
-        _syntaxInfos = syntaxInfos;
-    }
-
-    public IReadOnlyList<ISyntaxInfo> Items => _syntaxInfos;
-
-    public bool Equals(AggregateInfo? other)
-    {
-        if (ReferenceEquals(null, other))
-        {
-            return false;
-        }
-
-        if (ReferenceEquals(this, other))
-        {
-            return true;
-        }
-
-        return Equals(_syntaxInfos, other._syntaxInfos);
-    }
-
-    public override bool Equals(object? obj)
-        => ReferenceEquals(this, obj) ||
-            obj is AggregateInfo other && Equals(other);
-
-    public override int GetHashCode()
-        => _syntaxInfos.GetHashCode();
-
-    public static bool operator ==(AggregateInfo? left, AggregateInfo? right)
-        => Equals(left, right);
-
-    public static bool operator !=(AggregateInfo? left, AggregateInfo? right)
-        => !Equals(left, right);
-}

diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/ClassBaseClassInspector.cs b/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/ClassBaseClassInspector.cs
@@ -11,11 +11,7 @@ public bool TryHandle(
         GeneratorSyntaxContext context,
         [NotNullWhen(true)] out ISyntaxInfo? syntaxInfo)
     {
-        if (context.Node is ClassDeclarationSyntax
-            {
-                BaseList.Types.Count: > 0,
-                TypeParameterList: null
-            } possibleType)
+        if (context.Node is ClassDeclarationSyntax { BaseList.Types.Count: > 0, TypeParameterList: null } possibleType)
         {
             var model = context.SemanticModel.GetDeclaredSymbol(possibleType);
             if (model is { IsAbstract: false } type)

diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/DataLoaderDefaultsInfo.cs b/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/DataLoaderDefaultsInfo.cs
@@ -39,6 +39,21 @@ public bool Equals(DataLoaderDefaultsInfo? other)
             RegisterServices == other.RegisterServices;
     }
 
+    public bool Equals(ISyntaxInfo other)
+    {
+        if (ReferenceEquals(null, other))
+        {
+            return false;
+        }
+
+        if (ReferenceEquals(this, other))
+        {
+            return true;
+        }
+
+        return other is DataLoaderDefaultsInfo info && Equals(info);
+    }
+
     public override bool Equals(object? obj)
         => ReferenceEquals(this, obj) ||
             (obj is ModuleInfo other && Equals(other));

diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/DataLoaderInfo.cs b/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/DataLoaderInfo.cs
@@ -76,6 +76,21 @@ public bool Equals(DataLoaderInfo? other)
         return AttributeSyntax.Equals(other.AttributeSyntax) &&
             MethodSyntax.Equals(other.MethodSyntax);
     }
+
+    public bool Equals(ISyntaxInfo other)
+    {
+        if (ReferenceEquals(null, other))
+        {
+            return false;
+        }
+
+        if (ReferenceEquals(this, other))
+        {
+            return true;
+        }
+
+        return other is DataLoaderInfo info && Equals(info);
+    }
 
     public override bool Equals(object? obj)
     {

diff --git a/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/ISyntaxInfo.cs b/src/HotChocolate/Core/src/Types.Analyzers/Inspectors/ISyntaxInfo.cs
@@ -1,5 +1,5 @@
 namespace HotChocolate.Types.Analyzers.Inspectors;
 
-public interface ISyntaxInfo
+public interface ISyntaxInfo : IEquatable<ISyntaxInfo>
 {
 }