Skip to content

Commit e1f48e3

Browse files
TimothyMakkisonTimothyMakkison
committed
perf: cache GetAttributeObjectInitializer, make AttributeWriter stateful
1 parent 3b9dcaa commit e1f48e3

3 files changed

Lines changed: 62 additions & 30 deletions

File tree

‎TUnit.Core.SourceGenerator/CodeGenerators/Writers/AttributeWriter.cs‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@
66

77
namespace TUnit.Core.SourceGenerator.CodeGenerators.Writers;
88

9-
public class AttributeWriter
9+
public class AttributeWriter(Compilation compilation)
1010
{
11-
public static void WriteAttributes(ICodeWriter sourceCodeWriter, Compilation compilation,
11+
private readonly Dictionary<AttributeData, string> _attributeObjectInitializerCache = new();
12+
13+
public void WriteAttributes(ICodeWriter sourceCodeWriter,
1214
ImmutableArray<AttributeData> attributeDatas)
1315
{
1416
var attributesToWrite = new List<AttributeData>();
@@ -49,7 +51,7 @@ public static void WriteAttributes(ICodeWriter sourceCodeWriter, Compilation com
4951
{
5052
var attributeData = attributesToWrite[index];
5153

52-
WriteAttribute(sourceCodeWriter, compilation, attributeData);
54+
WriteAttribute(sourceCodeWriter, attributeData);
5355

5456
if (index != attributesToWrite.Count - 1)
5557
{
@@ -58,8 +60,7 @@ public static void WriteAttributes(ICodeWriter sourceCodeWriter, Compilation com
5860
}
5961
}
6062

61-
public static void WriteAttribute(ICodeWriter sourceCodeWriter, Compilation compilation,
62-
AttributeData attributeData)
63+
public void WriteAttribute(ICodeWriter sourceCodeWriter, AttributeData attributeData)
6364
{
6465
if (attributeData.ApplicationSyntaxReference is null)
6566
{
@@ -70,12 +71,23 @@ public static void WriteAttribute(ICodeWriter sourceCodeWriter, Compilation comp
7071
else
7172
{
7273
// For attributes from the current compilation, use the syntax-based approach
73-
sourceCodeWriter.Append(GetAttributeObjectInitializer(compilation, attributeData));
74+
sourceCodeWriter.Append(GetAttributeObjectInitializer(attributeData));
75+
}
76+
}
77+
78+
public string GetAttributeObjectInitializer(AttributeData attributeData)
79+
{
80+
if (_attributeObjectInitializerCache.TryGetValue(attributeData, out var initializer))
81+
{
82+
return initializer;
7483
}
84+
85+
initializer = GetAttributeObjectInitializerInner(compilation, attributeData);
86+
_attributeObjectInitializerCache.Add(attributeData, initializer);
87+
return initializer;
7588
}
7689

77-
public static string GetAttributeObjectInitializer(Compilation compilation,
78-
AttributeData attributeData)
90+
private static string GetAttributeObjectInitializerInner(Compilation compilation, AttributeData attributeData)
7991
{
8092
var sourceCodeWriter = new CodeWriter("", includeHeader: false);
8193

@@ -123,7 +135,6 @@ public static string GetAttributeObjectInitializer(Compilation compilation,
123135
return sourceCodeWriter.ToString();
124136
}
125137

126-
127138
private static string FormatConstructorArgument(Compilation compilation, AttributeArgumentSyntax attributeArgumentSyntax)
128139
{
129140
if (attributeArgumentSyntax.NameColon is not null)

‎TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs‎

Lines changed: 36 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -28,19 +28,31 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
2828
return !string.Equals(value, "false", StringComparison.OrdinalIgnoreCase);
2929
});
3030

31+
var compilationContext = context
32+
.CompilationProvider
33+
.Select(static (c, _) =>
34+
new CompilationContext(
35+
(CSharpCompilation)c,
36+
new AttributeWriter(c)
37+
));
38+
3139
var testMethodsProvider = context.SyntaxProvider
3240
.ForAttributeWithMetadataName(
3341
"TUnit.Core.TestAttribute",
3442
predicate: static (node, _) => node is MethodDeclarationSyntax,
35-
transform: static (ctx, _) => GetTestMethodMetadata(ctx))
43+
transform: static (ctx, _) => ctx)
44+
.Combine(compilationContext)
45+
.Select(static (ctx, _) => GetTestMethodMetadata(ctx.Left, ctx.Right))
3646
.Where(static m => m is not null)
3747
.Combine(enabledProvider);
3848

3949
var inheritsTestsClassesProvider = context.SyntaxProvider
4050
.ForAttributeWithMetadataName(
4151
"TUnit.Core.InheritsTestsAttribute",
4252
predicate: static (node, _) => node is ClassDeclarationSyntax,
43-
transform: static (ctx, _) => GetInheritsTestsClassMetadata(ctx))
53+
transform: static (ctx, _) => ctx)
54+
.Combine(compilationContext)
55+
.Select(static (ctx, _) => GetInheritsTestsClassMetadata(ctx.Left, ctx.Right))
4456
.Where(static m => m is not null)
4557
.Combine(enabledProvider);
4658

@@ -67,7 +79,7 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
6779
});
6880
}
6981

70-
private static InheritsTestsClassMetadata? GetInheritsTestsClassMetadata(GeneratorAttributeSyntaxContext context)
82+
private static InheritsTestsClassMetadata? GetInheritsTestsClassMetadata(GeneratorAttributeSyntaxContext context, CompilationContext compilationContext)
7183
{
7284
var classSyntax = (ClassDeclarationSyntax)context.TargetNode;
7385

@@ -85,11 +97,12 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
8597
{
8698
TypeSymbol = classSymbol,
8799
ClassSyntax = classSyntax,
88-
Context = context
100+
Context = context,
101+
CompilationContext = compilationContext
89102
};
90103
}
91104

92-
private static TestMethodMetadata? GetTestMethodMetadata(GeneratorAttributeSyntaxContext context)
105+
private static TestMethodMetadata? GetTestMethodMetadata(GeneratorAttributeSyntaxContext context, CompilationContext compilationContext)
93106
{
94107
var methodSyntax = (MethodDeclarationSyntax)context.TargetNode;
95108
var methodSymbol = context.TargetSymbol as IMethodSymbol;
@@ -121,6 +134,7 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
121134
LineNumber = lineNumber,
122135
TestAttribute = context.Attributes.First(),
123136
Context = context,
137+
CompilationContext = compilationContext,
124138
MethodSyntax = methodSyntax,
125139
IsGenericType = isGenericType,
126140
IsGenericMethod = isGenericMethod,
@@ -186,6 +200,7 @@ private static void GenerateInheritedTestSources(SourceProductionContext context
186200
LineNumber = lineNumber,
187201
TestAttribute = testAttribute,
188202
Context = classInfo.Context, // Use class context to access Compilation
203+
CompilationContext = classInfo.CompilationContext,
189204
MethodSyntax = null, // No syntax for inherited methods
190205
IsGenericType = typeForMetadata.IsGenericType,
191206
IsGenericMethod = (concreteMethod ?? method).IsGenericMethod,
@@ -458,7 +473,7 @@ private static void GenerateMetadata(CodeWriter writer, TestMethodMetadata testM
458473
.Concat(testMethod.TypeSymbol.ContainingAssembly.GetAttributes())
459474
.ToImmutableArray();
460475

461-
AttributeWriter.WriteAttributes(writer, compilation, attributes);
476+
testMethod.CompilationContext.AttributeWriter.WriteAttributes(writer, attributes);
462477

463478
writer.Unindent();
464479
writer.AppendLine("],");
@@ -504,7 +519,7 @@ private static void GenerateMetadataForConcreteInstantiation(CodeWriter writer,
504519
.Concat(testMethod.TypeSymbol.ContainingAssembly.GetAttributes())
505520
.ToImmutableArray();
506521

507-
AttributeWriter.WriteAttributes(writer, compilation, attributes);
522+
testMethod.CompilationContext.AttributeWriter.WriteAttributes(writer, attributes);
508523

509524
writer.Unindent();
510525
writer.AppendLine("],");
@@ -564,7 +579,7 @@ private static void GenerateDataSources(CodeWriter writer, TestMethodMetadata te
564579

565580
foreach (var attr in methodDataSources)
566581
{
567-
GenerateDataSourceAttribute(writer, compilation, attr, methodSymbol, typeSymbol);
582+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, attr, methodSymbol, typeSymbol);
568583
}
569584

570585
writer.Unindent();
@@ -584,7 +599,7 @@ private static void GenerateDataSources(CodeWriter writer, TestMethodMetadata te
584599

585600
foreach (var attr in classDataSources)
586601
{
587-
GenerateDataSourceAttribute(writer, compilation, attr, methodSymbol, typeSymbol);
602+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, attr, methodSymbol, typeSymbol);
588603
}
589604

590605
writer.Unindent();
@@ -595,7 +610,7 @@ private static void GenerateDataSources(CodeWriter writer, TestMethodMetadata te
595610
GeneratePropertyDataSources(writer, testMethod);
596611
}
597612

598-
private static void GenerateDataSourceAttribute(CodeWriter writer, Compilation compilation, AttributeData attr, IMethodSymbol methodSymbol, INamedTypeSymbol typeSymbol)
613+
private static void GenerateDataSourceAttribute(CodeWriter writer, CompilationContext compilationContext, AttributeData attr, IMethodSymbol methodSymbol, INamedTypeSymbol typeSymbol)
599614
{
600615
var attrClass = attr.AttributeClass;
601616
if (attrClass == null)
@@ -613,18 +628,18 @@ private static void GenerateDataSourceAttribute(CodeWriter writer, Compilation c
613628
{
614629
try
615630
{
616-
GenerateArgumentsAttributeWithParameterTypes(writer, compilation, attr, methodSymbol);
631+
GenerateArgumentsAttributeWithParameterTypes(writer, compilationContext.Compilation, attr, methodSymbol);
617632
}
618633
catch
619634
{
620635
// Fall back to default behavior if parameter type matching fails
621-
AttributeWriter.WriteAttribute(writer, compilation, attr);
636+
compilationContext.AttributeWriter.WriteAttribute(writer, attr);
622637
writer.AppendLine(",");
623638
}
624639
}
625640
else
626641
{
627-
AttributeWriter.WriteAttribute(writer, compilation, attr);
642+
compilationContext.AttributeWriter.WriteAttribute(writer, attr);
628643
writer.AppendLine(",");
629644
}
630645
}
@@ -1535,7 +1550,7 @@ private static void GeneratePropertyDataSources(CodeWriter writer, TestMethodMet
15351550
writer.AppendLine($"PropertyName = \"{property.Name}\",");
15361551
writer.AppendLine($"PropertyType = typeof({property.Type.GloballyQualified()}),");
15371552
writer.Append("DataSource = ");
1538-
GenerateDataSourceAttribute(writer, compilation, dataSourceAttr, testMethod.MethodSymbol, typeSymbol);
1553+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, dataSourceAttr, testMethod.MethodSymbol, typeSymbol);
15391554
writer.Unindent();
15401555
writer.AppendLine("},");
15411556
}
@@ -4805,7 +4820,7 @@ private static void GenerateConcreteMetadataWithFilteredDataSources(
48054820
writer.AppendLine("AttributeFactory = static () =>");
48064821
writer.AppendLine("[");
48074822
writer.Indent();
4808-
AttributeWriter.WriteAttributes(writer, compilation, filteredAttributes.ToImmutableArray());
4823+
testMethod.CompilationContext.AttributeWriter.WriteAttributes(writer, filteredAttributes.ToImmutableArray());
48094824
writer.Unindent();
48104825
writer.AppendLine("],");
48114826

@@ -4866,7 +4881,7 @@ private static void GenerateConcreteMetadataWithFilteredDataSources(
48664881

48674882
foreach (var attr in methodDataSources)
48684883
{
4869-
GenerateDataSourceAttribute(writer, compilation, attr, methodSymbol, concreteTypeSymbol);
4884+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, attr, methodSymbol, concreteTypeSymbol);
48704885
}
48714886

48724887
writer.Unindent();
@@ -4886,7 +4901,7 @@ private static void GenerateConcreteMetadataWithFilteredDataSources(
48864901

48874902
foreach (var attr in classDataSources)
48884903
{
4889-
GenerateDataSourceAttribute(writer, compilation, attr, methodSymbol, concreteTypeSymbol);
4904+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, attr, methodSymbol, concreteTypeSymbol);
48904905
}
48914906

48924907
writer.Unindent();
@@ -5133,7 +5148,7 @@ private static void GenerateConcreteTestMetadataForNonGeneric(
51335148
.Concat(testMethod.TypeSymbol.ContainingAssembly.GetAttributes())
51345149
.ToImmutableArray();
51355150

5136-
AttributeWriter.WriteAttributes(writer, compilation, attributes);
5151+
testMethod.CompilationContext.AttributeWriter.WriteAttributes(writer, attributes);
51375152

51385153
writer.Unindent();
51395154
writer.AppendLine("],");
@@ -5154,7 +5169,7 @@ private static void GenerateConcreteTestMetadataForNonGeneric(
51545169
writer.AppendLine("DataSources = new global::TUnit.Core.IDataSourceAttribute[]");
51555170
writer.AppendLine("{");
51565171
writer.Indent();
5157-
GenerateDataSourceAttribute(writer, compilation, methodDataSourceAttribute, testMethod.MethodSymbol, testMethod.TypeSymbol);
5172+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, methodDataSourceAttribute, testMethod.MethodSymbol, testMethod.TypeSymbol);
51585173
writer.Unindent();
51595174
writer.AppendLine("},");
51605175
}
@@ -5169,7 +5184,7 @@ private static void GenerateConcreteTestMetadataForNonGeneric(
51695184
writer.AppendLine("ClassDataSources = new global::TUnit.Core.IDataSourceAttribute[]");
51705185
writer.AppendLine("{");
51715186
writer.Indent();
5172-
GenerateDataSourceAttribute(writer, compilation, classDataSourceAttribute, testMethod.MethodSymbol, testMethod.TypeSymbol);
5187+
GenerateDataSourceAttribute(writer, testMethod.CompilationContext, classDataSourceAttribute, testMethod.MethodSymbol, testMethod.TypeSymbol);
51735188
writer.Unindent();
51745189
writer.AppendLine("},");
51755190
}
@@ -5282,5 +5297,6 @@ public class InheritsTestsClassMetadata
52825297
public required INamedTypeSymbol TypeSymbol { get; init; }
52835298
public required ClassDeclarationSyntax ClassSyntax { get; init; }
52845299
public GeneratorAttributeSyntaxContext Context { get; init; }
5300+
public required CompilationContext CompilationContext { get; init; }
52855301
}
52865302

‎TUnit.Core.SourceGenerator/Models/TestMethodMetadata.cs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
using System.Collections.Immutable;
22
using Microsoft.CodeAnalysis;
3+
using Microsoft.CodeAnalysis.CSharp;
34
using Microsoft.CodeAnalysis.CSharp.Syntax;
5+
using TUnit.Core.SourceGenerator.CodeGenerators.Writers;
46

57
namespace TUnit.Core.SourceGenerator.Models;
68

9+
public record CompilationContext(CSharpCompilation Compilation, AttributeWriter AttributeWriter);
10+
711
/// <summary>
812
/// Contains all the metadata about a test method discovered by the source generator.
913
/// </summary>
@@ -15,6 +19,7 @@ public class TestMethodMetadata : IEquatable<TestMethodMetadata>
1519
public required int LineNumber { get; init; }
1620
public required AttributeData TestAttribute { get; init; }
1721
public GeneratorAttributeSyntaxContext? Context { get; init; }
22+
public required CompilationContext CompilationContext { get; init; }
1823
public required MethodDeclarationSyntax? MethodSyntax { get; init; }
1924
public bool IsGenericType { get; init; }
2025
public bool IsGenericMethod { get; init; }
@@ -23,7 +28,7 @@ public class TestMethodMetadata : IEquatable<TestMethodMetadata>
2328
/// All attributes on the method, stored for later use during data combination generation
2429
/// </summary>
2530
public ImmutableArray<AttributeData> MethodAttributes { get; init; } = ImmutableArray<AttributeData>.Empty;
26-
31+
2732
/// <summary>
2833
/// The inheritance depth of this test method.
2934
/// 0 = method is declared directly in the test class

0 commit comments

Comments
 (0)