@@ -10,48 +10,35 @@ namespace Microsoft.AspNetCore.SourceGenerators;
1010[ Generator ]
1111public class PublicProgramSourceGenerator : IIncrementalGenerator
1212{
13+ private const string PublicPartialProgramClassSource = """
14+ // <auto-generated />
15+ public partial class Program { }
16+ """ ;
17+
1318public void Initialize ( IncrementalGeneratorInitializationContext context )
1419{
15- var internalGeneratedProgramClass = context . CompilationProvider . Select ( ( compilation , cancellationToken ) =>
16- {
17- var program = compilation . GetTypeByMetadataName ( "Program" ) ;
20+ var internalGeneratedProgramClass = context . CompilationProvider
21+ // Get the entry point associated with the compilation, this maps to the Main method definition
22+ . Select ( ( compilation , cancellationToken ) => compilation . GetEntryPoint ( cancellationToken ) )
23+ // Get the containing symbol of the entry point, this maps to the Program class
24+ . Select ( ( symbol , _ ) => symbol ? . ContainingSymbol )
1825// If the program class is already public, we don't need to generate anything.
19- if ( program is null || program . DeclaredAccessibility == Accessibility . Public )
20- {
21- return null ;
22- }
23- // If the discovered `Program` type is an interface, struct or generic type then its not
26+ . Select ( ( symbol , _ ) => symbol ? . DeclaredAccessibility == Accessibility . Public ? null : symbol )
27+ // If the discovered `Program` type is not a class then its not
2428// generated and has been defined in source, so we can skip it
25- if ( program . TypeKind == TypeKind . Struct || program . TypeKind == TypeKind . Interface || program . IsGenericType )
26- {
27- return null ;
28- }
29+ . Select ( ( symbol , _ ) => symbol is INamedTypeSymbol { TypeKind : TypeKind . Class } ? symbol : null )
2930// If there are multiple partial declarations, then do nothing since we don't want
3031// to trample on visibility explicitly set by the user
31- if ( program . DeclaringSyntaxReferences . Length > 1 )
32- {
33- return null ;
34- }
32+ . Select ( ( symbol , _ ) => symbol is { DeclaringSyntaxReferences : { Length : 1 } declaringSyntaxReferences } ? declaringSyntaxReferences . Single ( ) : null )
3533// If the `Program` class is already declared in user code, we don't need to generate anything.
36- if ( program . DeclaringSyntaxReferences . SingleOrDefault ( ) ? . GetSyntax ( cancellationToken ) is ClassDeclarationSyntax )
37- {
38- return null ;
39- }
40- return program ;
41- } ) ;
34+ . Select ( ( declaringSyntaxReference , cancellationToken ) => declaringSyntaxReference ? . GetSyntax ( cancellationToken ) is ClassDeclarationSyntax ? null : declaringSyntaxReference ) ;
4235
43- context . RegisterSourceOutput ( internalGeneratedProgramClass , ( context , symbol ) =>
36+ context . RegisterSourceOutput ( internalGeneratedProgramClass , ( context , result ) =>
4437{
45- if ( symbol is null )
38+ if ( result is not null )
4639{
47- return ;
40+ context . AddSource ( "PublicTopLevelProgram.Generated.g.cs" , PublicPartialProgramClassSource ) ;
4841}
49-
50- var output = """
51- // <auto-generated />
52- public partial class Program { }
53- """ ;
54- context . AddSource ( "PublicTopLevelProgram.Generated.cs" , output ) ;
5542} ) ;
5643}
5744}