perf: optimize TUnit.Mocks hot paths

[thomhurst]

Summary

• Eliminated ConditionalWeakTable — replaced with IMockObject interface for O(1) mock-to-wrapper reverse lookup, avoiding expensive GC handle allocation on every mock creation
• Replaced Lazy<T> with nullable fields — saves 48+ bytes per setup object by using manual null-check initialization pattern
• Deferred collection allocation in MockEngine — _setupsByMember, _callHistory, and other dictionaries are now null until first use, reducing mock creation cost to near-zero
• Removed factory delegate wrapper closures — leverages delegate covariance so Func<MockBehavior, Mock<T>> registers directly as Func<MockBehavior, object> without allocating a wrapping lambda
• Lock-free invocation hot path — RecordCall uses only ConcurrentQueue.Enqueue with no lock contention
• Direct lock reference in read paths — GetNextBehavior/GetEventRaises use _behaviorLock! directly instead of EnsureBehaviorLock() when the lock is guaranteed to exist

Test plan

• All 64 source generator snapshot tests pass
• All 672 TUnit.Mocks.Tests pass on net10.0
• No new warnings introduced

[claude]

Performance Optimization Review

This PR makes a series of targeted micro-optimizations to TUnit.Mocks. The intent is sound and most changes are improvements, but a few of them introduce new correctness risks that aren't present in the replaced code. Below is a detailed breakdown.

---

1. IMockObject interface: O(1) lookup is right, but the mutable public property is a design hazard

What changed: ConditionalWeakTable<object, IMock> is replaced by an IMockObject interface that each generated mock class implements, with a MockWrapper { get; set; } property set in Mock<T>'s constructor.

Why the intent is correct: ConditionalWeakTable does allocate a GC handle per entry, and casting directly to a known interface is clearly faster than a table lookup.

The concern — public mutability leaks an invariant into user space. IMockObject.MockWrapper is get; set; and the interface is public. Even with [EditorBrowsable(Never)], a user can write:

((IMockObject)myMock.Object).MockWrapper = null; // or another Mock<T>

This would silently corrupt Mock.Get<T>() for that mock instance. The original ConditionalWeakTable approach was immune to this because the mapping was stored outside the mock object itself. The internal static void Register(...) method, while also callable from external code (via reflection), at least communicated clearly that this was a one-way registration.

Suggestion: Make the setter init-only or internal. Since the implementing class is generated (internal sealed), an internal set would work and would entirely prevent external mutation without requiring any other change to the design.

---

2. EnsureSetup() null-check pattern: thread-safety regression relative to Lazy<T>

What changed: In MockMethodCall<TReturn>, VoidMockMethodCall, MockMethodCall.cs, and all generated *_MockCall classes, Lazy<T> is replaced with a nullable field and a manual null-check pattern:

private MethodSetupBuilder<TReturn>? _builder;

private MethodSetupBuilder<TReturn> EnsureSetup()
{
    if (_builder is { } b) return b;
    var setup = new MethodSetup(_memberId, _matchers, _memberName);
    _engine.AddSetup(setup);
    return _builder = new MethodSetupBuilder<TReturn>(setup);
}

The concern — lost thread-safety guarantee. The original Lazy<T> (default LazyThreadSafetyMode.ExecutionAndPublication) guarantees that the factory runs exactly once across all threads. The new pattern has a classic TOCTOU race: two threads can both observe _builder == null, both call _engine.AddSetup(setup), and create two distinct MethodSetup objects. This results in a duplicate setup being registered with the engine, which could cause a method to match twice or produce confusing diagnostics.

For VoidMockMethodCall (eagerly registered in the constructor), the race only exists between construction and a concurrent EnsureSetup() call. For MockMethodCall<TReturn>, the window is wider because lazy initialization is explicitly deferred to first use.

The PR description acknowledges this implicitly ("Thread safety note: EnsureSetup() is intentionally not thread-safe; MockMethodCall objects are per-test and single-threaded by design.") but this assumption is not enforced by the type system and is not documented in the code itself. If the assumption is actually guaranteed by TUnit's execution model, a comment at the field declaration stating the invariant would eliminate the concern without any code change.

Suggestion: Either (a) add a // Not thread-safe: MockMethodCall instances are owned by a single test and never shared across threads comment at the field declaration to capture the invariant, or (b) use Interlocked.CompareExchange for the assignment if the single-ownership invariant ever needs to be relaxed.

---

3. RecordCall in MockEngine — LazyInitializer.EnsureInitialized on the hot path

What changed:

// Before
_callHistory.Enqueue(record);

// After
var history = _callHistory ?? LazyInitializer.EnsureInitialized(ref _callHistory)!;
history.Enqueue(record);

The concern — LazyInitializer.EnsureInitialized can allocate multiple ConcurrentQueue<CallRecord> instances. Without a lock, two concurrent calls to RecordCall can both reach LazyInitializer.EnsureInitialized. The method uses Interlocked.CompareExchange internally, so only one wins, but the losing thread's freshly constructed ConcurrentQueue<CallRecord> is discarded. This is correct, but at the expense of allocating an extra object. More importantly, _callHistory is also set to null in Reset():

Volatile.Write(ref _callHistory, null);

If Reset() is called concurrently with a method invocation (which can happen if a test is torn down while still running), _callHistory can be nulled out between the ?? check and the Enqueue. The new LazyInitializer.EnsureInitialized approach handles this correctly because EnsureInitialized re-reads the field atomically. But the optimization in _callHistory ?? LazyInitializer.EnsureInitialized(...) is redundant: when _callHistory is non-null, the ?? short-circuits correctly; when it is null, both sides of ?? effectively do the same thing. The ?? does avoid the method call overhead on the fast path, which is the point — just worth being explicit about this.

---

4. MethodSetup — inconsistent locking in GetNextBehavior and GetEventRaises after removing the always-initialized lock

What changed: _behaviorLock is now Lock? (nullable), initialized lazily via EnsureBehaviorLock(). Read paths use _behaviorLock! directly with a null-check guard on _behaviors or _eventRaises:

public IBehavior? GetNextBehavior()
{
    if (_behaviors is null)
        return null;

    lock (_behaviorLock!)   // _behaviorLock must be non-null if _behaviors is non-null
    { ... }
}

The concern — the invariant "if _behaviors is non-null then _behaviorLock is non-null" is not compiler-enforced. It is true by construction: AddBehavior calls EnsureBehaviorLock() before assigning _behaviors. But it relies on write ordering that is not explicitly fenced. On weakly-ordered architectures (not x86, but the code targets net8.0;net9.0;net10.0 which includes ARM), the write to _behaviors inside the lock could be observed by another thread before the write to _behaviorLock (both are ordinary heap writes). The null-check on _behaviors is a read outside any lock, making this a potential data race under the C# memory model.

In practice, on .NET the JIT typically emits barriers around lock statements that prevent this, and LazyInitializer.EnsureInitialized uses Interlocked.CompareExchange which includes a full memory barrier. But the correctness argument is non-trivial and should be documented.

---

5. RegisterFactory / RegisterPartialFactory / RegisterDelegateFactory — delegate covariance via type erasure

What changed:

// Before
_factories[typeof(T)] = behavior => factory(behavior);

// After
_factories[typeof(T)] = factory;

This is clean and correct. Func<MockBehavior, Mock<T>> is covariant in its return type relative to Func<MockBehavior, object> — Mock<T> is a reference type and object is its base class, so delegate covariance applies. This genuinely eliminates the wrapper closure allocation. Good change.

One remaining inconsistency: RegisterWrapFactory still keeps its wrapper lambda:

_wrapFactories[typeof(T)] = (behavior, instance) => factory(behavior, (T)instance);
// Note: wrap factory requires cast so the delegate wrapper is retained.

The cast on instance is unavoidable since the stored delegate signature is Func<MockBehavior, object, object> and the original factory takes T. The comment added in this PR explains this correctly, which is appreciated.

---

6. Reset() — null assignment instead of Clear() for _setupsByMember

What changed:

// Before
_setupsByMember.Clear();

// After
_setupsByMember = null;

This is a correct improvement. Assigning null allows the GC to collect the dictionary and all the List<MethodSetup> entries it holds, rather than keeping the dictionary alive. Since Reset() is used between test runs in MockRepository, this reduces memory pressure. The guard on _setupsByMember is not null in downstream reads handles the null case correctly.

---

7. Deferred collection allocation in MockEngine — consistency issue

The PR defers _setupsByMember and _callHistory to first use. However, _setupLock (a Lock) is still eagerly allocated in the constructor:

private readonly Lock _setupLock = new();

Lock is a struct in .NET 9+, so this is actually fine — it is allocated as part of the MockEngine object itself, not as a separate heap allocation. But the analogous _behaviorLock in MethodSetup is Lock? and lazily initialized. The inconsistency is understandable (different threat models) but worth noting: MockEngine._setupLock guards _setupsByMember, while MethodSetup._behaviorLock guards _behaviors. If the pattern is "defer the lock until the protected collection is needed," they should be consistent. If the difference is intentional (e.g., _setupLock guards state-machine fields too), a comment would clarify.

---

Summary

Change	Verdict
IMockObject interface for O(1) reverse lookup	Correct improvement; setter mutability is a minor concern
Lazy<T> → manual null-check in EnsureSetup()	Performance win; thread-safety invariant needs documentation
ConcurrentQueue deferred init via LazyInitializer	Correct; fast-path short-circuit is intentional and useful
MethodSetup lazy lock / behaviors	Correct by construction; cross-thread visibility argument non-obvious
Delegate covariance eliminating factory closures	Clean and correct
Reset() assigning null instead of Clear()	Correct improvement
Void method eager EnsureSetup() call in constructor	Correct; preserves prior behavior while removing the Lazy<T> overhead

The biggest ask is: document the single-threaded ownership invariant for MockMethodCall and generated *_MockCall classes, either as a code comment or in the class-level XML doc. Without it, the removal of Lazy<T> looks like a thread-safety regression to reviewers unfamiliar with TUnit's execution model.

[claude]

Performance optimization review: TUnit.Mocks hot paths

Overall this is a solid, well-targeted performance PR. The motivations are sound and the allocations genuinely removed are real. A few concerns worth discussing before merge.

---

Race condition risk in EnsureSetup() — not thread-safe

The Lazy<T> replacement in MockMethodCall<TReturn>, VoidMockMethodCall, and the source-generated call classes introduces a TOCTOU race:

private MethodSetupBuilder<TReturn> EnsureSetup()
{
    if (_builder is { } b) return b;           // (1) reads null
    // --- another thread arrives here ---
    var setup = new MethodSetup(_memberId, _matchers, _memberName);
    _engine.AddSetup(setup);                    // (2) both threads add a duplicate setup!
    return _builder = new MethodSetupBuilder<TReturn>(setup);
}

Lazy<T> (default LazyThreadSafetyMode.ExecutionAndPublication) guaranteed exactly one AddSetup call per mock-call object. The new code has no synchronisation — two threads calling a method setup concurrently can each pass the null-check and register duplicate MethodSetup instances. While mock-call objects are typically used single-threaded (setup is usually done before tests run), the absence of a guarantee is a regression that would surface in parallel test scenarios.

Suggestion: use LazyInitializer.EnsureInitialized with a factory, or a lock, or at minimum a Volatile.Read/Interlocked.CompareExchange pattern. Alternatively, document that MockMethodCall is not thread-safe during construction/setup. If single-threaded use is the intended invariant, add a comment to that effect so the contract is explicit.

---

IMockObject.MockWrapper is a public mutable property on a public interface

IMockObject is public, and MockWrapper has a public setter. This means user code can call ((IMockObject)myMock).MockWrapper = null and silently break Mock.Get(...) lookups with no compile-time or run-time guard. The old ConditionalWeakTable approach was fully encapsulated inside Mock.

The [EditorBrowsable(Never)] attribute hides it from IDE autocompletion but does not prevent intentional (or accidental) use.

Suggestion: consider making the interface internal (since it is only implemented by generated classes in the same assembly or in the user's test assembly via source gen). If it must remain public (e.g., for source-generated code in a different assembly), adding a init-only setter — if the generated code can be restructured to set it during construction — would at least remove the mutable footgun. Alternatively, IMockObject could expose only a getter and the setter could be surfaced through a separate internal helper.

---

MethodSetup.GetNextBehavior / GetEventRaises: non-atomic check-then-lock

public IBehavior? GetNextBehavior()
{
    if (_behaviors is null)   // unsynchronised read
        return null;

    lock (_behaviorLock!)     // non-null assumed
    { ... }
}

This is the same pattern for GetEventRaises. The guard reads _behaviors without a lock; another thread could be inside AddBehavior (which acquires the lock) writing _behaviors concurrently. On sufficiently aggressive CPU architectures or JIT reorderings this is a data race on the reference itself. Because both reads and writes happen on plain fields without volatile, the null-check can observe a stale/teared value.

The previous code always acquired the lock before touching _behaviors, which was safe. The new "fast path" skips the lock, giving a potential improvement for the common case but at the cost of correctness.

Suggestion: mark _behaviors and _eventRaises as volatile, or use Volatile.Read in the fast-path null check (consistent with how the PR handles _callHistory and _autoTrackValues elsewhere in MockEngine).

---

OutRefAssignments property: lock skipped when _behaviorLock is null

public Dictionary<int, object?>? OutRefAssignments
{
    get
    {
        if (_behaviorLock is not { } lck)
            return _outRefAssignments;   // no lock
        lock (lck) { return _outRefAssignments; }
    }
}

If _behaviorLock is null it means no behaviors have been registered, so _outRefAssignments should also be null — but SetOutRefValue creates the lock via EnsureBehaviorLock(), so the ordering is: lock is created, then dictionary is written. There is a window where _behaviorLock becomes non-null but _outRefAssignments hasn't been written yet from another thread's perspective without a volatile fence. Same concern as above — this is safe in practice today because SetOutRefValue is called during setup (single-threaded), but the volatile annotation would make the guarantee explicit and future-proof.

---

Minor: Reset() nulls _callHistory under a different lock than it's written

RecordCall writes _callHistory via LazyInitializer.EnsureInitialized (no lock, relies on LazyInitializer's CAS), while Reset() nulls it with Volatile.Write (under no lock). This is safe because Volatile.Write guarantees the store is visible, but a concurrent RecordCall after Reset() starts could re-initialize the field before Reset() stores null, causing the reset queue to be re-populated immediately. This was also the case before (draining a queue while new items are enqueued has the same problem), so it's not a regression — just worth a comment.

---

Positive observations

• ConditionalWeakTable removal is the biggest win. ConditionalWeakTable has non-trivial overhead (GC handle registration, separate hash table, thread-safe internal locking on every TryGetValue). Routing through a direct interface property is strictly better for the common path.
• Delegate covariance on factory registration (RegisterFactory<T> storing Func<MockBehavior, Mock<T>> directly as Func<MockBehavior, object>) is a correct and clean use of variance — no lambda wrapper allocation.
• Deferred _setupsByMember / _callHistory allocation in MockEngine is a sensible optimisation for short-lived or unconfigured mocks.
• Snapshot tests all updated — the generated code diffs are consistent and the .verified.txt files look correct.

[claude]

Code Review: perf/mock-optimizations

Overall this is a well-motivated PR with clear performance wins. The design decisions are mostly sound, but there are a few correctness concerns worth discussing.

---

Issue 1 (Correctness): is no longer thread-safe

Where: MockMethodCall.cs, VoidMockMethodCall.cs, and all generated *_MockCall classes.

The replacement of Lazy<T> with the null-check pattern is not thread-safe:

private MethodSetupBuilder<TReturn> EnsureSetup()
{
    if (_builder is { } b) return b;   // both threads see null
    var setup = new MethodSetup(...);
    _engine.AddSetup(setup);           // both threads call AddSetup — duplicate setup registered
    return _builder = new MethodSetupBuilder<TReturn>(setup);
}

If two threads call EnsureSetup() concurrently (e.g., parallel test setup or auto-mocked return values resolved concurrently), both can pass the null check and call _engine.AddSetup(setup) twice, registering duplicate setups. The original Lazy<T> prevented this exactly.

The previous Lazy<T> had a closure allocation, but that's avoidable without losing thread safety. Consider LazyInitializer.EnsureInitialized (already used in MockEngine.RecordCall) or a dedicated factory:

private MethodSetupBuilder<TReturn>? _builder;
private MethodSetupBuilder<TReturn> EnsureSetup() =>
    LazyInitializer.EnsureInitialized(ref _builder, () =>
    {
        var setup = new MethodSetup(_memberId, _matchers, _memberName);
        _engine.AddSetup(setup);
        return new MethodSetupBuilder<TReturn>(setup);
    })!;

This avoids the Lazy<T> object allocation while preserving the thread-safety guarantee. The same applies to MethodSetup._behaviors, ._eventRaises, etc.

Why this matters: Mock setup is often done in [Before(Test)] hooks that can run in parallel with other test lifecycle methods. The regression is subtle and won't be caught by unit tests that run setup on a single thread.

---

Issue 2 (Correctness): Lock-free reads of _behaviors/_eventRaises violate the C# memory model

Where: MethodSetup.GetNextBehavior(), GetEventRaises(), OutRefAssignments getter.

// GetNextBehavior
if (_behaviors is null) return null;  // read without lock or Volatile.Read
lock (_behaviorLock!) { ... }

_behaviors is written inside a lock (which provides a memory barrier on exit), but the null check here is a plain read without any barrier. The JIT/CPU is theoretically allowed to cache null in a register. The correct pattern is:

if (Volatile.Read(ref _behaviors) is not { } behaviors) return null;
lock (_behaviorLock!) { ... }

In practice on x64 this is unlikely to manifest, but it's important for correctness on ARM and for AOT scenarios (which TUnit explicitly targets).

---

Issue 3 (Design): IMockObject is public with a settable MockWrapper

Where: IMockObject.cs, generated mock impls.

The interface is explicitly implemented (good — requires cast to access), but the property is declared with a public setter:

IMock? MockWrapper { get; set; }

Since set is accessible through an IMockObject cast, test code could write:

((IMockObject)(object)myMock).MockWrapper = null;  // breaks Mock.Get<T>()

This is a sabotage-only scenario and the [EditorBrowsable(Never)] discourages it, but an init-only setter (or internal setter on a concrete type) would be more robust. Alternatively, since the only setter call is mockObj.MockWrapper = this in the Mock<T> constructor, you could make the property { get; } and wire it in the generated constructor instead.

---

Issue 4 (Minor): Reset() replaces queue drain with null assignment — one edge case

Where: MockEngine.cs:Reset()

// Before:
while (_callHistory.TryDequeue(out _)) { }
// After:
Volatile.Write(ref _callHistory, null);

This is cleaner and correct for the common case. The one edge: if a test calls Reset() while invocations are still in flight (concurrent test teardown), calls recorded between the Volatile.Write(null) and the next RecordCall reinitializing via LazyInitializer.EnsureInitialized will end up in the new queue — correct behavior. The old drain-loop had the same window. ✓

---

What's excellent in this PR

• IMockObject interface for reverse lookup is architecturally superior to ConditionalWeakTable. It trades a GC-handle per mock for a single interface property slot, which is always cheaper. The explicit interface implementation keeps it invisible to users. This is the right design.
• Deferred allocation in MockEngine is a solid optimization — most mocks in a test suite have zero calls or one behavior, so pre-allocating Dictionary and ConcurrentQueue was wasteful. The LazyInitializer.EnsureInitialized usage in RecordCall is the correct pattern here.
• Factory covariance (Func<MockBehavior, Mock<T>> stored as Func<MockBehavior, object>) is a correct use of delegate covariance and saves one closure allocation per registered type — clean.
• _behaviorLock!. used in read paths when existence is guaranteed is sound reasoning, given that _behaviors being non-null implies the lock was already created.

---

Summary: Issues 1 and 2 are correctness concerns worth addressing before merge. Issue 1 especially is a behavioral regression in concurrent scenarios. Issue 3 is a minor design concern. The overall direction and most of the changes are well-executed.

[claude]

Code Review: perf/mock-optimizations

Overall this is a well-motivated PR with clear performance wins. The design decisions are mostly sound, but there are a few correctness concerns worth discussing.

---

Issue 1 (Correctness): EnsureSetup() is no longer thread-safe

Where: MockMethodCall.cs, VoidMockMethodCall.cs, and all generated *_MockCall classes.

The replacement of Lazy<T> with the null-check pattern is not thread-safe:

private MethodSetupBuilder<TReturn> EnsureSetup()
{
    if (_builder is { } b) return b;   // both threads see null
    var setup = new MethodSetup(...);
    _engine.AddSetup(setup);           // both threads call AddSetup — duplicate setup registered
    return _builder = new MethodSetupBuilder<TReturn>(setup);
}

If two threads call EnsureSetup() concurrently (e.g., parallel test setup or auto-mocked return values resolved concurrently), both can pass the null check and call _engine.AddSetup(setup) twice, registering duplicate setups. The original Lazy<T> prevented this.

The previous Lazy<T> had a closure allocation, but that's avoidable without losing thread safety. LazyInitializer.EnsureInitialized (already used in MockEngine.RecordCall) is the right tool:

private MethodSetupBuilder<TReturn>? _builder;
private MethodSetupBuilder<TReturn> EnsureSetup() =>
    LazyInitializer.EnsureInitialized(ref _builder, () =>
    {
        var setup = new MethodSetup(_memberId, _matchers, _memberName);
        _engine.AddSetup(setup);
        return new MethodSetupBuilder<TReturn>(setup);
    })!;

This avoids the Lazy<T> object allocation while preserving the thread-safety guarantee. The same applies to MethodSetup._behaviors, ._eventRaises, etc.

Why this matters: Mock setup is often done in [Before(Test)] hooks that can run in parallel with other test lifecycle methods. The regression is subtle and won't be caught by tests that run setup single-threaded.

---

Issue 2 (Correctness): Lock-free reads of _behaviors/_eventRaises violate the C# memory model

Where: MethodSetup.GetNextBehavior(), GetEventRaises(), OutRefAssignments getter.

// GetNextBehavior
if (_behaviors is null) return null;  // plain read without lock or Volatile.Read
lock (_behaviorLock!) { ... }

_behaviors is written inside a lock (which provides a memory barrier on exit), but this null check is a plain read with no barrier. The JIT/CPU is theoretically allowed to cache null in a register. The correct pattern is:

if (Volatile.Read(ref _behaviors) is not { } behaviors) return null;
lock (_behaviorLock!) { ... }

In practice on x64 this is unlikely to manifest, but it matters for correctness on ARM and for AOT scenarios which TUnit explicitly targets.

---

Issue 3 (Design): IMockObject.MockWrapper is settable by anyone who casts

Where: IMockObject.cs

The interface is explicitly implemented (good — requires cast to access), but the property has a public setter. Test code can write:

((IMockObject)(object)myMock).MockWrapper = null;  // breaks Mock.Get<T>()

The [EditorBrowsable(Never)] attribute discourages this but doesn't prevent it. Since the only legitimate setter call is mockObj.MockWrapper = this in the Mock<T> constructor, consider wiring it in the generated mock constructor instead and making the property getter-only. This would make the invariant structurally enforced rather than relying on convention.

---

Issue 4 (Minor): Reset() null-assignment edge case is fine

Where: MockEngine.cs:Reset()

// Before: while (_callHistory.TryDequeue(out _)) { }
// After:  Volatile.Write(ref _callHistory, null);

Clean improvement. Calls recorded between the Volatile.Write(null) and LazyInitializer.EnsureInitialized reinitializing in RecordCall go into the new queue — which is correct post-reset behavior. The old drain-loop had the same window. ✓

---

What's excellent in this PR

• IMockObject interface for reverse lookup is architecturally superior to ConditionalWeakTable. It trades a GC-handle per mock for a single interface property slot, which is always cheaper. The explicit interface implementation keeps it invisible to users. This is the right design.
• Deferred allocation in MockEngine is solid — most mocks have zero calls or one behavior, so pre-allocating Dictionary and ConcurrentQueue was wasteful. The LazyInitializer.EnsureInitialized usage in RecordCall is the correct pattern.
• Factory delegate covariance (Func<MockBehavior, Mock<T>> stored as Func<MockBehavior, object>) is a correct use of covariance that eliminates one closure allocation per registered factory type — clean.
• Direct _behaviorLock! in read paths is sound reasoning given that _behaviors being non-null structurally implies the lock was already initialized.

---

Summary: Issues 1 and 2 are correctness concerns worth addressing before merge — Issue 1 especially represents a behavioral regression in concurrent scenarios. Issue 3 is a minor design hardening suggestion. The overall direction and the majority of changes are well-executed.

[claude]

Code Review: perf/mock-optimizations

Previous reviews correctly identified the main concerns. The fix commit addressed two of three:

✅ Addressed: Volatile.Read for lock-free null checks

MethodSetup.GetNextBehavior(), GetEventRaises(), and OutRefAssignments now use Volatile.Read before taking the lock. This correctly fixes the memory-model concern for ARM/AOT targets.

✅ Addressed: IMockObject.MockWrapper design rationale documented

The XML doc explains the set-once invariant. Given that generated code in user assemblies must implement this interface, keeping it public is necessary — the documentation is the right approach here.

⚠️ Partially addressed: EnsureSetup() thread safety — documentation-only fix

The fix commit added a comment documenting that MockMethodCall is not thread-safe by design. I want to probe this assumption.

// Not thread-safe: MockMethodCall instances are owned by a single test and never shared across threads.
private MethodSetupBuilder<TReturn> EnsureSetup()
{
    if (_builder is { } b) return b;
    var setup = new MethodSetup(_memberId, _matchers, _memberName);
    _engine.AddSetup(setup);
    return _builder = new MethodSetupBuilder<TReturn>(setup);
}

The "single test thread" invariant holds for typical setup code. However, there's one scenario that could violate it: auto-mocked return values. When a loose mock returns an auto-mocked interface for an unconfigured method, the mock engine creates a new mock and the corresponding MockMethodCall objects lazily. If the method is invoked concurrently from different threads (e.g., the SUT itself is multi-threaded), two threads can race on EnsureSetup().

This would register duplicate MethodSetup entries for the same member — the engine's "last wins" dispatch would still produce correct results (the duplicate is shadowed), but VerifyAll / VerifyNoOtherCalls counts would be off because InvokeCount is tracked per-MethodSetup instance.

Concrete suggestion: Since LazyInitializer.EnsureInitialized is already used elsewhere in this PR (e.g., MockEngine.RecordCall), it would be consistent and sufficient to apply it here too. It avoids the Lazy<T> object allocation while providing the safety net:

private MethodSetupBuilder<TReturn> EnsureSetup() =>
    LazyInitializer.EnsureInitialized(ref _builder, () =>
    {
        var setup = new MethodSetup(_memberId, _matchers, _memberName);
        _engine.AddSetup(setup);
        return new MethodSetupBuilder<TReturn>(setup);
    })!;

The closure allocation here is bounded — one per MockMethodCall instance, not per invocation — and is far cheaper than the original Lazy<T>.

---

New observation: MockEngine._setupLock vs. MethodSetup._behaviorLock inconsistency in deferred init

MockEngine._setupLock is a readonly Lock (eagerly allocated), while MethodSetup._behaviorLock is Lock? (lazily allocated). The PR defers all other collections in both types. Lock in .NET 9+ is a value type (struct wrapping a Monitor), so the eager allocation in MockEngine carries no heap cost — it's just inline storage in the object. The inconsistency isn't a bug but a quick comment like // Lock is a struct; no heap alloc where _setupLock is declared would make the intent clear to future readers.

---

What's working well

• ConditionalWeakTable → IMockObject is the strongest change. Eliminating GC-handle registration per mock object is a real win, and the explicit interface implementation keeps the plumbing invisible to users. Well done.
• Deferred collections in MockEngine are clean and effective. The LazyInitializer.EnsureInitialized pattern used for RecordCall is the right idiom.
• Delegate covariance for factory registration is a correct, zero-overhead improvement.
• Snapshot tests fully updated — the diffs are consistent and correct across all 16 snapshot files.

Overall this PR is in good shape. The one actionable item is replacing the EnsureSetup() documentation workaround with a LazyInitializer.EnsureInitialized call to eliminate the concurrent auto-mock race condition.

[claude]

Code Review — perf: optimize TUnit.Mocks hot paths

Overall this is a well-motivated set of micro-optimisations. The goals are sensible and the implementation is mostly correct. I have a few observations ranging from potential correctness issues to design questions worth considering before merging.

---

1. RecordCall — double-checked locking without a Volatile.Read fence (correctness concern)

// MockEngine.cs
var history = _callHistory ?? LazyInitializer.EnsureInitialized(ref _callHistory)!;
history.Enqueue(record);

The left-hand side of ?? reads _callHistory as a plain field read with no memory barrier. On weakly-ordered architectures (ARM, for example) the runtime can reorder this so that a stale null is seen even after another thread has already initialised it, causing EnsureInitialized to be entered redundantly. ConcurrentQueue is safe to enqueue from multiple threads, so the redundancy is not a data-corruption issue here — but the intent is clearly "read once and re-use the reference". Using Volatile.Read for the fast-path read would be consistent with how _autoMockCache, _eventSubscriptions, etc. are already read elsewhere in the same file:

var history = Volatile.Read(ref _callHistory) ?? LazyInitializer.EnsureInitialized(ref _callHistory)!;

---

2. GetEventRaises / GetNextBehavior — Volatile read then bare lock (correctness concern)

// MethodSetup.cs — GetEventRaises
if (Volatile.Read(ref _eventRaises) is null)
    return [];

lock (_behaviorLock!)   // ← bare dereference; no Volatile.Read
{
    return _eventRaises!.ToList();
}

The Volatile.Read on _eventRaises guards the early-exit correctly. However, when the fast path is not taken, _behaviorLock is read without a fence. If a concurrent writer calls AddEventRaise / AddBehavior at exactly this moment (which calls EnsureBehaviorLock() → LazyInitializer.EnsureInitialized), this thread may see a stale null and NullReferenceException, despite the ! suppression telling the compiler "trust me, it's non-null". The comment in the PR description says "lock is guaranteed to exist", but that guarantee relies on the memory model in a way that is not obvious. A Volatile.Read on _behaviorLock here (or reusing the EnsureBehaviorLock() helper for consistency) would make the contract explicit.

---

3. IMockObject.MockWrapper — mutable public setter on a public interface (design concern)

The new interface is correctly hidden from IDE autocompletion via [EditorBrowsable(Never)]. However, the property has a public set accessor on a public interface. Anyone who obtains a reference typed as IMockObject can overwrite the wrapper reference:

// Valid at compile time today:
((IMockObject)myMock).MockWrapper = someOtherMock;

Because the interface must be public (generated code in the user's assembly implements it), one option that would eliminate the mutable surface without breaking the generated-code assignment scenario is to use an explicit initialiser method instead of a property setter:

public interface IMockObject
{
    [EditorBrowsable(EditorBrowsableState.Never)]
    void SetMockWrapper(IMock wrapper);   // write-only by design
    [EditorBrowsable(EditorBrowsableState.Never)]
    IMock? GetMockWrapper();             // read needed by Mock.Get
}

Alternatively, if keeping a property, making the setter internal and using [InternalsVisibleTo] for the user's generated assembly (which already references the TUnit.Mocks assembly) is another option — though that adds configuration overhead. The current approach is pragmatic and the XML doc comment clearly says "should not be reassigned", so this is a low-severity concern.

---

4. LazyInitializer.EnsureInitialized with three refs — thread-safety semantics change from Lazy<T> (correctness concern)

The three-argument overload (ref T value, ref bool initialized, ref object syncRoot) provides ExecutionAndPublication semantics (equivalent to the Lazy<T> default). However, if the factory delegate throws, LazyInitializer does not cache the exception — on the next call the factory will be retried. Lazy<T> with ExecutionAndPublication mode also does not cache exceptions by default (since .NET Framework 4.5 this changed), so the observable behaviour is the same. This is worth confirming in the context of AddSetup side-effects: if the factory throws halfway through (e.g., AddSetup throws), the setup will not have been registered, and the next call will attempt registration again. This is probably acceptable, but it is a subtle behavioural difference from the old approach that should be documented or tested explicitly.

---

5. Reset() — _setupsByMember = null inside the lock, but readers use optimistic no-lock fast paths

// Reset()
lock (_setupLock) { _setupsByMember = null; ... }

// FindMatchingSetup() (also under _setupLock — fine)
lock (_setupLock) { if (_setupsByMember is not { } setupDict || ...) }

FindMatchingSetup takes the same _setupLock, so there is no race with Reset there. However GetCallsFor, GetUnverifiedCalls, and GetDiagnostics read _callHistory without any lock after Reset assigns null via Volatile.Write. The code is correct for _callHistory because Volatile.Write provides a release fence and the subsequent Volatile.Read (or plain field read on the reader side) provides an acquire fence on platforms that require it. The pattern is internally consistent but mixing lock-protected fields with Volatile-protected fields within the same class is a source of future confusion; a brief comment in the file noting which fields use which synchronisation strategy would help maintainability.

---

6. _builderLock as object? — per-instance boxing on first use (minor)

Each generated mock call class now carries three fields: _builder, _builderInitialized, _builderLock. The lock object is null until EnsureInitialized needs it, at which point the framework allocates a plain object for the monitor. Since the PR description explicitly targets allocation reduction, it is worth noting that the three-field pattern allocates a heap object on first contended access (boxing the sync root), whereas Lazy<T> always allocated its state object. For the typical single-threaded test scenario the three-field pattern is an improvement. If zero-allocation even under contention is desired, Lock (the new .NET 9+ struct value type, already used for _setupLock and _behaviorLock elsewhere) could be used here too — though as a non-nullable struct it cannot be deferred.

---

7. Source-generator EmitEnsureSetup — raw string alignment in generated output (style/readability)

// MockMembersBuilder.cs
writer.AppendLine($"private {builderType} EnsureSetup() =>");
writer.AppendLine($"    global::System.Threading.LazyInitializer.EnsureInitialized(ref _builder, ref _builderInitialized, ref _builderLock, () =>");
writer.AppendLine("    {");
writer.AppendLine("        var setup = new global::TUnit.Mocks.Setup.MethodSetup(_memberId, _matchers, _memberName);");

The indentation is hardcoded with leading spaces rather than using the CodeWriter's indent machinery (which is used everywhere else in the same builder). This makes the generated code look correct in isolation but will be misaligned if the surrounding block depth changes in the future. Using writer.Indent() / writer.Dedent() (or whatever the CodeWriter abstraction provides) would be more resilient.

---

8. Delegate covariance assignment — minor type-system subtlety worth a comment

// Mock.cs
_factories[typeof(T)] = factory;   // Func<MockBehavior, Mock<T>> stored as Func<MockBehavior, object>

This relies on the fact that Mock<T> is a reference type, so Func<MockBehavior, Mock<T>> is implicitly covariant to Func<MockBehavior, object> in C#. This is valid and a nice clean-up. A one-line comment noting the covariance reliance would prevent a future reader from questioning whether the implicit conversion is safe.

---

Summary

#	Severity	Category
1	Medium	Correctness (missing Volatile fence in RecordCall fast path)
2	Medium	Correctness (bare lock dereference after Volatile read)
3	Low	Design (mutable public setter on public interface)
4	Low	Correctness (LazyInitializer exception-retry semantics)
5	Low	Maintainability (mixed sync strategies, no comment)
6	Info	Allocation (per-instance object allocation for sync root)
7	Info	Code quality (hardcoded indent in code generator)
8	Info	Readability (implicit delegate covariance)

Items 1 and 2 are the ones most worth addressing before merging. The rest are informational or low-priority.