perf(assertions-source-gen): make assertion generators properly incremental

src/TUnit.Assertions.SourceGenerator/Generators/MethodAssertionGenerator.cs

@@ -4,6 +4,7 @@
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Microsoft.CodeAnalysis.Text;
 using TUnit.Assertions.SourceGenerator.Models;
 
 namespace TUnit.Assertions.SourceGenerator.Generators;
@@ -18,9 +19,10 @@
 public sealed class MethodAssertionGenerator : IIncrementalGenerator
 {
     public static string BuildAssertion = "MethodAssertionData";
+    public const string BuildAssertionGroup = "MethodAssertionGroup";
 
     private static readonly DiagnosticDescriptor MethodMustBeStaticRule = new DiagnosticDescriptor(
         id: "TUNITGEN001",
         title: "Method must be static",
         messageFormat: "Method '{0}' decorated with [GenerateAssertion] must be static",
         category: "TUnit.Assertions.SourceGenerator",
@@ -29,7 +31,7 @@
         description: "Methods decorated with [GenerateAssertion] must be static to be used in generated assertions.");
 
     private static readonly DiagnosticDescriptor MethodMustHaveParametersRule = new DiagnosticDescriptor(
         id: "TUNITGEN002",
         title: "Method must have at least one parameter",
         messageFormat: "Method '{0}' decorated with [GenerateAssertion] must have at least one parameter (the value to assert)",
         category: "TUnit.Assertions.SourceGenerator",
@@ -38,18 +40,18 @@
         description: "Methods decorated with [GenerateAssertion] must have at least one parameter representing the value being asserted.");
 
     private static readonly DiagnosticDescriptor UnsupportedReturnTypeRule = new DiagnosticDescriptor(
         id: "TUNITGEN003",
         title: "Unsupported return type",
         messageFormat: "Method '{0}' decorated with [GenerateAssertion] has unsupported return type '{1}'. Supported types are: bool, AssertionResult, AssertionResult<T>, Task<bool>, Task<AssertionResult>, Task<AssertionResult<T>>",
         category: "TUnit.Assertions.SourceGenerator",
         defaultSeverity: DiagnosticSeverity.Error,
         isEnabledByDefault: true,
         description: "Methods decorated with [GenerateAssertion] must return bool, AssertionResult, Task<bool>, or Task<AssertionResult>.");
 
     private static readonly DiagnosticDescriptor RefStructRequiresInliningRule = new DiagnosticDescriptor(
         id: "TUNITGEN004",
         title: "Ref struct parameter requires method body inlining",
         messageFormat: "Method '{0}' has ref struct parameter '{1}' of type '{2}'. Use InlineMethodBody = true and ensure the method has a single-expression or single-return-statement body",
         category: "TUnit.Assertions.SourceGenerator",
         defaultSeverity: DiagnosticSeverity.Error,
         isEnabledByDefault: true,
@@ -74,20 +76,51 @@
             .Where(x => x.Diagnostic != null)
             .Select((x, _) => x.Diagnostic!);
 
-        // Report diagnostics
-        context.RegisterSourceOutput(diagnostics, static (context, diagnostic) =>
+        // Report diagnostics. Each one is located in the current compilation's syntax tree so that
+        // #pragma and per-file .editorconfig severities still apply; an external-file location
+        // has no SourceTree to read them from. Pairing each diagnostic (not the collected array)
+        // with the compilation means this output only re-runs while there are diagnostics.
+        context.RegisterSourceOutput(diagnostics.Combine(context.CompilationProvider), static (context, data) =>
         {
-            context.ReportDiagnostic(diagnostic);
+            context.ReportDiagnostic(data.Left.ToDiagnostic(data.Right));
         });
 
+        // Group methods by containing type so an edit to one type only re-emits that type's file.
+        // The grouping itself is cheap; each group compares by value, so unchanged groups stay cached.
+        var methodGroups = methods.Collect()
+            .SelectMany(static (methods, _) => GroupByContainingType(methods))
+            .WithTrackingName(BuildAssertionGroup);
+
         // Generate assertion classes and extension methods
-        context.RegisterSourceOutput(methods.Collect(), static (context, methods) =>
+        context.RegisterSourceOutput(methodGroups, static (context, group) =>
         {
-            GenerateAssertions(context, methods);
+            GenerateAssertions(context, group);
         });
     }
 
-    private static (AssertionMethodData? Data, Diagnostic? Diagnostic) GetAssertionMethodData(
+    private static ImmutableArray<ContainingTypeGroup> GroupByContainingType(ImmutableArray<AssertionMethodData> methods)
+    {
+        if (methods.IsEmpty)
+        {
+            return ImmutableArray<ContainingTypeGroup>.Empty;
+        }
+
+        var builder = ImmutableArray.CreateBuilder<ContainingTypeGroup>();
+        foreach (var methodGroup in methods.GroupBy(m => m.Method.ContainingType?.FullContainingType))
+        {
+            var containingType = methodGroup.First().Method.ContainingType;
+            if (containingType == null)
+            {
+                continue;
+            }
+
+            builder.Add(new ContainingTypeGroup(containingType, methodGroup.ToImmutableEquatableArray()));
+        }
+
+        return builder.ToImmutable();
+    }
+
+    private static (AssertionMethodData? Data, DiagnosticInfo? Diagnostic) GetAssertionMethodData(
         GeneratorAttributeSyntaxContext context,
         CancellationToken cancellationToken)
     {
@@ -101,7 +134,7 @@
         // Validate method is static
         if (!methodSymbol.IsStatic)
         {
-            var diagnostic = Diagnostic.Create(
+            var diagnostic = DiagnosticInfo.Create(
                 MethodMustBeStaticRule,
                 location,
                 methodSymbol.Name);
@@ -111,7 +144,7 @@
         // Validate method has at least one parameter
         if (methodSymbol.Parameters.Length == 0)
         {
-            var diagnostic = Diagnostic.Create(
+            var diagnostic = DiagnosticInfo.Create(
                 MethodMustHaveParametersRule,
                 location,
                 methodSymbol.Name);
@@ -122,7 +155,7 @@
         var returnTypeInfo = AnalyzeReturnType(methodSymbol.ReturnType);
         if (returnTypeInfo == null)
         {
-            var diagnostic = Diagnostic.Create(
+            var diagnostic = DiagnosticInfo.Create(
                 UnsupportedReturnTypeRule,
                 location,
                 methodSymbol.Name,
@@ -301,7 +334,7 @@
         {
             if (param.IsRefStruct && string.IsNullOrEmpty(methodBody))
             {
-                var diagnostic = Diagnostic.Create(
+                var diagnostic = DiagnosticInfo.Create(
                     RefStructRequiresInliningRule,
                     location,
                     methodSymbol.Name,
@@ -546,86 +579,74 @@
 
     private static void GenerateAssertions(
         SourceProductionContext context,
-        ImmutableArray<AssertionMethodData> methods)
+        ContainingTypeGroup group)
     {
-        if (methods.IsEmpty)
-        {
-            return;
-        }
+        var containingType = group.ContainingType;
+        var methodGroup = group.Methods;
 
-        // Group by containing class to generate one file per class
-        foreach (var methodGroup in methods.GroupBy(m => m.Method.ContainingType?.FullContainingType))
-        {
-            var containingType = methodGroup.First().Method.ContainingType;
-            if (containingType == null)
-            {
-                continue;
-            }
+        var sourceBuilder = new StringBuilder();
+        // Always generate extension methods in TUnit.Assertions.Extensions namespace
+        // so they're available via implicit usings in consuming projects
+        var namespaceName = "TUnit.Assertions.Extensions";
 
-            var sourceBuilder = new StringBuilder();
-            // Always generate extension methods in TUnit.Assertions.Extensions namespace
-            // so they're available via implicit usings in consuming projects
-            var namespaceName = "TUnit.Assertions.Extensions";
+        // Get the original namespace where the helper methods are defined
+        var originalNamespace = containingType.ContainingNamespace;
 
-            // Get the original namespace where the helper methods are defined
-            var originalNamespace = containingType.ContainingNamespace;
+        // File header
+        sourceBuilder.AppendLine("// <auto-generated/>");
+        sourceBuilder.AppendLine("#pragma warning disable");
+        sourceBuilder.AppendLine("#nullable enable");
+        sourceBuilder.AppendLine();
+        sourceBuilder.AppendLine("using System;");
+        sourceBuilder.AppendLine("using System.Runtime.CompilerServices;");
+        sourceBuilder.AppendLine("using System.Threading.Tasks;");
+        sourceBuilder.AppendLine("using TUnit.Assertions.Core;");
 
-            // File header
-            sourceBuilder.AppendLine("// <auto-generated/>");
-            sourceBuilder.AppendLine("#pragma warning disable");
-            sourceBuilder.AppendLine("#nullable enable");
-            sourceBuilder.AppendLine();
-            sourceBuilder.AppendLine("using System;");
-            sourceBuilder.AppendLine("using System.Runtime.CompilerServices;");
-            sourceBuilder.AppendLine("using System.Threading.Tasks;");
-            sourceBuilder.AppendLine("using TUnit.Assertions.Core;");
+        // Add using for the original namespace to access helper methods
+        if (!string.IsNullOrEmpty(originalNamespace) && originalNamespace != namespaceName)
+        {
+            sourceBuilder.AppendLine($"using {originalNamespace};");
+        }
 
-            // Add using for the original namespace to access helper methods
-            if (!string.IsNullOrEmpty(originalNamespace) && originalNamespace != namespaceName)
-            {
-                sourceBuilder.AppendLine($"using {originalNamespace};");
-            }
+        sourceBuilder.AppendLine();
 
+        if (!string.IsNullOrEmpty(namespaceName))
+        {
+            sourceBuilder.AppendLine($"namespace {namespaceName};");
             sourceBuilder.AppendLine();
+        }
 
-            if (!string.IsNullOrEmpty(namespaceName))
-            {
-                sourceBuilder.AppendLine($"namespace {namespaceName};");
-                sourceBuilder.AppendLine();
-            }
-
-            // Generate assertion classes
-            foreach (var methodData in methodGroup)
-            {
-                GenerateAssertionClass(sourceBuilder, methodData);
-                sourceBuilder.AppendLine();
-            }
+        // Generate assertion classes
+        foreach (var methodData in methodGroup)
+        {
+            GenerateAssertionClass(sourceBuilder, methodData);
+            sourceBuilder.AppendLine();
+        }
 
-            // Generate extension methods class
-            // For file-scoped types, we can't use partial classes, so create a standalone public class
-            // For non-file-scoped types, we use partial classes to combine with the source definition
-            var isFileScopedType = methodGroup.Any(m => m.IsFileScoped);
+        // Generate extension methods class
+        // For file-scoped types, we can't use partial classes, so create a standalone public class
+        // For non-file-scoped types, we use partial classes to combine with the source definition
+        var isFileScopedType = methodGroup.Any(m => m.IsFileScoped);
 
-            var extensionClassName = isFileScopedType
-                ? $"{containingType.Name}Extensions"  // Standalone class for file-scoped types
-                : containingType.Name;                 // Partial class for public types
+        var extensionClassName = isFileScopedType
+            ? $"{containingType.Name}Extensions"  // Standalone class for file-scoped types
+            : containingType.Name;                 // Partial class for public types
 
-            var partialModifier = isFileScopedType ? "" : "partial ";
-            sourceBuilder.AppendLine($"public static {partialModifier}class {extensionClassName}");
-            sourceBuilder.AppendLine("{");
+        var partialModifier = isFileScopedType ? "" : "partial ";
+        sourceBuilder.AppendLine($"public static {partialModifier}class {extensionClassName}");
+        sourceBuilder.AppendLine("{");
 
-            foreach (var methodData in methodGroup)
-            {
-                GenerateExtensionMethod(sourceBuilder, methodData);
-                sourceBuilder.AppendLine();
-            }
+        foreach (var methodData in methodGroup)
+        {
+            GenerateExtensionMethod(sourceBuilder, methodData);
+            sourceBuilder.AppendLine();
+        }
 
-            sourceBuilder.AppendLine("}");
+        sourceBuilder.AppendLine("}");
 
-            // Add source to compilation
-            var fileName = $"{containingType.Name}.GeneratedAssertions.g.cs";
-            context.AddSource(fileName, sourceBuilder.ToString());
-        }
+        // Add source to compilation
+        var fileName = $"{containingType.Name}.GeneratedAssertions.g.cs";
+        context.AddSource(fileName, sourceBuilder.ToString());
     }
 
     private static void GenerateAssertionClass(StringBuilder sb, AssertionMethodData data)
@@ -929,7 +950,7 @@
         // Add null-forgiving operator for reference types if not already present
         // This is safe because we've already checked for null above
         var isNullable = data.TargetType.IsNullable;
         if (isNullable && !string.IsNullOrEmpty(inlinedBody) && !inlinedBody.StartsWith("value!"))
         {
             // Replace null-conditional operators with null-forgiving + regular operators
             // value?.Member becomes value!.Member (safe because we already null-checked)
@@ -1351,6 +1372,68 @@
         ImmutableEquatableArray<string> DiagnosticAttributesForExtensionMethod
     );
 
+    /// <summary>
+    /// All [GenerateAssertion] methods of one containing type; emitted as a single file.
+    /// </summary>
+    private sealed record ContainingTypeGroup(
+        ContainingTypeData ContainingType,
+        ImmutableEquatableArray<AssertionMethodData> Methods
+    );
+
+    /// <summary>
+    /// Equatable stand-in for a <see cref="Diagnostic"/>. A <see cref="Diagnostic"/> holds a
+    /// <see cref="Location"/> that references its <see cref="SyntaxTree"/>, which would keep old
+    /// compilations alive in the incremental cache; this captures only primitive location data.
+    /// </summary>
+    private sealed record DiagnosticInfo(
+        DiagnosticDescriptor Descriptor,
+        string? FilePath,
+        TextSpan TextSpan,
+        LinePositionSpan LineSpan,
+        ImmutableEquatableArray<string> MessageArgs)
+    {
+        public static DiagnosticInfo Create(DiagnosticDescriptor descriptor, Location location, params string[] messageArgs)
+        {
+            var lineSpan = location.GetLineSpan();
+            return new DiagnosticInfo(
+                descriptor,
+                lineSpan.Path,
+                location.SourceSpan,
+                lineSpan.Span,
+                messageArgs.ToImmutableEquatableArray());
+        }
+
+        public Diagnostic ToDiagnostic(Compilation compilation)
+        {
+            if (FilePath is null)
+            {
+                return Diagnostic.Create(Descriptor, Location.None, MessageArgs.Cast<object?>().ToArray());
+            }
+
+            // Reattach to the tree only when the path identifies it; trees parsed without a path
+            // all share an empty one.
+            SyntaxTree? match = null;
+            foreach (var tree in compilation.SyntaxTrees)
+            {
+                if (string.Equals(tree.FilePath, FilePath, StringComparison.Ordinal))
+                {
+                    if (match is not null)
+                    {
+                        match = null;
+                        break;
+                    }
+                    match = tree;
+                }
+            }
+
+            var location = match is not null && TextSpan.End <= match.Length
+                ? Location.Create(match, TextSpan)
+                : Location.Create(FilePath, TextSpan, LineSpan);
+
+            return Diagnostic.Create(Descriptor, location, MessageArgs.Cast<object?>().ToArray());
+        }
+    }
+
      private record ContainingTypeData(
         string Name,
         string FullContainingType,