|
| 1 | +using System.Collections.Generic; |
| 2 | +using System.IO; |
| 3 | +using System.Text; |
| 4 | +using Microsoft.CodeAnalysis; |
| 5 | +using Microsoft.CodeAnalysis.CSharp; |
| 6 | +using Microsoft.CodeAnalysis.CSharp.Syntax; |
| 7 | + |
| 8 | +using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory; |
| 9 | + |
| 10 | +namespace System.Linq.Async.SourceGenerator |
| 11 | +{ |
| 12 | + [Generator] |
| 13 | + public sealed class AsyncOverloadsGenerator : ISourceGenerator |
| 14 | + { |
| 15 | + private const string AttributeSource = |
| 16 | + "using System;\n" + |
| 17 | + "using System.Diagnostics;\n" + |
| 18 | + "namespace System.Linq\n" + |
| 19 | + "{\n" + |
| 20 | + " [AttributeUsage(AttributeTargets.Method)]\n" + |
| 21 | + " [Conditional(\"COMPILE_TIME_ONLY\")]\n" + |
| 22 | + " internal sealed class GenerateAsyncOverloadAttribute : Attribute { }\n" + |
| 23 | + "}\n"; |
| 24 | + |
| 25 | + public void Initialize(GeneratorInitializationContext context) |
| 26 | + { |
| 27 | + context.RegisterForSyntaxNotifications(() => new SyntaxReceiver()); |
| 28 | + context.RegisterForPostInitialization(c => c.AddSource("GenerateAsyncOverloadAttribute", AttributeSource)); |
| 29 | + } |
| 30 | + |
| 31 | + public void Execute(GeneratorExecutionContext context) |
| 32 | + { |
| 33 | + if (context.SyntaxReceiver is not SyntaxReceiver syntaxReceiver) return; |
| 34 | + |
| 35 | + var options = GetGenerationOptions(context); |
| 36 | + var methodsBySyntaxTree = GetMethodsGroupedBySyntaxTree(context, syntaxReceiver); |
| 37 | + |
| 38 | + foreach (var grouping in methodsBySyntaxTree) |
| 39 | + context.AddSource( |
| 40 | + $"{Path.GetFileNameWithoutExtension(grouping.SyntaxTree.FilePath)}.AsyncOverloads", |
| 41 | + GenerateOverloads(grouping, options)); |
| 42 | + } |
| 43 | + |
| 44 | + private static GenerationOptions GetGenerationOptions(GeneratorExecutionContext context) |
| 45 | + => new(SupportFlatAsyncApi: context.ParseOptions.PreprocessorSymbolNames.Contains("SUPPORT_FLAT_ASYNC_API")); |
| 46 | + |
| 47 | + private static IEnumerable<AsyncMethodGrouping> GetMethodsGroupedBySyntaxTree(GeneratorExecutionContext context, SyntaxReceiver syntaxReceiver) |
| 48 | + => GetMethodsGroupedBySyntaxTree( |
| 49 | + context, |
| 50 | + syntaxReceiver, |
| 51 | + GetAsyncOverloadAttributeSymbol(context)); |
| 52 | + |
| 53 | + private static string GenerateOverloads(AsyncMethodGrouping grouping, GenerationOptions options) |
| 54 | + { |
| 55 | + var usings = grouping.SyntaxTree.GetRoot() is CompilationUnitSyntax compilationUnit |
| 56 | + ? compilationUnit.Usings.ToString() |
| 57 | + : string.Empty; |
| 58 | + |
| 59 | + var overloads = new StringBuilder(); |
| 60 | + overloads.AppendLine("#nullable enable"); |
| 61 | + overloads.AppendLine(usings); |
| 62 | + overloads.AppendLine("namespace System.Linq"); |
| 63 | + overloads.AppendLine("{"); |
| 64 | + overloads.AppendLine(" partial class AsyncEnumerable"); |
| 65 | + overloads.AppendLine(" {"); |
| 66 | + |
| 67 | + foreach (var method in grouping.Methods) |
| 68 | + overloads.AppendLine(GenerateOverload(method, options)); |
| 69 | + |
| 70 | + overloads.AppendLine(" }"); |
| 71 | + overloads.AppendLine("}"); |
| 72 | + |
| 73 | + return overloads.ToString(); |
| 74 | + } |
| 75 | + |
| 76 | + private static string GenerateOverload(AsyncMethod method, GenerationOptions options) |
| 77 | + => MethodDeclaration(method.Syntax.ReturnType, GetMethodName(method.Symbol, options)) |
| 78 | + .WithModifiers(TokenList(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.StaticKeyword))) |
| 79 | + .WithTypeParameterList(method.Syntax.TypeParameterList) |
| 80 | + .WithParameterList(method.Syntax.ParameterList) |
| 81 | + .WithConstraintClauses(method.Syntax.ConstraintClauses) |
| 82 | + .WithExpressionBody(ArrowExpressionClause( |
| 83 | + InvocationExpression( |
| 84 | + IdentifierName(method.Symbol.Name), |
| 85 | + ArgumentList( |
| 86 | + SeparatedList( |
| 87 | + method.Syntax.ParameterList.Parameters |
| 88 | + .Select(p => Argument(IdentifierName(p.Identifier)))))))) |
| 89 | + .WithSemicolonToken(Token(SyntaxKind.SemicolonToken)) |
| 90 | + .WithLeadingTrivia(method.Syntax.GetLeadingTrivia().Where(t => t.GetStructure() is not DirectiveTriviaSyntax)) |
| 91 | + .NormalizeWhitespace() |
| 92 | + .ToFullString(); |
| 93 | + |
| 94 | + private static INamedTypeSymbol GetAsyncOverloadAttributeSymbol(GeneratorExecutionContext context) |
| 95 | + => context.Compilation.GetTypeByMetadataName("System.Linq.GenerateAsyncOverloadAttribute") ?? throw new InvalidOperationException(); |
| 96 | + |
| 97 | + private static IEnumerable<AsyncMethodGrouping> GetMethodsGroupedBySyntaxTree(GeneratorExecutionContext context, SyntaxReceiver syntaxReceiver, INamedTypeSymbol attributeSymbol) |
| 98 | + => from candidate in syntaxReceiver.Candidates |
| 99 | + group candidate by candidate.SyntaxTree into grouping |
| 100 | + let model = context.Compilation.GetSemanticModel(grouping.Key) |
| 101 | + select new AsyncMethodGrouping( |
| 102 | + grouping.Key, |
| 103 | + from methodSyntax in grouping |
| 104 | + let methodSymbol = model.GetDeclaredSymbol(methodSyntax) ?? throw new InvalidOperationException() |
| 105 | + where methodSymbol.GetAttributes().Any(a => SymbolEqualityComparer.Default.Equals(a.AttributeClass!, attributeSymbol)) |
| 106 | + select new AsyncMethod(methodSymbol, methodSyntax)); |
| 107 | + |
| 108 | + private static string GetMethodName(IMethodSymbol methodSymbol, GenerationOptions options) |
| 109 | + { |
| 110 | + var methodName = methodSymbol.Name.Replace("Core", ""); |
| 111 | + return options.SupportFlatAsyncApi |
| 112 | + ? methodName.Replace("Await", "").Replace("WithCancellation", "") |
| 113 | + : methodName; |
| 114 | + } |
| 115 | + } |
| 116 | +} |
0 commit comments