From b682d7971edbe7174b3da87269fcbb6dad8e9935 Mon Sep 17 00:00:00 2001
From: Steven Giesel <stgiesel35@gmail.com>
Date: Sun, 26 Nov 2023 13:22:12 +0100
Subject: [PATCH] collect everything - then generate

---
 .../Web.Stubs/StubGenerator.cs                | 69 +++++++++++--------
 1 file changed, 42 insertions(+), 27 deletions(-)

diff --git a/src/bunit.generators/Web.Stubs/StubGenerator.cs b/src/bunit.generators/Web.Stubs/StubGenerator.cs
index 0e1959e5a..76d7cac24 100644
--- a/src/bunit.generators/Web.Stubs/StubGenerator.cs
+++ b/src/bunit.generators/Web.Stubs/StubGenerator.cs
@@ -1,3 +1,4 @@
+using System.Collections.Immutable;
 using System.Linq;
 using System.Text;
 using Microsoft.CodeAnalysis;
@@ -18,7 +19,8 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 			.CreateSyntaxProvider(
 				predicate: static (s, _) => s is InvocationExpressionSyntax,
 				transform: static (ctx, _) => GetStubClassInfo(ctx))
-			.Where(static m => m is not null);
+			.Where(static m => m is not null)
+			.Collect();
 
 		context.RegisterSourceOutput(
 			classesToStub,
@@ -78,17 +80,21 @@ static string GetInterceptorFilePath(SyntaxTree tree, Compilation compilation)
 		}
 	}
 
-	private static void Execute(StubClassInfo classInfo, SourceProductionContext context)
+	private static void Execute(ImmutableArray<StubClassInfo> classInfos, SourceProductionContext context)
 	{
-		var didStubComponent = GenerateStubComponent(classInfo, context);
-		if (didStubComponent is false)
+		foreach (var stubClassGrouped in classInfos.GroupBy(c => c.UniqueQualifier))
 		{
-			return;
-		}
+			var stubbedComponentGroup = stubClassGrouped.First();
+			var didStubComponent = GenerateStubComponent(stubbedComponentGroup, context);
+			if (didStubComponent is false)
+			{
+				return;
+			}
 
-		// Generate the attribute
-		const string attribute = @"namespace System.Runtime.CompilerServices
+			// Generate the attribute
+			const string attribute = @"namespace System.Runtime.CompilerServices
 {
+	[AttributeUsage(AttributeTargets.Method, AllowMultiple = true)]
 	sealed file class InterceptsLocationAttribute : Attribute
 	{
 		public InterceptsLocationAttribute(string filePath, int line, int column)
@@ -100,25 +106,33 @@ public InterceptsLocationAttribute(string filePath, int line, int column)
 	}
 }";
 
-		// Generate the interceptor
-		var interceptorSource = new StringBuilder();
-		interceptorSource.AppendLine(attribute);
-		interceptorSource.AppendLine();
-		interceptorSource.AppendLine("namespace Bunit");
-		interceptorSource.AppendLine("{");
-		interceptorSource.AppendLine($"\tstatic class Interceptor{classInfo.StubClassName}");
-		interceptorSource.AppendLine("\t{");
-		interceptorSource.AppendLine(
-			$"\t\t[System.Runtime.CompilerServices.InterceptsLocationAttribute(\"{classInfo.Path}\", {classInfo.Line}, {classInfo.Column})]");
-		interceptorSource.AppendLine("\t\tpublic static global::Bunit.ComponentFactoryCollection AddGeneratedStubInterceptor<TComponent>(this global::Bunit.ComponentFactoryCollection factories)");
-		interceptorSource.AppendLine("\t\t\twhere TComponent : Microsoft.AspNetCore.Components.IComponent");
-		interceptorSource.AppendLine("\t\t{");
-		interceptorSource.AppendLine($"\t\t\treturn factories.Add<global::{classInfo.TargetType.ToDisplayString()}, {classInfo.TargetTypeNamespace}.{classInfo.StubClassName}>();");
-		interceptorSource.AppendLine("\t\t}");
-		interceptorSource.AppendLine("\t}");
-		interceptorSource.AppendLine("}");
-
-		context.AddSource($"Interceptor{classInfo.StubClassName}.g.cs", interceptorSource.ToString());
+			// Generate the interceptor
+			var interceptorSource = new StringBuilder();
+			interceptorSource.AppendLine(attribute);
+			interceptorSource.AppendLine();
+			interceptorSource.AppendLine("namespace Bunit");
+			interceptorSource.AppendLine("{");
+			interceptorSource.AppendLine($"\tstatic class Interceptor{stubbedComponentGroup.StubClassName}");
+			interceptorSource.AppendLine("\t{");
+
+			foreach (var hit in stubClassGrouped)
+			{
+				interceptorSource.AppendLine(
+					$"\t\t[System.Runtime.CompilerServices.InterceptsLocationAttribute(\"{hit.Path}\", {hit.Line}, {hit.Column})]");
+			}
+
+			interceptorSource.AppendLine(
+				"\t\tpublic static global::Bunit.ComponentFactoryCollection AddGeneratedStubInterceptor<TComponent>(this global::Bunit.ComponentFactoryCollection factories)");
+			interceptorSource.AppendLine("\t\t\twhere TComponent : Microsoft.AspNetCore.Components.IComponent");
+			interceptorSource.AppendLine("\t\t{");
+			interceptorSource.AppendLine(
+				$"\t\t\treturn factories.Add<global::{stubbedComponentGroup.TargetType.ToDisplayString()}, {stubbedComponentGroup.TargetTypeNamespace}.{stubbedComponentGroup.StubClassName}>();");
+			interceptorSource.AppendLine("\t\t}");
+			interceptorSource.AppendLine("\t}");
+			interceptorSource.AppendLine("}");
+
+			context.AddSource($"Interceptor{stubbedComponentGroup.StubClassName}.g.cs", interceptorSource.ToString());
+		}
 	}
 
 	private static bool GenerateStubComponent(StubClassInfo classInfo, SourceProductionContext context)
@@ -172,6 +186,7 @@ internal sealed class StubClassInfo
 {
 	public string StubClassName { get; set; }
 	public string TargetTypeNamespace { get; set; }
+	public string UniqueQualifier => $"{TargetTypeNamespace}.{StubClassName}";
 	public ITypeSymbol TargetType { get; set; }
 	public string Path { get; set; }
 	public int Line { get; set; }