Index pre-generated types per assembly to avoid O(n) ExportedTypes scans (#189)

src/CodegenTests/CodeGenerationExtensionsTests.cs

@@ -0,0 +1,60 @@
+using System.Reflection;
+using JasperFx.CodeGeneration;
+using Shouldly;
+using Xunit;
+
+namespace CodegenTests;
+
+public class CodeGenerationExtensionsTests
+{
+    [Fact]
+    public void FindPreGeneratedType_returns_a_known_exported_type()
+    {
+        var assembly = typeof(CodeGenerationExtensionsTests).Assembly;
+        var type = assembly.FindPreGeneratedType(
+            typeof(CodeGenerationExtensionsTests).Namespace!,
+            nameof(CodeGenerationExtensionsTests));
+
+        type.ShouldBe(typeof(CodeGenerationExtensionsTests));
+    }
+
+    [Fact]
+    public void FindPreGeneratedType_returns_null_for_unknown_name()
+    {
+        var assembly = typeof(CodeGenerationExtensionsTests).Assembly;
+        var type = assembly.FindPreGeneratedType("Some.Made.Up.Namespace", "DefinitelyDoesNotExist");
+
+        type.ShouldBeNull();
+    }
+
+    [Fact]
+    public void FindPreGeneratedType_is_repeatable_against_the_same_assembly()
+    {
+        // Lookups are now backed by a per-assembly indexed cache. The cache should
+        // be transparent: repeated calls return the same Type without enumerating
+        // ExportedTypes again.
+        var assembly = typeof(CodeGenerationExtensionsTests).Assembly;
+
+        var first = assembly.FindPreGeneratedType(
+            typeof(CodeGenerationExtensionsTests).Namespace!,
+            nameof(CodeGenerationExtensionsTests));
+        var second = assembly.FindPreGeneratedType(
+            typeof(CodeGenerationExtensionsTests).Namespace!,
+            nameof(CodeGenerationExtensionsTests));
+
+        first.ShouldBeSameAs(second);
+        first.ShouldBe(typeof(CodeGenerationExtensionsTests));
+    }
+
+    [Fact]
+    public void FindPreGeneratedType_handles_a_distinct_assembly_independently()
+    {
+        // Smoke test that the per-assembly index keys correctly: looking up a
+        // type defined in System.Private.CoreLib via the BCL assembly should
+        // find it without the cache crossing wires with the test assembly above.
+        var coreLib = typeof(string).Assembly;
+        var stringType = coreLib.FindPreGeneratedType("System", "String");
+
+        stringType.ShouldBe(typeof(string));
+    }
+}

src/JasperFx/CodeGeneration/CodeGenerationExtensions.cs

@@ -1,4 +1,6 @@
+using System.Collections.Concurrent;
 using System.Reflection;
+using System.Runtime.CompilerServices;
 using JasperFx.CodeGeneration.Model;
 using JasperFx.Core;
 
@@ -13,6 +15,13 @@ public MissingTypeException(string? message) : base(message)
 
 public static class CodeGenerationExtensions
 {
+    // Per-assembly cache of pre-generated types indexed by full name. Used by
+    // FindPreGeneratedType so static-mode code-file attachment does not pay an
+    // O(ExportedTypes) scan per lookup. ConditionalWeakTable means we don't
+    // root assemblies that the host might unload (plugin scenarios).
+    private static readonly ConditionalWeakTable<Assembly, IReadOnlyDictionary<string, Type>> _exportedTypeIndex
+        = new();
+
     /// <summary>
     ///     Try to locate a pre-generated type by namespace and type name in the
     ///     supplied assembly
@@ -24,7 +33,26 @@ public static class CodeGenerationExtensions
     public static Type? FindPreGeneratedType(this Assembly assembly, string @namespace, string typeName)
     {
         var fullName = $"{@namespace}.{typeName}";
-        return assembly.ExportedTypes.FirstOrDefault(x => x.FullName == fullName);
+        var index = _exportedTypeIndex.GetValue(assembly, BuildExportedTypeIndex);
+        return index.TryGetValue(fullName, out var type) ? type : null;
+    }
+
+    private static IReadOnlyDictionary<string, Type> BuildExportedTypeIndex(Assembly assembly)
+    {
+        // ExportedTypes can contain duplicates only in pathological cases (forwarded
+        // types crossing module boundaries with name conflicts). Use the first
+        // occurrence to mirror the historical FirstOrDefault behavior.
+        var dict = new Dictionary<string, Type>(StringComparer.Ordinal);
+        foreach (var type in assembly.ExportedTypes)
+        {
+            var fullName = type.FullName;
+            if (fullName != null && !dict.ContainsKey(fullName))
+            {
+                dict[fullName] = type;
+            }
+        }
+
+        return dict;
     }
 
     public static GeneratedAssembly StartAssembly(this ICodeFileCollection generator, GenerationRules rules)
@@ -126,4 +154,4 @@ public static void AssertPreBuildTypesExist(this ICodeFileCollection collection,
                 missing.Select(x => x.ToString())!.Join("\n"));
         }
     }
-}
+}

src/JasperFx/JasperFx.csproj

@@ -3,7 +3,7 @@
         <Description>Foundational helpers and command line support used by JasperFx and the Critter Stack projects</Description>
         <AssemblyName>JasperFx</AssemblyName>
         <PackageId>JasperFx</PackageId>
-        <Version>1.26.0</Version>
+        <Version>1.27.0</Version>
     </PropertyGroup>
 
     <ItemGroup Condition="'$(TargetFramework)' == 'net8.0'">