Incremental generator for illink analyzer tests

src/tools/illink/test/ILLink.RoslynAnalyzer.Tests.Generator/TestCaseGenerator.cs

@@ -1,10 +1,14 @@
 // Copyright (c) .NET Foundation and contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
+using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Linq;
 using System.Reflection.Metadata;
 using System.Text;
+using System.Threading;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 
@@ -95,7 +99,7 @@ public static string GenerateClassFooter ()
 	}
 
 	[Generator]
-	public class TestCaseGenerator : ISourceGenerator
+	public class TestCaseGenerator : IIncrementalGenerator
 	{
 		public const string TestCaseAssembly = "Mono.Linker.Tests.Cases";
 
@@ -119,7 +123,7 @@ public void Execute (GeneratorExecutionContext context)
 			TestCases testCases = new ();
 			string suiteName;
 
-			// Find test classes
+			// Find test classes to generate
 			foreach (var typeDefHandle in metadataReader.TypeDefinitions) {
 				TypeDefinition typeDef = metadataReader.GetTypeDefinition (typeDefHandle);
 				// Must not be a nested type
@@ -148,65 +152,205 @@ public void Execute (GeneratorExecutionContext context)
 
 				testCases.Add (suiteName, testName);
 			}
+		}
+
+		static string GetFullName(ClassDeclarationSyntax classSyntax)
+		{
+			return GetFullName(classSyntax)!;
 
-			TestCases existingTestCases = ((ExistingTestCaseDiscoverer) context.SyntaxContextReceiver!).ExistingTestCases;
+			static string? GetFullName(SyntaxNode? node) {
+				if (node == null)
+					return null;
 
-			// Generate test facts
-			foreach (var kvp in testCases.Suites) {
-				suiteName = kvp.Key;
-				var cases = kvp.Value;
+				var name = node switch {
+					ClassDeclarationSyntax classSyntax => $"{classSyntax.Identifier.ValueText}",
+					NamespaceDeclarationSyntax namespaceSyntax => $"{namespaceSyntax.Name}",
+					CompilationUnitSyntax => null,
+					_ => throw new NotImplementedException($"GetFullName for node type {node.GetType()}")
+				};
 
-				bool newTestSuite = !existingTestCases.Suites.TryGetValue (suiteName, out HashSet<string> existingCases);
-				var newCases = newTestSuite ? cases : cases.Except (existingCases);
-				// Skip generating a test class if all testcases in the suite already exist.
-				if (!newCases.Any ())
-					continue;
+				if (name == null)
+					return null;
+
+				return GetFullName(node.Parent) is string parentName
+					? $"{parentName}.{name}"
+					: name;
+			}
+		}
+
+		public void Initialize (IncrementalGeneratorInitializationContext context)
+		{
+			IncrementalValuesProvider<MetadataReference> metadataReferences = context.MetadataReferencesProvider;
+
+			IncrementalValueProvider<Compilation> compilation = context.CompilationProvider;
+
+			// For any steps below which can fail (for example, getting metadata might return null),
+			// we use SelectMany to return zero or one results.
+			IncrementalValuesProvider<IAssemblySymbol> assemblySymbol = metadataReferences
+				.Combine(compilation)
+				.SelectMany(static (combined, _) => {
+					var (reference, compilation) = combined;
+					return compilation.GetAssemblyOrModuleSymbol (reference) is IAssemblySymbol asmSym && asmSym.Name == TestCaseAssembly
+						? ImmutableArray.Create(asmSym)
+						: ImmutableArray<IAssemblySymbol>.Empty;
+				});
+
+			IncrementalValuesProvider<AssemblyMetadata> assemblyMetadata = assemblySymbol
+				.SelectMany(static (symbol, _) =>
+					symbol.GetMetadata() is AssemblyMetadata metadata
+						? ImmutableArray.Create(metadata)
+						: ImmutableArray<AssemblyMetadata>.Empty);
+
+			IncrementalValuesProvider<ModuleMetadata> moduleMetadata = assemblyMetadata
+				.Select(static (metadata, _) =>
+					metadata.GetModules().Single());
 
+			IncrementalValuesProvider<MetadataReader> metadataReader = moduleMetadata
+				.Select(static (metadata, _) =>
+					metadata.GetMetadataReader());
+
+			// Find all test cases (some of which may need generated facts)
+			IncrementalValuesProvider<TestCases> testCases = metadataReader
+				.Select(static (reader, cancellationToken) =>
+					FindTestCases(reader, cancellationToken));
+
+			// Find already-generated test types
+			IncrementalValuesProvider<INamedTypeSymbol?> existingTestTypes = context.SyntaxProvider.CreateSyntaxProvider(
+				static (node, cancellationToken) => {
+					if (node is not ClassDeclarationSyntax classSyntax)
+						return false;
+
+					var typeFullName = GetFullName(classSyntax);
+
+					// Ignore types not in the testcase namespace or that don't end with "Tests"
+					if (!typeFullName.StartsWith (TestClassGenerator.TestNamespace))
+						return false;
+					if (!typeFullName.EndsWith ("Tests"))
+						return false;
+
+					return true;
+				},
+				static (generatorSyntaxContext, cancellationToken) => {
+					var node = generatorSyntaxContext.Node;
+					return generatorSyntaxContext.SemanticModel.GetDeclaredSymbol (node, cancellationToken) is INamedTypeSymbol typeSymbol
+						? typeSymbol
+						: null;
+				});
+
+			// Find already-generated test cases
+			IncrementalValueProvider<TestCases> existingTestCases = existingTestTypes
+				// Find test methods
+				.SelectMany(static (typeSymbol, cancellationToken) => 
+					FindExistingTestCases(typeSymbol, cancellationToken))
+				.Collect()
+				.Select(static (existingTestCases, cancellationToken) => {
+					var testCases = new TestCases();
+					foreach (var (suiteName, testName) in existingTestCases) {
+						testCases.Add(suiteName, testName);
+					}
+					return testCases;
+				});
+
+			// Find the new test cases that need generated facts
+			IncrementalValuesProvider<(string suiteName, IEnumerable<string> newCases, bool newTestSuite)> newTestCases = testCases
+				.Combine(existingTestCases)
+				.SelectMany(static (combined, cancellationToken) => {
+					var (testCases, existingTestCases) = combined;
+					return FindNewTestCases(testCases, existingTestCases, cancellationToken);
+				});
+
+
+			// Generate facts for all testcases that don't already exist
+			context.RegisterSourceOutput(newTestCases, static (sourceProductionContext, newTestCases) => {
 				StringBuilder sourceBuilder = new ();
+				var suiteName = newTestCases.suiteName;
+				bool newTestSuite = newTestCases.newTestSuite;
+
 				sourceBuilder.Append (TestClassGenerator.GenerateClassHeader (suiteName, newTestSuite));
-				// Generate facts for all testcases that don't already exist
-				foreach (var testCase in newCases)
+				foreach (var testCase in newTestCases.newCases)
 					sourceBuilder.Append (TestClassGenerator.GenerateFact (testCase));
 				sourceBuilder.Append (TestClassGenerator.GenerateClassFooter ());
 
-				context.AddSource ($"{suiteName}Tests.g.cs", sourceBuilder.ToString ());
+				sourceProductionContext.AddSource ($"{suiteName}Tests.g.cs", sourceBuilder.ToString ());
+			});
+
+			static IEnumerable<(string SuiteName, IEnumerable<string> NewCases, bool newTestSuite)> FindNewTestCases (TestCases testCases, TestCases existingTestCases, CancellationToken cancellationToken) {
+				foreach (var kvp in testCases.Suites) {
+					cancellationToken.ThrowIfCancellationRequested ();
+
+					string suiteName = kvp.Key;
+					var cases = kvp.Value;
+
+					bool newTestSuite = !existingTestCases.Suites.TryGetValue (suiteName, out HashSet<string> existingCases);
+					var newCases = newTestSuite ? cases : cases.Except (existingCases);
+					// Skip generating a test class if all testcases in the suite already exist.
+					if (!newCases.Any ())
+						continue;
+
+					yield return (suiteName, newCases, newTestSuite);
+				}
 			}
-		}
 
-		public void Initialize (GeneratorInitializationContext context)
-		{
-			context.RegisterForSyntaxNotifications (() => new ExistingTestCaseDiscoverer ());
-		}
-	}
+			static IEnumerable<(string SuiteName, string TestName)> FindExistingTestCases (INamedTypeSymbol? typeSymbol, CancellationToken cancellationToken) {
+				if (typeSymbol is null)
+					yield break;
 
-	sealed class ExistingTestCaseDiscoverer : ISyntaxContextReceiver
-	{
-		public readonly TestCases ExistingTestCases = new TestCases ();
+				var displayFormat = new SymbolDisplayFormat (
+					typeQualificationStyle: SymbolDisplayTypeQualificationStyle.NameAndContainingTypesAndNamespaces);
+				string typeFullName = typeSymbol.ToDisplayString (displayFormat);
 
-		public void OnVisitSyntaxNode (GeneratorSyntaxContext context)
-		{
-			if (context.Node is not ClassDeclarationSyntax classSyntax)
-				return;
+				// Ignore types not in the testcase namespace or that don't end with "Tests"
 
-			if (context.SemanticModel.GetDeclaredSymbol (classSyntax) is not INamedTypeSymbol typeSymbol)
-				return;
+				string suiteName = typeFullName.Substring (TestClassGenerator.TestNamespace.Length + 1);
+				suiteName = suiteName.Substring (0, suiteName.Length - "Tests".Length);
+				foreach (var member in typeSymbol.GetMembers ()) {
+					cancellationToken.ThrowIfCancellationRequested ();
 
-			var displayFormat = new SymbolDisplayFormat (
-				typeQualificationStyle: SymbolDisplayTypeQualificationStyle.NameAndContainingTypesAndNamespaces);
-			string typeFullName = typeSymbol.ToDisplayString (displayFormat);
+					if (member is not IMethodSymbol methodSymbol)
+						continue;
+					yield return (suiteName, methodSymbol.Name);
+				}
+			}
 
-			// Ignore types not in the testcase namespace or that don't end with "Tests"
-			if (!typeFullName.StartsWith (TestClassGenerator.TestNamespace))
-				return;
-			if (!typeFullName.EndsWith ("Tests"))
-				return;
+			static TestCases FindTestCases (MetadataReader metadataReader, CancellationToken cancellationToken) {
+				TestCases testCases = new ();
+				string suiteName;
 
-			string suiteName = typeFullName.Substring (TestClassGenerator.TestNamespace.Length + 1);
-			suiteName = suiteName.Substring (0, suiteName.Length - 5);
-			foreach (var memberSymbol in typeSymbol.GetMembers ()) {
-				if (memberSymbol is not IMethodSymbol methodSymbol)
-					continue;
-				ExistingTestCases.Add (suiteName, methodSymbol.Name);
+				// Find test classes to generate
+				foreach (var typeDefHandle in metadataReader.TypeDefinitions) {
+					cancellationToken.ThrowIfCancellationRequested ();
+
+					TypeDefinition typeDef = metadataReader.GetTypeDefinition (typeDefHandle);
+					// Must not be a nested type
+					if (typeDef.IsNested)
+						continue;
+
+					string ns = metadataReader.GetString (typeDef.Namespace);
+					// Must be in the testcases namespace
+					if (!ns.StartsWith (TestCaseAssembly))
+						continue;
+
+					// Must have a Main method
+					bool hasMain = false;
+					foreach (var methodDefHandle in typeDef.GetMethods ()) {
+						cancellationToken.ThrowIfCancellationRequested ();
+
+						MethodDefinition methodDef = metadataReader.GetMethodDefinition (methodDefHandle);
+						if (metadataReader.GetString (methodDef.Name) == "Main") {
+							hasMain = true;
+							break;
+						}
+					}
+					if (!hasMain)
+						continue;
+
+					string testName = metadataReader.GetString (typeDef.Name);
+					suiteName = ns.Substring (TestCaseAssembly.Length + 1);
+
+					testCases.Add (suiteName, testName);
+				}
+
+				return testCases;
 			}
 		}
 	}