From 7ed8433f613476eafe15b2409c9097b6948babd6 Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Thu, 28 May 2026 23:30:07 +0100 Subject: [PATCH 1/2] perf(sourcegen): real InterfaceCache + single-pass attribute classification Closes #6106 - InterfaceCache: add a real per-type cache of implemented interface names backed by a ConditionalWeakTable>. Keying on the symbol ties each entry's lifetime to the symbol/Compilation, so the cache is reclaimed on collection and cannot leak symbols across compilations in long IDE sessions (the failure mode a static dictionary would cause). ImplementsInterface now does an O(1) set lookup instead of walking AllInterfaces on every call. - TestMetadataGenerator.CollectConcreteInstantiations: replace the 6-8 repeated `.Where(...).ToArray()` scans over MethodAttributes with a single classification foreach partitioning into null-init buckets. Buckets are independent so multi-membership (e.g. MethodDataSourceAttribute is also an IDataSourceAttribute) is preserved exactly. - MethodExtensions.GetTestAttribute: FirstOrDefault(lambda) -> foreach + early return, hoisting the well-known name out of the loop. - CodeWriter.GetIndentation: drop Enumerable.Repeat on the cache-miss path; fast path for single-space indent, StringBuilder loop otherwise. Output is byte-identical. Generated source output is unchanged - verified by the snapshot tests in TUnit.Core.SourceGenerator.Tests (116 passed, 1 pre-existing skip, 0 failed, no .received.txt produced). No ISymbol is stored in incremental pipeline nodes; only post-Collect consumers touch symbols. --- TUnit.Core.SourceGenerator/CodeWriter.cs | 20 ++++++- .../Extensions/MethodExtensions.cs | 14 ++++- .../Generators/TestMetadataGenerator.cs | 58 +++++++++++++++---- .../Helpers/InterfaceCache.cs | 36 +++++++++--- 4 files changed, 105 insertions(+), 23 deletions(-) diff --git a/TUnit.Core.SourceGenerator/CodeWriter.cs b/TUnit.Core.SourceGenerator/CodeWriter.cs index 04163d5e28..f8d77b5bcf 100644 --- a/TUnit.Core.SourceGenerator/CodeWriter.cs +++ b/TUnit.Core.SourceGenerator/CodeWriter.cs @@ -48,7 +48,25 @@ public ICodeWriter SetIndentLevel(int level) private string GetIndentation(int level) { var key = (_indentString, level); - return _indentCache.GetOrAdd(key, static k => string.Concat(Enumerable.Repeat(k.Item1, k.Item2))); + return _indentCache.GetOrAdd(key, static k => + { + var (indent, count) = k; + + // Fast path: a single-space indent (or any whitespace-only indent) can be built + // directly without allocating an intermediate sequence. + if (indent == " ") + { + return new string(' ', count); + } + + var builder = new StringBuilder(indent.Length * count); + for (var i = 0; i < count; i++) + { + builder.Append(indent); + } + + return builder.ToString(); + }); } /// diff --git a/TUnit.Core.SourceGenerator/Extensions/MethodExtensions.cs b/TUnit.Core.SourceGenerator/Extensions/MethodExtensions.cs index d28002c83e..3d067ce47c 100644 --- a/TUnit.Core.SourceGenerator/Extensions/MethodExtensions.cs +++ b/TUnit.Core.SourceGenerator/Extensions/MethodExtensions.cs @@ -19,8 +19,16 @@ public static AttributeData GetRequiredTestAttribute(this IMethodSymbol methodSy return null; } - return attributes - .FirstOrDefault(x => x.AttributeClass?.BaseType?.GloballyQualified() - == WellKnownFullyQualifiedClassNames.BaseTestAttribute.WithGlobalPrefix); + var baseTestAttribute = WellKnownFullyQualifiedClassNames.BaseTestAttribute.WithGlobalPrefix; + + foreach (var attribute in attributes) + { + if (attribute.AttributeClass?.BaseType?.GloballyQualified() == baseTestAttribute) + { + return attribute; + } + } + + return null; } } diff --git a/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs b/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs index 92f353e4f2..8c8d40dc14 100644 --- a/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs +++ b/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs @@ -620,9 +620,46 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt }); } - var methodArgumentsAttributes = testMethod.MethodAttributes - .Where(a => a.AttributeClass?.Name == "ArgumentsAttribute") - .ToArray(); + // Single-pass classification of method attributes into named buckets, replacing the + // repeated `.Where(...).ToArray()` scans that previously walked MethodAttributes 6-8 times. + List? methodArgumentsBucket = null; + List? methodDataSourceBucket = null; + List? typedDataSourceBucket = null; + List? generateGenericTestBucket = null; + + foreach (var attribute in testMethod.MethodAttributes) + { + var attributeClass = attribute.AttributeClass; + if (attributeClass is null) + { + continue; + } + + // Buckets are independent: an attribute may match more than one (e.g. + // MethodDataSourceAttribute is also an IDataSourceAttribute), matching the original + // behaviour where each `.Where(...)` scan was evaluated separately. + switch (attributeClass.Name) + { + case "ArgumentsAttribute": + (methodArgumentsBucket ??= []).Add(attribute); + break; + case "MethodDataSourceAttribute": + (methodDataSourceBucket ??= []).Add(attribute); + break; + } + + if (DataSourceAttributeHelper.IsDataSourceAttribute(attributeClass)) + { + (typedDataSourceBucket ??= []).Add(attribute); + } + + if (attributeClass.IsOrInherits("global::TUnit.Core.GenerateGenericTestAttribute")) + { + (generateGenericTestBucket ??= []).Add(attribute); + } + } + + var methodArgumentsAttributes = methodArgumentsBucket ?? []; var classArgumentsAttributes = testMethod.IsGenericType ? testMethod.TypeSymbol.GetAttributes() @@ -679,7 +716,7 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt } // Handle generic classes with non-generic methods that have method-level Arguments - if (testMethod is { IsGenericType: true, IsGenericMethod: false } && methodArgumentsAttributes.Length > 0) + if (testMethod is { IsGenericType: true, IsGenericMethod: false } && methodArgumentsAttributes.Count > 0) { foreach (var methodArgAttr in methodArgumentsAttributes) { @@ -692,7 +729,7 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt } // Process typed data source attributes - foreach (var dataSourceAttr in testMethod.MethodAttributes.Where(a => DataSourceAttributeHelper.IsDataSourceAttribute(a.AttributeClass))) + foreach (var dataSourceAttr in typedDataSourceBucket ?? Enumerable.Empty()) { var inferredTypes = InferTypesFromDataSourceAttribute(testMethod.MethodSymbol, dataSourceAttr); if (inferredTypes is { Length: > 0 }) @@ -714,7 +751,7 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt // Process MethodDataSource attributes for generic classes (non-generic methods) if (testMethod is { IsGenericType: true, IsGenericMethod: false }) { - foreach (var mdsAttr in testMethod.MethodAttributes.Where(a => a.AttributeClass?.Name == "MethodDataSourceAttribute")) + foreach (var mdsAttr in methodDataSourceBucket ?? Enumerable.Empty()) { var inferredTypes = InferClassTypesFromMethodDataSource(testMethod, mdsAttr); if (inferredTypes is { Length: > 0 }) @@ -734,7 +771,7 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt // Process MethodDataSource attributes for generic methods if (testMethod.IsGenericMethod) { - foreach (var mdsAttr in testMethod.MethodAttributes.Where(a => a.AttributeClass?.Name == "MethodDataSourceAttribute")) + foreach (var mdsAttr in methodDataSourceBucket ?? Enumerable.Empty()) { var inferredTypes = InferTypesFromMethodDataSource(testMethod, mdsAttr); if (inferredTypes is { Length: > 0 }) @@ -754,7 +791,7 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt { if (testMethod.IsGenericMethod) { - foreach (var methodArgAttr in testMethod.MethodAttributes.Where(a => a.AttributeClass?.Name == "ArgumentsAttribute")) + foreach (var methodArgAttr in methodArgumentsAttributes) { var methodInferredTypes = InferTypesFromArgumentsAttribute(testMethod.MethodSymbol, methodArgAttr, compilation); if (methodInferredTypes is { Length: > 0 }) @@ -774,9 +811,8 @@ void TryAddInstantiation(ITypeSymbol[] typeArguments, AttributeData? specificAtt // Process GenerateGenericTest attributes // GenerateGenericTestAttribute takes params Type[] in its constructor, so extract from constructor args { - var methodGenericTestAttrs = testMethod.IsGenericMethod - ? testMethod.MethodAttributes - .Where(a => a.AttributeClass?.IsOrInherits("global::TUnit.Core.GenerateGenericTestAttribute") is true) + var methodGenericTestAttrs = testMethod.IsGenericMethod && generateGenericTestBucket is not null + ? generateGenericTestBucket .Select(ExtractTypeArgsFromGenerateGenericTestAttribute) .Where(t => t is { Length: > 0 }) .ToList() diff --git a/TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs b/TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs index 5951b48b56..1d61defe86 100644 --- a/TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs +++ b/TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs @@ -1,27 +1,47 @@ +using System.Collections.Immutable; +using System.Runtime.CompilerServices; using Microsoft.CodeAnalysis; using TUnit.Core.SourceGenerator.Extensions; namespace TUnit.Core.SourceGenerator.Helpers; /// -/// Caches interface implementation checks to avoid repeated AllInterfaces traversals +/// Caches interface implementation checks to avoid repeated AllInterfaces traversals. /// +/// +/// Cache entries are keyed by the itself and held in a +/// . This ties each entry's lifetime to the +/// symbol's lifetime, so the cache is reclaimed automatically when a +/// is collected. This avoids the cross-compilation symbol leak a long-lived static dictionary +/// would cause in extended IDE sessions. +/// public static class InterfaceCache { /// - /// Checks if a type implements a specific interface + /// Per-type cache of the globally-qualified names of every interface the type implements. /// - public static bool ImplementsInterface(ITypeSymbol type, string fullyQualifiedInterfaceName) + private static readonly ConditionalWeakTable> InterfaceNames = new(); + + private static ImmutableHashSet GetInterfaceNames(ITypeSymbol type) { - foreach (var i in type.AllInterfaces) + return InterfaceNames.GetValue(type, static t => { - if (i.GloballyQualified() == fullyQualifiedInterfaceName) + var builder = ImmutableHashSet.CreateBuilder(StringComparer.Ordinal); + foreach (var i in t.AllInterfaces) { - return true; + builder.Add(i.GloballyQualified()); } - } - return false; + return builder.ToImmutable(); + }); + } + + /// + /// Checks if a type implements a specific interface + /// + public static bool ImplementsInterface(ITypeSymbol type, string fullyQualifiedInterfaceName) + { + return GetInterfaceNames(type).Contains(fullyQualifiedInterfaceName); } /// From 0cbb716b250c96b634fc8951caab090617b33b79 Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Thu, 28 May 2026 23:38:38 +0100 Subject: [PATCH 2/2] refactor(sourcegen): drop Roslyn-type caching, rename InterfaceCache -> InterfaceHelper Reverts the ConditionalWeakTable interface-name cache: caching Roslyn symbols in persistent generator state is an anti-pattern (cross-compilation rooting risk), and AllInterfaces is already lazily cached on the symbol by Roslyn, so the win was marginal. The class no longer caches, so the name is corrected to InterfaceHelper. Keeps the real wins from this PR (single-pass attribute classification, FirstOrDefault/Enumerable.Repeat LINQ removals). Snapshots unchanged. --- .../Helpers/DataSourceAttributeHelper.cs | 10 ++--- .../Extensions/AttributeDataExtensions.cs | 9 ++--- .../Generators/TestMetadataGenerator.cs | 7 ++-- .../{InterfaceCache.cs => InterfaceHelper.cs} | 38 +++++-------------- 4 files changed, 20 insertions(+), 44 deletions(-) rename TUnit.Core.SourceGenerator/Helpers/{InterfaceCache.cs => InterfaceHelper.cs} (64%) diff --git a/TUnit.Core.SourceGenerator/CodeGenerators/Helpers/DataSourceAttributeHelper.cs b/TUnit.Core.SourceGenerator/CodeGenerators/Helpers/DataSourceAttributeHelper.cs index 935fec6b69..729f89dfb2 100644 --- a/TUnit.Core.SourceGenerator/CodeGenerators/Helpers/DataSourceAttributeHelper.cs +++ b/TUnit.Core.SourceGenerator/CodeGenerators/Helpers/DataSourceAttributeHelper.cs @@ -12,8 +12,8 @@ public static bool IsDataSourceAttribute(INamedTypeSymbol? attributeClass) return false; } - // Check if the attribute implements IDataSourceAttribute (using cache) - return InterfaceCache.ImplementsInterface(attributeClass, "global::TUnit.Core.IDataSourceAttribute"); + // Check if the attribute implements IDataSourceAttribute + return InterfaceHelper.ImplementsInterface(attributeClass, "global::TUnit.Core.IDataSourceAttribute"); } public static bool IsTypedDataSourceAttribute(INamedTypeSymbol? attributeClass) @@ -23,8 +23,8 @@ public static bool IsTypedDataSourceAttribute(INamedTypeSymbol? attributeClass) return false; } - // Check if the attribute implements ITypedDataSourceAttribute (using cache) - return InterfaceCache.ImplementsGenericInterface(attributeClass, "global::TUnit.Core.ITypedDataSourceAttribute`1"); + // Check if the attribute implements ITypedDataSourceAttribute + return InterfaceHelper.ImplementsGenericInterface(attributeClass, "global::TUnit.Core.ITypedDataSourceAttribute`1"); } public static ITypeSymbol? GetTypedDataSourceType(INamedTypeSymbol? attributeClass) @@ -34,7 +34,7 @@ public static bool IsTypedDataSourceAttribute(INamedTypeSymbol? attributeClass) return null; } - var typedInterface = InterfaceCache.GetGenericInterface(attributeClass, "global::TUnit.Core.ITypedDataSourceAttribute`1"); + var typedInterface = InterfaceHelper.GetGenericInterface(attributeClass, "global::TUnit.Core.ITypedDataSourceAttribute`1"); return typedInterface?.TypeArguments.FirstOrDefault(); } diff --git a/TUnit.Core.SourceGenerator/Extensions/AttributeDataExtensions.cs b/TUnit.Core.SourceGenerator/Extensions/AttributeDataExtensions.cs index 68348a2c4b..282b1a4369 100644 --- a/TUnit.Core.SourceGenerator/Extensions/AttributeDataExtensions.cs +++ b/TUnit.Core.SourceGenerator/Extensions/AttributeDataExtensions.cs @@ -22,8 +22,7 @@ public static bool IsDataSourceAttribute(this AttributeData? attributeData) return false; } - // Use InterfaceCache instead of AllInterfaces.Any() for better performance - return InterfaceCache.ImplementsInterface(attributeData.AttributeClass, + return InterfaceHelper.ImplementsInterface(attributeData.AttributeClass, WellKnownFullyQualifiedClassNames.IDataSourceAttribute.WithGlobalPrefix); } @@ -34,8 +33,7 @@ public static bool IsTypedDataSourceAttribute(this AttributeData? attributeData) return false; } - // Use InterfaceCache instead of AllInterfaces.Any() for better performance - return InterfaceCache.ImplementsGenericInterface(attributeData.AttributeClass, + return InterfaceHelper.ImplementsGenericInterface(attributeData.AttributeClass, WellKnownFullyQualifiedClassNames.ITypedDataSourceAttribute.WithGlobalPrefix + "`1"); } @@ -46,8 +44,7 @@ public static bool IsTypedDataSourceAttribute(this AttributeData? attributeData) return null; } - // Use InterfaceCache instead of AllInterfaces.FirstOrDefault() for better performance - var typedInterface = InterfaceCache.GetGenericInterface(attributeData.AttributeClass, + var typedInterface = InterfaceHelper.GetGenericInterface(attributeData.AttributeClass, WellKnownFullyQualifiedClassNames.ITypedDataSourceAttribute.WithGlobalPrefix + "`1"); return typedInterface?.TypeArguments.FirstOrDefault(); diff --git a/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs b/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs index 8c8d40dc14..89e74e7ad3 100644 --- a/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs +++ b/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs @@ -2004,8 +2004,7 @@ private static void GeneratePropertyDataSourceFactory(CodeWriter writer, IProper private static bool IsAsyncEnumerable(ITypeSymbol type) { - // Use cached interface check - return InterfaceCache.IsAsyncEnumerable(type); + return InterfaceHelper.IsAsyncEnumerable(type); } private static bool IsTask(ITypeSymbol type) @@ -2017,8 +2016,8 @@ private static bool IsTask(ITypeSymbol type) private static bool IsEnumerable(ITypeSymbol type) { - // Use cached interface check (already handles string exclusion) - return InterfaceCache.IsEnumerable(type); + // Already handles string exclusion + return InterfaceHelper.IsEnumerable(type); } private static void WriteTypedConstant(CodeWriter writer, TypedConstant constant) diff --git a/TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs b/TUnit.Core.SourceGenerator/Helpers/InterfaceHelper.cs similarity index 64% rename from TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs rename to TUnit.Core.SourceGenerator/Helpers/InterfaceHelper.cs index 1d61defe86..84dcf5b7ae 100644 --- a/TUnit.Core.SourceGenerator/Helpers/InterfaceCache.cs +++ b/TUnit.Core.SourceGenerator/Helpers/InterfaceHelper.cs @@ -1,47 +1,27 @@ -using System.Collections.Immutable; -using System.Runtime.CompilerServices; using Microsoft.CodeAnalysis; using TUnit.Core.SourceGenerator.Extensions; namespace TUnit.Core.SourceGenerator.Helpers; /// -/// Caches interface implementation checks to avoid repeated AllInterfaces traversals. +/// Helpers for interface implementation checks over . /// -/// -/// Cache entries are keyed by the itself and held in a -/// . This ties each entry's lifetime to the -/// symbol's lifetime, so the cache is reclaimed automatically when a -/// is collected. This avoids the cross-compilation symbol leak a long-lived static dictionary -/// would cause in extended IDE sessions. -/// -public static class InterfaceCache +public static class InterfaceHelper { /// - /// Per-type cache of the globally-qualified names of every interface the type implements. + /// Checks if a type implements a specific interface /// - private static readonly ConditionalWeakTable> InterfaceNames = new(); - - private static ImmutableHashSet GetInterfaceNames(ITypeSymbol type) + public static bool ImplementsInterface(ITypeSymbol type, string fullyQualifiedInterfaceName) { - return InterfaceNames.GetValue(type, static t => + foreach (var i in type.AllInterfaces) { - var builder = ImmutableHashSet.CreateBuilder(StringComparer.Ordinal); - foreach (var i in t.AllInterfaces) + if (i.GloballyQualified() == fullyQualifiedInterfaceName) { - builder.Add(i.GloballyQualified()); + return true; } + } - return builder.ToImmutable(); - }); - } - - /// - /// Checks if a type implements a specific interface - /// - public static bool ImplementsInterface(ITypeSymbol type, string fullyQualifiedInterfaceName) - { - return GetInterfaceNames(type).Contains(fullyQualifiedInterfaceName); + return false; } ///