Async overloads source generator

Ix.NET/Source/Ix.NET.sln

@@ -62,6 +62,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Benchmarks.System.Interacti
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "System.Linq.Async.Ref", "refs\System.Linq.Async.Ref\System.Linq.Async.Ref.csproj", "{1754B36C-D0DB-4E5D-8C30-1F116046DC0F}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "System.Linq.Async.SourceGenerator", "System.Linq.Async.SourceGenerator\System.Linq.Async.SourceGenerator.csproj", "{5C26D649-5ED4-49EE-AFBD-8FA8F12C4AE4}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -136,6 +138,10 @@ Global
 		{1754B36C-D0DB-4E5D-8C30-1F116046DC0F}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{1754B36C-D0DB-4E5D-8C30-1F116046DC0F}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{1754B36C-D0DB-4E5D-8C30-1F116046DC0F}.Release|Any CPU.Build.0 = Release|Any CPU
+		{5C26D649-5ED4-49EE-AFBD-8FA8F12C4AE4}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{5C26D649-5ED4-49EE-AFBD-8FA8F12C4AE4}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{5C26D649-5ED4-49EE-AFBD-8FA8F12C4AE4}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{5C26D649-5ED4-49EE-AFBD-8FA8F12C4AE4}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -158,6 +164,7 @@ Global
 		{2EC0C302-B029-4DDB-AC91-000BF11006AD} = {A3D72E6E-4ADA-42E0-8B2A-055B1F244281}
 		{5DF341BE-B369-4250-AFD4-604DE8C95E45} = {A3D72E6E-4ADA-42E0-8B2A-055B1F244281}
 		{1754B36C-D0DB-4E5D-8C30-1F116046DC0F} = {A3D72E6E-4ADA-42E0-8B2A-055B1F244281}
+		{5C26D649-5ED4-49EE-AFBD-8FA8F12C4AE4} = {80EFE3A1-1414-42EA-949B-1B5370A1B2EA}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {AF70B0C6-C9D9-43B1-9BE4-08720EC1B7B7}

Ix.NET/Source/System.Linq.Async.SourceGenerator/AsyncMethod.cs

@@ -0,0 +1,7 @@
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+namespace System.Linq.Async.SourceGenerator
+{
+    internal sealed record AsyncMethod(IMethodSymbol Symbol, MethodDeclarationSyntax Syntax);
+}

Ix.NET/Source/System.Linq.Async.SourceGenerator/AsyncMethodGrouping.cs

@@ -0,0 +1,8 @@
+using System.Collections.Generic;
+
+using Microsoft.CodeAnalysis;
+
+namespace System.Linq.Async.SourceGenerator
+{
+    internal sealed record AsyncMethodGrouping(SyntaxTree SyntaxTree, IEnumerable<AsyncMethod> Methods);
+}

Ix.NET/Source/System.Linq.Async.SourceGenerator/AsyncOverloadsGenerator.cs

@@ -0,0 +1,116 @@
+using System.Collections.Generic;
+using System.IO;
+using System.Text;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
+
+namespace System.Linq.Async.SourceGenerator
+{
+    [Generator]
+    public sealed class AsyncOverloadsGenerator : ISourceGenerator
+    {
+        private const string AttributeSource =
+            "using System;\n" +
+            "using System.Diagnostics;\n" +
+            "namespace System.Linq\n" +
+            "{\n" +
+            "    [AttributeUsage(AttributeTargets.Method)]\n" +
+            "    [Conditional(\"COMPILE_TIME_ONLY\")]\n" +
+            "    internal sealed class GenerateAsyncOverloadAttribute : Attribute { }\n" +
+            "}\n";
+
+        public void Initialize(GeneratorInitializationContext context)
+        {
+            context.RegisterForSyntaxNotifications(() => new SyntaxReceiver());
+            context.RegisterForPostInitialization(c => c.AddSource("GenerateAsyncOverloadAttribute", AttributeSource));
+        }
+
+        public void Execute(GeneratorExecutionContext context)
+        {
+            if (context.SyntaxReceiver is not SyntaxReceiver syntaxReceiver) return;
+
+            var options = GetGenerationOptions(context);
+            var methodsBySyntaxTree = GetMethodsGroupedBySyntaxTree(context, syntaxReceiver);
+
+            foreach (var grouping in methodsBySyntaxTree)
+                context.AddSource(
+                    $"{Path.GetFileNameWithoutExtension(grouping.SyntaxTree.FilePath)}.AsyncOverloads",
+                    GenerateOverloads(grouping, options));
+        }
+
+        private static GenerationOptions GetGenerationOptions(GeneratorExecutionContext context)
+            => new(SupportFlatAsyncApi: context.ParseOptions.PreprocessorSymbolNames.Contains("SUPPORT_FLAT_ASYNC_API"));
+
+        private static IEnumerable<AsyncMethodGrouping> GetMethodsGroupedBySyntaxTree(GeneratorExecutionContext context, SyntaxReceiver syntaxReceiver)
+            => GetMethodsGroupedBySyntaxTree(
+                context,
+                syntaxReceiver,
+                GetAsyncOverloadAttributeSymbol(context));
+
+        private static string GenerateOverloads(AsyncMethodGrouping grouping, GenerationOptions options)
+        {
+            var usings = grouping.SyntaxTree.GetRoot() is CompilationUnitSyntax compilationUnit
+                ? compilationUnit.Usings.ToString()
+                : string.Empty;
+
+            var overloads = new StringBuilder();
+            overloads.AppendLine("#nullable enable");
+            overloads.AppendLine(usings);
+            overloads.AppendLine("namespace System.Linq");
+            overloads.AppendLine("{");
+            overloads.AppendLine("    partial class AsyncEnumerable");
+            overloads.AppendLine("    {");
+
+            foreach (var method in grouping.Methods)
+                overloads.AppendLine(GenerateOverload(method, options));
+
+            overloads.AppendLine("    }");
+            overloads.AppendLine("}");
+
+            return overloads.ToString();
+        }
+
+        private static string GenerateOverload(AsyncMethod method, GenerationOptions options)
+            => MethodDeclaration(method.Syntax.ReturnType, GetMethodName(method.Symbol, options))
+                .WithModifiers(TokenList(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.StaticKeyword)))
+                .WithTypeParameterList(method.Syntax.TypeParameterList)
+                .WithParameterList(method.Syntax.ParameterList)
+                .WithConstraintClauses(method.Syntax.ConstraintClauses)
+                .WithExpressionBody(ArrowExpressionClause(
+                    InvocationExpression(
+                        IdentifierName(method.Symbol.Name),
+                        ArgumentList(
+                            SeparatedList(
+                                method.Syntax.ParameterList.Parameters
+                                    .Select(p => Argument(IdentifierName(p.Identifier))))))))
+                .WithSemicolonToken(Token(SyntaxKind.SemicolonToken))
+                .WithLeadingTrivia(method.Syntax.GetLeadingTrivia().Where(t => t.GetStructure() is not DirectiveTriviaSyntax))
+                .NormalizeWhitespace()
+                .ToFullString();
+
+        private static INamedTypeSymbol GetAsyncOverloadAttributeSymbol(GeneratorExecutionContext context)
+            => context.Compilation.GetTypeByMetadataName("System.Linq.GenerateAsyncOverloadAttribute") ?? throw new InvalidOperationException();
+
+        private static IEnumerable<AsyncMethodGrouping> GetMethodsGroupedBySyntaxTree(GeneratorExecutionContext context, SyntaxReceiver syntaxReceiver, INamedTypeSymbol attributeSymbol)
+            => from candidate in syntaxReceiver.Candidates
+               group candidate by candidate.SyntaxTree into grouping
+               let model = context.Compilation.GetSemanticModel(grouping.Key)
+               select new AsyncMethodGrouping(
+                  grouping.Key,
+                  from methodSyntax in grouping
+                  let methodSymbol = model.GetDeclaredSymbol(methodSyntax) ?? throw new InvalidOperationException()
+                  where methodSymbol.GetAttributes().Any(a => SymbolEqualityComparer.Default.Equals(a.AttributeClass!, attributeSymbol))
+                  select new AsyncMethod(methodSymbol, methodSyntax));
+
+        private static string GetMethodName(IMethodSymbol methodSymbol, GenerationOptions options)
+        {
+            var methodName = methodSymbol.Name.Replace("Core", "");
+            return options.SupportFlatAsyncApi
+                ? methodName.Replace("Await", "").Replace("WithCancellation", "")
+                : methodName;
+        }
+    }
+}

Ix.NET/Source/System.Linq.Async.SourceGenerator/GenerationOptions.cs

@@ -0,0 +1,4 @@
+namespace System.Linq.Async.SourceGenerator
+{
+    internal sealed record GenerationOptions(bool SupportFlatAsyncApi);
+}

Ix.NET/Source/System.Linq.Async.SourceGenerator/SyntaxReceiver.cs

@@ -0,0 +1,20 @@
+using System.Collections.Generic;
+
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+namespace System.Linq.Async.SourceGenerator
+{
+    internal sealed class SyntaxReceiver : ISyntaxReceiver
+    {
+        public IList<MethodDeclarationSyntax> Candidates { get; } = new List<MethodDeclarationSyntax>();
+
+        public void OnVisitSyntaxNode(SyntaxNode syntaxNode)
+        {
+            if (syntaxNode is MethodDeclarationSyntax { AttributeLists: { Count: >0 } } methodDeclarationSyntax)
+            {
+                Candidates.Add(methodDeclarationSyntax);
+            }
+        }
+    }
+}

Ix.NET/Source/System.Linq.Async.SourceGenerator/System.Linq.Async.SourceGenerator.csproj

@@ -0,0 +1,10 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFramework>netstandard2.0</TargetFramework>
+    <LangVersion>9.0</LangVersion>
+  </PropertyGroup>
+  <ItemGroup>
+    <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="3.9.0" PrivateAssets="all" />
+    <PackageReference Include="IsExternalInit" Version="1.0.0" PrivateAssets="all" />
+  </ItemGroup>
+</Project>

Ix.NET/Source/System.Linq.Async.slnf

@@ -0,0 +1,11 @@
+{
+  "solution": {
+    "path": "Ix.NET.sln",
+    "projects": [
+      "System.Linq.Async\\System.Linq.Async.csproj",
+      "System.Linq.Async.Tests\\System.Linq.Async.Tests.csproj",
+      "System.Linq.Async.SourceGenerator\\System.Linq.Async.SourceGenerator.csproj",
+      "refs\\System.Linq.Async.Ref\\System.Linq.Async.Ref.csproj"
+    ]
+  }
+}

Ix.NET/Source/System.Linq.Async/System.Linq.Async.csproj

@@ -28,6 +28,7 @@
   <ItemGroup>
     <PackageReference Condition="'$(TargetFramework)' != 'netcoreapp3.1' and '$(TargetFramework)' != 'netstandard2.1' " Include="Microsoft.Bcl.AsyncInterfaces" Version="5.0.0" />
     <ReferenceAssemblyProjectReference Include="..\refs\System.Linq.Async.Ref\System.Linq.Async.Ref.csproj" />
+    <ProjectReference Include="..\System.Linq.Async.SourceGenerator\System.Linq.Async.SourceGenerator.csproj" OutputItemType="Analyzer" ReferenceOutputAssembly="false" Private="false" />
   </ItemGroup>
 
   <ItemGroup>