diff --git a/TUnit.Core.SourceGenerator/CodeGenerators/Writers/AttributeWriter.cs b/TUnit.Core.SourceGenerator/CodeGenerators/Writers/AttributeWriter.cs index 052642e926..d3e0aa4c44 100644 --- a/TUnit.Core.SourceGenerator/CodeGenerators/Writers/AttributeWriter.cs +++ b/TUnit.Core.SourceGenerator/CodeGenerators/Writers/AttributeWriter.cs @@ -6,7 +6,7 @@ namespace TUnit.Core.SourceGenerator.CodeGenerators.Writers; -public class AttributeWriter(Compilation compilation) +public class AttributeWriter(Compilation compilation, TUnit.Core.SourceGenerator.Helpers.WellKnownTypes wellKnownTypes) { private readonly Dictionary _attributeObjectInitializerCache = new(); @@ -30,7 +30,7 @@ public void WriteAttributes(ICodeWriter sourceCodeWriter, // Skip framework-specific attributes when targeting older frameworks // We determine this by checking if we can compile the attribute - if (ShouldSkipFrameworkSpecificAttribute(compilation, attributeData)) + if (ShouldSkipFrameworkSpecificAttribute(attributeData)) { continue; } @@ -190,7 +190,7 @@ public static void WriteAttributeWithoutSyntax(ICodeWriter sourceCodeWriter, Att } - private static bool ShouldSkipFrameworkSpecificAttribute(Compilation compilation, AttributeData attributeData) + private bool ShouldSkipFrameworkSpecificAttribute(AttributeData attributeData) { if (attributeData.AttributeClass == null) { @@ -199,14 +199,14 @@ private static bool ShouldSkipFrameworkSpecificAttribute(Compilation compilation // Generic approach: Check if the attribute type is actually available in the target compilation // This works by seeing if we can resolve the type from the compilation's references - var fullyQualifiedName = attributeData.AttributeClass.ToDisplayString(); + var fullyQualifiedName = wellKnownTypes.GetDisplayString(attributeData.AttributeClass); // Check if this is a system/runtime attribute that might not exist on all frameworks if (fullyQualifiedName.StartsWith("System.") || fullyQualifiedName.StartsWith("Microsoft.")) { // Try to get the type from the compilation // If it doesn't exist in the compilation's references, we should skip it - var typeSymbol = compilation.GetTypeByMetadataName(fullyQualifiedName); + var typeSymbol = wellKnownTypes.TryGet(fullyQualifiedName); // If the type doesn't exist in the compilation, skip it if (typeSymbol == null) @@ -240,14 +240,14 @@ private static bool IsNullableAttribute(string fullyQualifiedName) fullyQualifiedName.Contains("NullablePublicOnlyAttribute"); } - private static bool ShouldSkipCompilerInternalAttribute(AttributeData attributeData) + private bool ShouldSkipCompilerInternalAttribute(AttributeData attributeData) { if (attributeData.AttributeClass == null) { return false; } - var fullyQualifiedName = attributeData.AttributeClass.ToDisplayString(); + var fullyQualifiedName = wellKnownTypes.GetDisplayString(attributeData.AttributeClass); // Skip compiler-internal attributes that should never be re-emitted // System.Runtime.CompilerServices contains compiler-generated and structural metadata attributes diff --git a/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs b/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs index 3e99cce971..eb10e428d2 100644 --- a/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs +++ b/TUnit.Core.SourceGenerator/Generators/TestMetadataGenerator.cs @@ -31,10 +31,14 @@ public void Initialize(IncrementalGeneratorInitializationContext context) var compilationContext = context .CompilationProvider .Select(static (c, _) => - new CompilationContext( + { + var wellKnownTypes = new WellKnownTypes(c); + return new CompilationContext( (CSharpCompilation)c, - new AttributeWriter(c) - )); + new AttributeWriter(c, wellKnownTypes), + wellKnownTypes + ); + }); var testMethodsProvider = context.SyntaxProvider .ForAttributeWithMetadataName( diff --git a/TUnit.Core.SourceGenerator/Helpers/WellKnownTypes.cs b/TUnit.Core.SourceGenerator/Helpers/WellKnownTypes.cs new file mode 100644 index 0000000000..55c24be027 --- /dev/null +++ b/TUnit.Core.SourceGenerator/Helpers/WellKnownTypes.cs @@ -0,0 +1,50 @@ +using Microsoft.CodeAnalysis; + +namespace TUnit.Core.SourceGenerator.Helpers; + +public class WellKnownTypes(Compilation compilation) +{ + private readonly Dictionary _cachedTypes = new(); + private readonly Dictionary _cachedToDisplayString = new(SymbolEqualityComparer.Default); + + public string GetDisplayString(INamedTypeSymbol type) + { + if (_cachedToDisplayString.TryGetValue(type, out var displayString)) + { + return displayString; + } + + displayString = type.ToDisplayString(); + _cachedToDisplayString.Add(type, displayString); + + return displayString; + } + + public INamedTypeSymbol Get() => Get(typeof(T)); + + public INamedTypeSymbol Get(Type type) + { + if (type.IsConstructedGenericType) + { + type = type.GetGenericTypeDefinition(); + } + + return Get(type.FullName ?? throw new InvalidOperationException("Could not get name of type " + type)); + } + + public INamedTypeSymbol? TryGet(string typeFullName) + { + if (_cachedTypes.TryGetValue(typeFullName, out var typeSymbol)) + { + return typeSymbol; + } + + typeSymbol = compilation.GetTypeByMetadataName(typeFullName); + _cachedTypes.Add(typeFullName, typeSymbol); + + return typeSymbol; + } + + private INamedTypeSymbol Get(string typeFullName) => + TryGet(typeFullName) ?? throw new InvalidOperationException("Could not get type " + typeFullName); +} diff --git a/TUnit.Core.SourceGenerator/Models/TestMethodMetadata.cs b/TUnit.Core.SourceGenerator/Models/TestMethodMetadata.cs index 0c0cc6bfbc..33616e1e7d 100644 --- a/TUnit.Core.SourceGenerator/Models/TestMethodMetadata.cs +++ b/TUnit.Core.SourceGenerator/Models/TestMethodMetadata.cs @@ -3,10 +3,11 @@ using Microsoft.CodeAnalysis.CSharp; using Microsoft.CodeAnalysis.CSharp.Syntax; using TUnit.Core.SourceGenerator.CodeGenerators.Writers; +using TUnit.Core.SourceGenerator.Helpers; namespace TUnit.Core.SourceGenerator.Models; -public record CompilationContext(CSharpCompilation Compilation, AttributeWriter AttributeWriter); +public record CompilationContext(CSharpCompilation Compilation, AttributeWriter AttributeWriter, WellKnownTypes WellKnownTypes); /// /// Contains all the metadata about a test method discovered by the source generator.