perf(source-gen): stop InfrastructureGenerator pinning an old Compilation

src/TUnit.Core.SourceGenerator/CodeGenerators/InfrastructureGenerator.cs

@@ -1,5 +1,6 @@
+using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
 using Microsoft.CodeAnalysis;
-using TUnit.Core.SourceGenerator.CodeGenerators.Equality;
 using TUnit.Core.SourceGenerator.Models;
 using TUnit.Core.SourceGenerator.Models.Extracted;
 
@@ -57,11 +58,13 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                 return !string.Equals(value, "false", StringComparison.OrdinalIgnoreCase);
             });
 
-        // Extract assembly names as primitives in the transform step
-        // This enables proper incremental caching
+        // Extract assembly names as primitives in the transform step. AssemblyInfoModel equality
+        // keeps the generated source cached across keystrokes. No custom Compilation comparer here:
+        // when a comparer reports "equal", the input node keeps the OLD compilation in its table,
+        // pinning it (trees, bound state) until references change. The reference walk is memoized
+        // on the reference list instead, so syntax-only edits still skip it.
         var assemblyInfoProvider = context.CompilationProvider
-            .WithComparer(new PreventCompilationTriggerOnEveryKeystrokeComparer())
-            .Select((compilation, _) => ExtractAssemblyInfo(compilation))
+            .Select((compilation, _) => GetAssemblyInfo(compilation))
             .WithTrackingName(ExtractAssemblyInfoStep)
             .Combine(enabledProvider);
 
@@ -77,6 +80,58 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
         });
     }
 
+    /// <summary>
+    /// Returns the memoized model while the compilation keeps the same reference list. Syntax-only
+    /// edits reuse the backing array of <see cref="Compilation.ExternalReferences"/>; adding, removing,
+    /// rebuilding or (for an IDE project reference) editing a reference produces a new array, so
+    /// extraction reruns. The table is keyed weakly on that array and its values hold only strings,
+    /// so no Compilation or MetadataReference is kept alive by the cache.
+    /// </summary>
+    private static AssemblyInfoModel GetAssemblyInfo(Compilation compilation)
+    {
+        // Directive references (#r in scripts) are not part of the key; skip memoization for them.
+        if (!compilation.DirectiveReferences.IsEmpty)
+        {
+            return ExtractAssemblyInfo(compilation);
+        }
+
+        var references = compilation.ExternalReferences;
+        var key = Unsafe.As<ImmutableArray<MetadataReference>, MetadataReference[]?>(ref references);
+        if (key is null)
+        {
+            return ExtractAssemblyInfo(compilation);
+        }
+
+        var holder = AssemblyInfoCache.GetValue(key, static _ => new AssemblyInfoCacheEntry());
+        var cached = holder.Value;
+        if (cached is not null
+            && cached.Language == compilation.Language
+            && cached.AssemblyName == compilation.AssemblyName)
+        {
+            return cached.Model;
+        }
+
+        var model = ExtractAssemblyInfo(compilation);
+        holder.Value = new CachedAssemblyInfo(compilation.Language, compilation.AssemblyName, model);
+        return model;
+    }
+
+    private static readonly ConditionalWeakTable<MetadataReference[], AssemblyInfoCacheEntry> AssemblyInfoCache = new();
+
+    private sealed class AssemblyInfoCacheEntry
+    {
+        public volatile CachedAssemblyInfo? Value;
+    }
+
+    private sealed class CachedAssemblyInfo(string language, string? assemblyName, AssemblyInfoModel model)
+    {
+        public string Language { get; } = language;
+
+        public string? AssemblyName { get; } = assemblyName;
+
+        public AssemblyInfoModel Model { get; } = model;
+    }
+
     /// <summary>
     /// Extracts all needed data as primitives in the transform step.
     /// This enables proper incremental caching - the model contains only strings.

tests/TUnit.SourceGenerator.IncrementalTests/InfrastructureGeneratorIncrementalTests.cs

@@ -1,3 +1,4 @@
+using System.Runtime.CompilerServices;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using TUnit.Core.SourceGenerator.CodeGenerators;
@@ -26,10 +27,15 @@ public void EditSource_ShouldNotRegenerate()
 
         var driver1 = TestHelper.GenerateTracked<InfrastructureGenerator>(compilation1);
         AssertRunReason(driver1, IncrementalStepRunReason.New);
+        var model1 = GetExtractedModel(driver1);
 
+        // Syntax-only edits rerun the cheap Select, which returns the memoized model (same instance,
+        // so the reference walk was skipped) and keeps the generated source cached.
         var compilation2 = compilation1.AddSyntaxTrees(CSharpSyntaxTree.ParseText("struct MyValue {}"));
         var driver2 = driver1.RunGenerators(compilation2);
-        AssertRunReason(driver2, IncrementalStepRunReason.Cached);
+        AssertRunReason(driver2, IncrementalStepRunReason.Unchanged);
+        Xunit.Assert.Same(model1, GetExtractedModel(driver2));
+        TestHelper.AssertSourceOutputsCached(driver2.GetRunResult().Results[0]);
 
         var compilation3 = TestHelper.ReplaceMethodDeclaration(compilation1, "Test1",
             """
@@ -40,9 +46,53 @@ public void Test1()
             }
             """);
         var driver3 = driver2.RunGenerators(compilation3);
-        AssertRunReason(driver3, IncrementalStepRunReason.Cached);
+        AssertRunReason(driver3, IncrementalStepRunReason.Unchanged);
+        Xunit.Assert.Same(model1, GetExtractedModel(driver3));
+        TestHelper.AssertSourceOutputsCached(driver3.GetRunResult().Results[0]);
+    }
+
+    [Fact]
+    public void EditSource_DoesNotRetainPreviousCompilation()
+    {
+        var (driver, firstCompilation, latestCompilation) = RunSourceEdits(10);
+
+        for (var i = 0; i < 3; i++)
+        {
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+        }
+
+        // A comparer that reports the new compilation as equal makes the input node keep the old
+        // one, pinning its syntax trees and bound state for as long as the driver lives.
+        Xunit.Assert.False(firstCompilation.IsAlive, "The generator driver kept the first compilation alive after syntax-only edits.");
+        GC.KeepAlive(driver);
+        GC.KeepAlive(latestCompilation);
     }
 
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static (GeneratorDriver Driver, WeakReference FirstCompilation, Compilation LatestCompilation) RunSourceEdits(int edits)
+    {
+        Compilation compilation = Fixture.CreateLibrary(CSharpSyntaxTree.ParseText(DefaultSource, CSharpParseOptions.Default));
+        var firstCompilation = new WeakReference(compilation);
+
+        var driver = CSharpGeneratorDriver
+            .Create([new InfrastructureGenerator().AsSourceGenerator()])
+            .RunGenerators(compilation);
+
+        for (var i = 0; i < edits; i++)
+        {
+            compilation = compilation.ReplaceSyntaxTree(
+                compilation.SyntaxTrees.First(),
+                CSharpSyntaxTree.ParseText(DefaultSource + $"\npublic class Edit{i} {{ }}", CSharpParseOptions.Default));
+            driver = driver.RunGenerators(compilation);
+        }
+
+        return (driver, firstCompilation, compilation);
+    }
+
+    private static object? GetExtractedModel(GeneratorDriver driver) =>
+        driver.GetRunResult().Results[0].TrackedSteps[InfrastructureGenerator.ExtractAssemblyInfoStep][0].Outputs[0].Value;
+
     [Fact]
     public void AddReference_ShouldRegenerate()
     {