From 0346f9d277e32eb79f4eba0ade7e1d4d9a3b8bc7 Mon Sep 17 00:00:00 2001
From: Stuart Turner <stuart@turner-isageek.com>
Date: Mon, 14 Jul 2025 09:41:40 -0500
Subject: [PATCH] Add generator cache tracking

---
 .../ImmediateHandlersGenerator.cs             | 21 +++++---
 .../GeneratorTests/GeneratorTestHelper.cs     | 52 +++++++++++++++++--
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/src/Immediate.Handlers.Generators/ImmediateHandlersGenerator.cs b/src/Immediate.Handlers.Generators/ImmediateHandlersGenerator.cs
index ee796fe5..14dcd778 100644
--- a/src/Immediate.Handlers.Generators/ImmediateHandlersGenerator.cs
+++ b/src/Immediate.Handlers.Generators/ImmediateHandlersGenerator.cs
@@ -15,14 +15,16 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 			.MetadataReferencesProvider
 			.Where(r => (r.Display ?? "").Contains("Microsoft.Extensions.DependencyInjection.Abstractions"))
 			.Collect()
-			.Select((refs, _) => refs.Any());
+			.Select((refs, _) => refs.Any())
+			.WithTrackingName("MsDi");
 
 		var assemblyName = context.CompilationProvider
 			.Select((cp, _) => cp.AssemblyName!
 				.Replace(".", string.Empty)
 				.Replace(" ", string.Empty)
 				.Trim()
-			);
+			)
+			.WithTrackingName("AssemblyName");
 
 		var @namespace = context
 			.AnalyzerConfigOptionsProvider
@@ -30,7 +32,8 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 				(c, _) => c.GlobalOptions
 					.TryGetValue("build_property.rootnamespace", out var ns)
 						? ns : null
-			);
+			)
+			.WithTrackingName("RootNamespace");
 
 		var behaviors = context.SyntaxProvider
 			.ForAttributeWithMetadataName(
@@ -39,18 +42,21 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 				TransformBehaviors.ParseBehaviors
 			)
 			.SelectMany((x, _) => x)
-			.Collect();
+			.Collect()
+			.WithTrackingName("Behaviors");
 
 		var handlers = context.SyntaxProvider
 			.ForAttributeWithMetadataName(
 				"Immediate.Handlers.Shared.HandlerAttribute",
 				predicate: (node, _) => node is TypeDeclarationSyntax,
 				TransformHandler.ParseHandler
-			);
+			)
+			.WithTrackingName("Handlers");
 
 		var handlerNodes = handlers
 			.Combine(behaviors)
-			.Combine(hasMsDi);
+			.Combine(hasMsDi)
+			.WithTrackingName("HandlersWithBehaviors");
 
 		var template = GetTemplate("Handler");
 		context.RegisterSourceOutput(
@@ -71,7 +77,8 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 			.Combine(@namespace
 				.Combine(hasMsDi)
 				.Combine(assemblyName)
-			);
+			)
+			.WithTrackingName("Registrations");
 
 		context.RegisterSourceOutput(
 			registrationNodes,
diff --git a/tests/Immediate.Handlers.Tests/GeneratorTests/GeneratorTestHelper.cs b/tests/Immediate.Handlers.Tests/GeneratorTests/GeneratorTestHelper.cs
index cf1952bb..ccfabcd8 100644
--- a/tests/Immediate.Handlers.Tests/GeneratorTests/GeneratorTestHelper.cs
+++ b/tests/Immediate.Handlers.Tests/GeneratorTests/GeneratorTestHelper.cs
@@ -1,3 +1,4 @@
+using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 using Immediate.Handlers.Generators;
 using Immediate.Handlers.Tests.Helpers;
@@ -28,10 +29,34 @@ .. assemblies.GetAdditionalReferences(),
 			)
 		);
 
-		var generator = new ImmediateHandlersGenerator();
+		var clone = compilation.Clone().AddSyntaxTrees(CSharpSyntaxTree.ParseText("// dummy"));
 
-		var driver = CSharpGeneratorDriver
-			.Create(generator)
+		GeneratorDriver driver = CSharpGeneratorDriver.Create(
+			generators: [new ImmediateHandlersGenerator().AsSourceGenerator()],
+			driverOptions: new GeneratorDriverOptions(default, trackIncrementalGeneratorSteps: true)
+		);
+
+		var result1 = RunGenerator(ref driver, compilation);
+		var result2 = RunGenerator(ref driver, clone);
+
+		foreach (var (_, step) in result2.Results[0].TrackedOutputSteps)
+			AssertSteps(step);
+
+		foreach (var step in TrackedSteps)
+		{
+			if (result2.Results[0].TrackedSteps.TryGetValue(step, out var outputs))
+				AssertSteps(outputs);
+		}
+
+		return result1;
+	}
+
+	private static GeneratorDriverRunResult RunGenerator(
+		ref GeneratorDriver driver,
+		Compilation compilation
+	)
+	{
+		driver = driver
 			.RunGeneratorsAndUpdateCompilation(
 				compilation,
 				out var outputCompilation,
@@ -47,4 +72,25 @@ out var diagnostics
 		Assert.Empty(diagnostics);
 		return driver.GetRunResult();
 	}
+
+	private static ReadOnlySpan<string> TrackedSteps =>
+		new string[]
+		{
+			"MsDi",
+			"AssemblyName",
+			"RootNamespace",
+			"Behaviors",
+			"Handlers",
+			"HandlersWithBehaviors",
+			"Registrations",
+		};
+
+	private static void AssertSteps(
+		ImmutableArray<IncrementalGeneratorRunStep> steps
+	)
+	{
+		var outputs = steps.SelectMany(o => o.Outputs);
+
+		Assert.All(outputs, o => Assert.True(o.Reason is IncrementalStepRunReason.Unchanged or IncrementalStepRunReason.Cached));
+	}
 }