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

[thomhurst]

Summary

InfrastructureGenerator used CompilationProvider.WithComparer(new PreventCompilationTriggerOnEveryKeystrokeComparer()) so syntax-only edits would skip the reference walk. The side effect: when a comparer says two inputs are equal, Roslyn's InputNode.UpdateStateTable calls TryUseCachedEntries and keeps the old item in the node's state table. The generator driver therefore kept the first Compilation (all its syntax trees and bound state) alive until a reference changed. In an IDE that can be the whole session.

What changed

• Removed the comparer, along with PreventCompilationTriggerOnEveryKeystrokeComparer, which nothing else used.
• CompilationProvider.Select now goes through a small per-driver memo (AssemblyInfoMemo), created in Initialize and captured by the transform. It holds one entry: the last reference list, options, assembly name and the extracted model.
  • References are compared with ImmutableArray<T> ==, which compares the backing array by reference. Syntax-only edits reuse that array, so the memo hits. Adding, removing, rebuilding or editing a reference (including an IDE project reference, per fix(source-gen): model equality covers every emitted field; infrastructure refreshes on reference changes #6912) produces a new array, so extraction reruns.
  • Options are compared by instance, so any options change re-extracts.
  • Script compilations (#r directives, previous submissions) always extract fresh.
  • A source type can shadow a selected referenced type (GetTypeByMetadataName prefers source). Results that depend on shadowing are never stored, and on each hit the selected types are re-checked against the current source.
• Every miss replaces or clears the slot. It only ever holds the latest compilation's reference list and options, which the live compilation holds anyway, and it dies with the driver. It never holds a Compilation or a syntax tree.

Why memoize at all

On TUnit.TestProject (211 references, 568 source files, Roslyn 4.14), the reference walk on a new compilation after a syntax edit takes about 6.5 ms median (9 ms p90). In the IDE that would run on every keystroke. On a memo hit, the shadow re-check takes about 0.7 ms, most of which is building the source namespace members, a cost the compiler pays anyway.

Retention measurements

The harness drove a GeneratorDriver over 20 syntax-only edits, then ran GC.Collect() and checked a WeakReference to the first compilation.

Generator	Roslyn 4.7	Roslyn 4.14	Roslyn 5.3
InfrastructureGenerator, before	first compilation alive	alive	alive
InfrastructureGenerator, after	collected	collected	collected
Control: CompilationProvider.Select without a comparer	collected	collected	collected
Control: WithComparer(always-equal)	alive	alive	alive

Results were the same with step tracking on and off.

The memo depends on the host keeping the same ExternalReferences array across edits. I checked that with an AdhocWorkspace on Roslyn 4.14:

• Editing the project's own document gave a new compilation with the same array, so the memo hits.
• Editing a referenced project gave a new array, so extraction reruns.

Command-line builds are unaffected, because they only run the generator once.

Testing done

• tests/TUnit.SourceGenerator.IncrementalTests (run with dotnet vstest; this project is not run in CI): 57/57 pass.
  • EditSource_DoesNotRetainPreviousCompilation fails on main and passes on this branch.
  • EditSource_ShouldNotRegenerate now expects the step to be Unchanged rather than Cached, since the cheap Select reruns. It also asserts that the source outputs stay cached.
  • The Memo_* tests call AssemblyInfoMemo directly (via InternalsVisibleTo). They cover reuse across syntax edits, and re-extraction on a reference change, an options change, an assembly-name change, and for scripts. Instance identity has to be checked on the memo itself, because the driver keeps the previous output instance whenever a rerun produces an equal model. With the memo disabled, the reuse tests fail.
  • Shadowing is covered inside one driver in both orders: EditSource_ShadowingSelectedType_ShouldRegenerate and EditSource_SourceSensitiveSelectionIsNotReused. FreshDriver_SameReferences_DoesNotShareMemo confirms that separate drivers do not share the memo.
  • AddReference_ShouldRegenerate and EditProjectReference_ShouldRegenerate (from fix(source-gen): model equality covers every emitted field; infrastructure refreshes on reference changes #6912) still pass.
• tests/TUnit.Core.SourceGenerator.Tests on net10.0: 172 passed, 1 skipped. Generated output is unchanged.

Summary by CodeRabbit

• Improvements
  • Syntax-only edits can reuse extracted assembly information when relevant inputs are unchanged, avoiding unnecessary regeneration.
  • Generated outputs remain unchanged after syntax-only edits, and repeated edits no longer keep earlier compilations alive.
  • Infrastructure type selection accounts for source types that shadow referenced types.
  • Changes to referenced types, assembly identity, or metadata import options refresh generated output. Script compilations continue to be handled without reusing cached assembly information.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: e3c25337-4ef0-4a70-b367-9b691a17b50d

📥 Commits

Reviewing files that changed from the base of the PR and between e93802c and 0fe438f.

📒 Files selected for processing (3)

• src/TUnit.Core.SourceGenerator/CodeGenerators/InfrastructureGenerator.cs
• src/TUnit.Core.SourceGenerator/Properties/AssemblyInfo.cs
• tests/TUnit.SourceGenerator.IncrementalTests/InfrastructureGeneratorIncrementalTests.cs

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 1 remain after this review.

---

📝 Walkthrough

Walkthrough

The generator now memoizes assembly-info extraction per driver when reference identity, compilation options, and assembly name permit reuse. It tracks selected metadata names and source shadowing. Tests cover reuse, invalidation, script compilations, generated outputs, and compilation collection.

Changes

Assembly Info Memoization

Layer / File(s)	Summary
Track selected types and source shadowing src/TUnit.Core.SourceGenerator/CodeGenerators/InfrastructureGenerator.cs	Extraction records selected metadata names and detects source types that shadow candidate metadata types. This applies to generic and non-generic type selection.
Memoize assembly information src/TUnit.Core.SourceGenerator/CodeGenerators/Equality/PreventCompilationTriggerOnEveryKeystrokeComparer.cs, src/TUnit.Core.SourceGenerator/CodeGenerators/InfrastructureGenerator.cs	Initialize routes compilations through a per-driver memo. The memo checks external-reference array identity, the compilation options instance, and assembly name. It bypasses caching for script compilations and directive references, and checks cached selections for source shadowing. The compilation comparer is removed.
Verify reuse and cache invalidation src/TUnit.Core.SourceGenerator/Properties/AssemblyInfo.cs, tests/TUnit.SourceGenerator.IncrementalTests/InfrastructureGeneratorIncrementalTests.cs	Tests check model reuse, generated outputs, compilation collection, source shadowing, and invalidation across references, options, assembly names, and script compilations. The source generator grants the test assembly access to internal members.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant CompilationProvider
  participant AssemblyInfoMemo
  participant ExtractAssemblyInfo
  CompilationProvider->>AssemblyInfoMemo: Provide compilation
  AssemblyInfoMemo->>AssemblyInfoMemo: Check reference identity and compilation settings
  AssemblyInfoMemo->>ExtractAssemblyInfo: Extract when reuse is not permitted
  ExtractAssemblyInfo-->>AssemblyInfoMemo: Return model and selected metadata names
  AssemblyInfoMemo-->>CompilationProvider: Return assembly information

Loading

Security Architecture Review

Security architecture risk: 🔵 Low · up to e9380

The new cache does not appear to expand which assemblies can be initialized. It checks compilation context and source-type shadowing before reuse, and reference changes trigger fresh extraction. Some unusual binding and concurrency cases remain untested.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — An incorrect reused model could affect the generated initializer’s referenced-type selection. The observed cache checks and reference-replacement test constrain that path; no broader service or tenant exposure is established.

Trust Boundaries and Controls

• observed — Compilation references supply candidate types, but assembly filters, current-source shadow checks, and the script bypass constrain when a previously selected type can be reused.

Resilience and Maintainability Implications

• inferred — Because entry assignment follows extraction and uses whole-entry publication, a failed extraction before assignment cannot publish a partial model. Concurrent-driver and cancellation behavior has not been directly stress-tested.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.64% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: preventing InfrastructureGenerator from retaining an old Compilation.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{A rabbit checks the cache at dawn
And hops through edits, one by one
The chosen types leave tracks behind
While shadowed names are checked in kind
Fresh scripts bring new paths to run
The rabbit bounds; the work is done}

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

Review

Good fix. The diagnosis is right: a WithComparer that reports two compilations as equal makes Roslyn keep the old Compilation in the input node's state table. That pins its syntax trees and bound state. Replacing the comparer with a memo keyed on the reference list keeps the perf win without the retention, and the new EditSource_DoesNotRetainPreviousCompilation test covers it. The updated EditSource_ShouldNotRegenerate assertions are also good. They check that the model instance is reused and that the source outputs stay cached.

I did not run the tests myself. I judged this from the diff and the PR's own measurements. (The code-review skill failed to run, so this was a manual review.)

Notes and suggestions, none blocking:

1. Unsafe.As<ImmutableArray<T>, T[]> relies on ImmutableArray<T>'s internal layout. The struct is documented as a single-array wrapper, and ImmutableCollectionsMarshal.AsArray does the same thing. That API is not available on netstandard2.0 analyzers, though, so the trick is reasonable. Two safeguards would help:

   • Add a comment saying the layout assumption is deliberate.
   • Keep the incremental test that asserts a reference change reruns extraction (AddReference_ShouldRegenerate, EditProjectReference_ShouldRegenerate). Those tests would catch a Roslyn change that silently breaks the assumption.

2. The memo's correctness depends on host behaviour. Your AdhocWorkspace check on Roslyn 4.14 is useful. A host that mutates references without producing a new array is unlikely, but the failure would be stale output with no error. Keying on the array identity is the right call. The comment on GetAssemblyInfo already states the assumption.

3. Simplification: a smaller cache entry. CachedAssemblyInfo plus the AssemblyInfoCacheEntry holder is two classes for one memo. A single holder class with Language, AssemblyName and Model fields would remove one type. Alternatively, make AssemblyInfoModel carry the language and name. The volatile field is fine for thread safety, since a race only duplicates work.

4. Static shared cache across generator instances. The ConditionalWeakTable is static, so entries are shared across all compilations in the process. Weak keys make this safe. It is also worth knowing that the array key stays alive as long as any compilation holds it, which is the intended lifetime.

5. Minor. The #r directive-reference bail-out is correct. Since it is conservative and cheap, no change is needed.

Overall this is a solid, well-tested fix and I'm happy to approve. The only substantive suggestion is documenting the Unsafe.As assumption.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-29T10:37:42.704835Z	0fe438f	New commits

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[greptile-apps]

Confidence Score: 5/5

[Medium risk] Refactors source generator compilation caching strategy.

The PR appears safe to merge; no outstanding finding or new actionable issue was identified.

Summary

The PR replaces a compilation comparer that retained old compilations with a per-driver memo of the reference walk.

• Syntax-only edits can reuse the extracted model without retaining a Compilation.
• Reference, options, assembly-name, and source-shadowing changes trigger fresh extraction when needed.
• Incremental tests cover memo reuse, invalidation, generated output, and compilation collection.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  C[Current compilation] --> K{Memo key matches and selected types remain unshadowed?}
  K -- Yes --> M[Reuse assembly-info model]
  K -- No --> E[Extract assembly info and update memo]
  M --> G[Generated infrastructure]
  E --> G

Loading

_{Reviews (5) · Last reviewed commit: "refactor(source-gen): memoize Infrastruc..."}

[github-actions]

Review

This fixes a real leak. WithComparer keeps the old Compilation in the input node's state table. Moving the memo to a ConditionalWeakTable keyed on the references array, with a shadow check, is sound. The tests are thorough: retention, shadowing in both directions, and same-instance reuse. I have not run them.

Concerns, most important first:

1. Unsafe.As<ImmutableArray<T>, T[]> is the riskiest part. It relies on ImmutableArray<T> being a struct with a single array field. That layout is not a contract, and the code comment admits it. A layout change would fail quietly, either as key always null (no memo, so a perf regression) or as memory misreads. Two safer options:

   • Multi-target and use ImmutableCollectionsMarshal.AsArray where it exists, with the Unsafe path only for netstandard2.0.
   • Key the memo on something public, such as a ConditionalWeakTable keyed on compilation.ExternalReferences[0] (a MetadataReference reference) plus a length and per-element ReferenceEquals check. This is O(n) but cheap next to the walk. It also avoids a static table keyed on a private array.

2. The memo is process-wide static and shared across drivers. The shadow check handles the known source dependency. Any future source-dependent input to ExtractAssemblyInfo will silently break correctness, and only the incremental tests would catch it. I suggest a comment on ExtractAssemblyInfo saying the result must depend only on references, language and assembly name.

3. Shadow check cost. IsAnyShadowedBySource calls GetTypeByMetadataName once per selected assembly on every keystroke. That is probably fine, but the PR doesn't measure it. Please confirm it is negligible against the old full walk on a large solution.

4. The DirectiveReferences skip is a small gap. Script compilations never memoize, which is acceptable. It is worth a test, or a note that it is intentionally uncovered.

Nothing blocking beyond the Unsafe.As fragility. I'd want that addressed, or at least guarded by a test that fails loudly if the layout assumption breaks, before merge.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 242761574f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[thomhurst]

Follow-up to the github-actions review points, in f5d8c55:

1. Unsafe.As layout. A one-time startup check (CheckImmutableArrayLayout) confirms that ImmutableArray<T> has the size of one array reference and that Unsafe.As returns its backing array with the same elements. If the check fails, the memo is disabled. The generator then extracts on every run, which is slower but still correct. It can never misread memory into a wrong model. EditSource_ShouldNotRegenerate asserts that the memoized instance is reused, so a disabled memo fails that test loudly. ImmutableCollectionsMarshal.AsArray is not available to the netstandard2.0 generator.
2. Reference-only contract. ExtractAssemblyInfo now documents that its result must depend only on references, their binding options, the language and the assembly name. The shadowing dependency is reported and re-checked on reuse. Any new input must be added to CachedAssemblyInfo.Matches. The Codex point about binding options is also fixed: the comparer, resolver and import options are now part of the match.
3. Shadow-check cost. On reuse the check does one IAssemblySymbol.GetTypeByMetadataName against the source assembly for each selected type, that is, one per loaded library. The old path walked every referenced assembly and enumerated public types. The check is a dictionary lookup in the source declaration table, so its cost is negligible by comparison.
4. Directive references. Compilations with #r directives skip the memo on purpose. They are script compilations, and the directive references are not part of the key. They take the normal extraction path, which the existing tests already cover. No dedicated test was added.

[github-actions]

Review of #6926

The diagnosis is right. A comparer that reports "equal" makes InputNode keep the old Compilation, which pins it. The fix is verified by a new WeakReference test, and the shadowing and binding-option edge cases each have tests. I have one main concern and two smaller ones.

1. Unsafe.As<ImmutableArray<T>, T[]> is a fragile way to build the cache key (main concern)

• The memo relies on the internal layout of ImmutableArray<T>, with a startup probe to check it. That is a lot of machinery for a source generator, and it is easy to break under a future runtime or Roslyn change.
• When the probe fails, the code falls back silently to no memoization. That is safe for correctness, but the only thing that would notice is EditSource_ShouldNotRegenerate, which runs in this repo's CI and not on user machines.
• Consider a design that needs no layout hack. Add an intermediate incremental step that projects the compilation to a small value type carrying Language, AssemblyName, the binding options, and the MetadataReference list. Give it an IEquatable implementation that compares references with ReferenceEquals. Roslyn would then keep the old small value rather than the compilation.
• The catch is that extraction needs a Compilation, and holding one in the key would reintroduce the leak. context.MetadataReferencesProvider may be worth a look, since it already compares by reference.
• If the layout hack stays, please say in the code why ExternalReferences itself cannot be the key.

2. Complexity and correctness surface

• The memo, the shadow check, the options matching and the selectedMetadataNames out-parameter add about 150 lines. The doc comment on ExtractAssemblyInfo admits that any new input read there must be added to Matches. That invariant is easy to violate later.
• The static ConditionalWeakTable is shared across drivers, which is what forced the shadow check. A per-generator-instance cache would avoid cross-driver sharing and remove some of the special cases.
• holder.Value is a benign race. It would be worth a comment saying two threads may both extract.

3. Per-edit cost

• IsAnyShadowedBySource calls GetTypeByMetadataName once per selected assembly on every keystroke. That is probably cheap, but the PR description has no timing for it. A before/after number for the syntax-edit path would help.

I found no functional bugs. I'd like #1 reconsidered before merging, or at least the reasoning for keeping it documented.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f5d8c5521c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[thomhurst]

Replies to the latest github-actions review. Changes are in e93802c.

1. The Unsafe.As key. I kept it and added a comment on GetReferencesKey that explains why:

• Compilation.ExternalReferences is an ImmutableArray<T>, which is a struct, so it cannot be a ConditionalWeakTable key. Boxing it creates a new object on every call, so the table would never hit. The backing array is the stable identity that Roslyn reuses across source-only edits.
• A value-type projection step, or MetadataReferencesProvider, cannot replace it. Extraction needs the Compilation to bind referenced assembly symbols. Any node that carries the Compilation into the transform is what pinned it in the first place.

The startup layout check disables the memo if the assumption ever breaks. That costs speed only and never gives a wrong result. EditSource_ShouldNotRegenerate fails loudly in this repo's CI if that happens.

2. Per-instance cache vs static. A per-generator-instance cache would not remove the shadow check. Within one driver a source edit can still shadow a selected type, and the in-driver test EditSource_ShadowingSelectedType_ShouldRegenerate covers that. The static table only widens reuse to fresh drivers over the same reference array, which is the build-server and IDE-reload case. I kept it static. Its correctness rules are in the ExtractAssemblyInfo doc comment: the result depends only on references, binding options, language and assembly name, and the one source dependency is re-checked.

Race on holder.Value. It is benign, and a comment now says so. Two threads can both miss, both extract and both store. Every stored model equals a fresh extraction, so the only cost is duplicate work.

3. Shadow-check cost. On reuse, IsAnyShadowedBySource calls GetTypeByMetadataName on compilation.Assembly (the source assembly only), once per selected type. N is the number of referenced assemblies that reference TUnit.Core, usually a few. These are declaration-table lookups with no binding and no walk over referenced assemblies. The full extraction they replace walks every referenced assembly's namespaces.

Also (codex follow-up): ReferencesSupersedeLowerVersions is internal and only set by the scripting host, so script compilations now skip the memo entirely. New test: ScriptCompilation_SameReferences_IsNeverMemoized.

[github-actions]

Review of PR 6926

The diagnosis is right. A WithComparer that returns "equal" makes the input node keep the old Compilation, so removing it fixes the leak. The before/after/control measurements across Roslyn 4.7, 4.14 and 5.3 are convincing. I did not build or run anything; this is from reading the diff. The code-review skill errored, so this is a manual review.

Concerns

1. The memo is much more complex than the bug it fixes.

   • InfrastructureGenerator now carries a static ConditionalWeakTable, a racy holder, a Matches check on four compilation options, a shadow-by-source re-check, a script/#r bypass and a layout probe.
   • The doc comments say any new input read by ExtractAssemblyInfo must be added to Matches. That invariant is enforced only by convention. A later edit that reads another compilation property would return stale models silently, and the symptom would be an IDE-only wrong AssemblyInfoModel.
   • The static cache is shared across independent drivers. That is why the shadowing and options checks are needed at all.
   • Is the reference walk expensive enough to justify this? If not, dropping the memo and relying on AssemblyInfoModel equality to keep the generated source cached is the most maintainable option. The description does not say what the walk costs.

2. Unsafe.As<ImmutableArray<T>, T[]> is fragile.

   • It relies on a private struct layout. The startup probe and the test mitigate this, but a failure only shows up as a silent perf regression.
   • ImmutableArray<T> == already compares the backing array by reference, and it is public.
   • A safer design is to make the per-driver state the cache. Hold the last (ImmutableArray<MetadataReference>, AssemblyInfoModel) pair in an instance object created in Initialize, and compare with == plus language and name. That removes the static weak table, the cross-driver sharing that forces the shadow and options checks, and the Unsafe hack.

3. Lifetime. CachedAssemblyInfo holds MetadataReferenceResolver and AssemblyIdentityComparer objects. Some host resolvers hold file-system or reference caches. Please confirm the claim that nothing keeps a MetadataReference alive for common hosts (workspace and scripting resolvers), or this may reintroduce a smaller leak.

Positives

• The comparer and its dead class are removed cleanly.
• The regression test is good, and so is the fallback that disables memoization rather than risking correctness.

The leak fix is correct and worth merging. I would simplify the memoization along the lines of point 2, or justify why a process-wide cache is needed. If this shape stays, please add a test that fails when ExtractAssemblyInfo reads an input that Matches does not cover.