From 3bc89cd4eb0e93eabaec6ec8189d703895b0cd7b Mon Sep 17 00:00:00 2001
From: Safia Abdalla <safia@safia.rocks>
Date: Fri, 4 Oct 2024 21:14:02 -0700
Subject: [PATCH] Use GetEntrypoint API and transformations for better caching

---
 .../PublicTopLevelProgramGenerator.cs         | 49 +++++++------------
 .../PublicTopLevelProgramGeneratorTests.cs    |  2 +-
 2 files changed, 19 insertions(+), 32 deletions(-)

diff --git a/src/Framework/AspNetCoreAnalyzers/src/SourceGenerators/PublicTopLevelProgramGenerator.cs b/src/Framework/AspNetCoreAnalyzers/src/SourceGenerators/PublicTopLevelProgramGenerator.cs
index 12d98325e72c..4ea333b96252 100644
--- a/src/Framework/AspNetCoreAnalyzers/src/SourceGenerators/PublicTopLevelProgramGenerator.cs
+++ b/src/Framework/AspNetCoreAnalyzers/src/SourceGenerators/PublicTopLevelProgramGenerator.cs
@@ -10,48 +10,35 @@ namespace Microsoft.AspNetCore.SourceGenerators;
 [Generator]
 public class PublicProgramSourceGenerator : IIncrementalGenerator
 {
+    private const string PublicPartialProgramClassSource = """
+// <auto-generated />
+public partial class Program { }
+""";
+
     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
-        var internalGeneratedProgramClass = context.CompilationProvider.Select((compilation, cancellationToken) =>
-        {
-            var program = compilation.GetTypeByMetadataName("Program");
+        var internalGeneratedProgramClass = context.CompilationProvider
+            // Get the entry point associated with the compilation, this maps to the Main method definition
+            .Select((compilation, cancellationToken) => compilation.GetEntryPoint(cancellationToken))
+            // Get the containing symbol of the entry point, this maps to the Program class
+            .Select((symbol, _) => symbol?.ContainingSymbol)
             // If the program class is already public, we don't need to generate anything.
-            if (program is null || program.DeclaredAccessibility == Accessibility.Public)
-            {
-                return null;
-            }
-            // If the discovered `Program` type is an interface, struct or generic type then its not
+            .Select((symbol, _) => symbol?.DeclaredAccessibility == Accessibility.Public ? null : symbol)
+            // If the discovered `Program` type is not a class then its not
             // generated and has been defined in source, so we can skip it
-            if (program.TypeKind == TypeKind.Struct || program.TypeKind == TypeKind.Interface || program.IsGenericType)
-            {
-                return null;
-            }
+            .Select((symbol, _) => symbol is INamedTypeSymbol { TypeKind: TypeKind.Class } ? symbol : null)
             // If there are multiple partial declarations, then do nothing since we don't want
             // to trample on visibility explicitly set by the user
-            if (program.DeclaringSyntaxReferences.Length > 1)
-            {
-                return null;
-            }
+            .Select((symbol, _) => symbol is { DeclaringSyntaxReferences: { Length: 1 } declaringSyntaxReferences } ? declaringSyntaxReferences.Single() : null)
             // If the `Program` class is already declared in user code, we don't need to generate anything.
-            if (program.DeclaringSyntaxReferences.SingleOrDefault()?.GetSyntax(cancellationToken) is ClassDeclarationSyntax)
-            {
-                return null;
-            }
-            return program;
-        });
+            .Select((declaringSyntaxReference, cancellationToken) => declaringSyntaxReference?.GetSyntax(cancellationToken) is ClassDeclarationSyntax ? null : declaringSyntaxReference);
 
-        context.RegisterSourceOutput(internalGeneratedProgramClass, (context, symbol) =>
+        context.RegisterSourceOutput(internalGeneratedProgramClass, (context, result) =>
         {
-            if (symbol is null)
+            if (result is not null)
             {
-                return;
+                context.AddSource("PublicTopLevelProgram.Generated.g.cs", PublicPartialProgramClassSource);
             }
-
-            var output = """
-// <auto-generated />
-public partial class Program { }
-""";
-            context.AddSource("PublicTopLevelProgram.Generated.cs", output);
         });
     }
 }
diff --git a/src/Framework/AspNetCoreAnalyzers/test/SourceGenerators/PublicTopLevelProgramGeneratorTests.cs b/src/Framework/AspNetCoreAnalyzers/test/SourceGenerators/PublicTopLevelProgramGeneratorTests.cs
index 3a7845b159b9..c91a110031a7 100644
--- a/src/Framework/AspNetCoreAnalyzers/test/SourceGenerators/PublicTopLevelProgramGeneratorTests.cs
+++ b/src/Framework/AspNetCoreAnalyzers/test/SourceGenerators/PublicTopLevelProgramGeneratorTests.cs
@@ -25,7 +25,7 @@ public async Task GeneratesSource_ProgramWithTopLevelStatements()
 public partial class Program { }
 """;
 
-        await VerifyCS.VerifyAsync(source, "PublicTopLevelProgram.Generated.cs", expected);
+        await VerifyCS.VerifyAsync(source, "PublicTopLevelProgram.Generated.g.cs", expected);
     }
 
     [Fact]