[C#]SourceGenerator实战: 对任意对象使用await吧!!!

前言

本文记录一次简单的 SourceGenerator 实战，最终实现可以在代码中 await 任意类型对象，仅供娱乐，请勿在生产环境中使用！！！

关键技术：

• SourceGenerator

  • 在编译时动态生成代码的技术。官方文档：源生成器
  • 关于 IncrementalGenerator 的基本使用可以参考
  • SourceGenerator 除了提供附加文件进行代码生成，还有丰富的语法树、类型等分析API
  • SourceGenerator 只能拓展代码，不能替换代码

• Await anything

  • C#中的 async/await 最终由编译器编译为状态机，其核心逻辑在于 await 对象需要实现符合要求的 GetAwaiter 方法，这个方法可以是 拓展方法
  • 参见官方博客 await anything;

那么要实现对任何对象的 await 我们的思路大概如下：

1. 找到所有的 await 语法
2. 检查 await 的对象是否有 GetAwaiter 方法
3. 为没有 GetAwaiter 方法的对象生成 GetAwaiter 拓展方法

得益于 SourceGenerator 丰富的分析API，我们可以很容易的办到这件事

实现源生成器

GetAwaiter拓展方法模板

我们先来实现一个可以让 TargetType 支持 await 的拓展方法类模板：

using System.Runtime.CompilerServices;

namespace System.Threading.Tasks
{
    public static class GetAwaiterExtension_TargetTypeName
    {
        public static TaskAwaiterFor_TargetTypeName GetAwaiter(this TargetType value)
        {
            return new TaskAwaiterFor_TargetTypeName(value);
        }

        public readonly struct TaskAwaiterFor_TargetTypeName : ICriticalNotifyCompletion, INotifyCompletion
        {
            private readonly TargetType _value;

            public bool IsCompleted { get; } = true;

            public TaskAwaiterFor_TargetTypeName(TargetType value)
            {
                _value = value;
            }

            public TargetType GetResult()
            {
                return _value;
            }

            public void OnCompleted(Action continuation)
            {
                continuation();
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                continuation();
            }
        }
    }
}

• 将类型放在命名空间 System.Threading.Tasks 下，可以在使用的时候不需要额外的命名空间引用；
• 由于我们已经有了需要返回的结果值，所以 Awaiter 的 IsCompleted 始终为 true，GetResult 直接返回结果即可；

分析所有 await 语法，并筛选出需要为其生成 GetAwaiter 方法的类型

1. 先建立一个 IncrementalGenerator

   [Generator(LanguageNames.CSharp)]
   public class GetAwaiterIncrementalGenerator : IIncrementalGenerator
   {
       public void Initialize(IncrementalGeneratorInitializationContext context)
       {
       }
   }
   

2. 在 Initialize 方法中筛选目标类型

   /// 使用语法提供器筛选出所有的 `await` 语法，并获取其类型
   var symbolProvider = context.SyntaxProvider.CreateSyntaxProvider((node, _) => node is AwaitExpressionSyntax //直接判断节点是否为 `AwaitExpressionSyntax` 即可筛选出所有 await 表达式
                                                                    , TransformAwaitExpressionSyntax)  //从 await 表达式中解析出其尚不支持 await 的对象类型符号
                                              .Where(m => m is not null)   //筛选掉无效的项
                                              .WithComparer(SymbolEqualityComparer.Default);   //使用默认的符号比较器进行比较
   

3. 直接使用表达式语法不太方便处理，我们实现表达式语法到类型符号的转换方法 TransformAwaitExpressionSyntax

   private static ITypeSymbol? TransformAwaitExpressionSyntax(GeneratorSyntaxContext generatorSyntaxContext, CancellationToken cancellationToken)
   {
       //经过筛选，到达此处的节点一定是 AwaitExpressionSyntax
       var awaitExpressionSyntax = (AwaitExpressionSyntax)generatorSyntaxContext.Node;
   
       //如果 await 表达式语法的 await 对象仍然是 AwaitExpressionSyntax ，那么跳过此条记录
       //类似 "await await await 1;" 我们直接忽略前两个 await 表达式
       if (awaitExpressionSyntax.Expression is AwaitExpressionSyntax)
       {
           return null;
       }
   
       //使用 `SemanticModel` 可以分析出更具体的符号信息，比如类型，方法等
       //直接使用其提供的 `GetAwaitExpressionInfo` 可以从表达式语法获取 await 的详细信息
       var awaitExpressionInfo = generatorSyntaxContext.SemanticModel.GetAwaitExpressionInfo(awaitExpressionSyntax);
   
       //判断分析结果中此表达式是否包含 `GetAwaiter` 方法，如果不包含，那么我们需要为其生成
       if (awaitExpressionInfo.GetAwaiterMethod is null)
       {
           //`SemanticModel` 的 GetTypeInfo 方法可以获取一个表达式的类型符号信息
           //返回 await 对象的类型符号
           return generatorSyntaxContext.SemanticModel.GetTypeInfo(awaitExpressionSyntax.Expression).Type;
       }
   
       return null;
   }
   

为所有目标类型生成 GetAwaiter 拓展方法

由于只需要为相同类型生成一次 GetAwaiter 方法，所以我们需要将类型符号去重之后进行生成

• 直接将上面的 symbolProvider 传递给 RegisterSourceOutput 方法的话，每次只会处理一个类型符号，我们无法去重
• 调用 symbolProvider 的 Collect 方法，可以将前面步骤筛选出的所有类型符号作为一个集合进行处理

所以注册源码生成器可以这样写：

context.RegisterSourceOutput(symbolProvider.Collect(),  //将筛选的结果作为整体传递
                            (ctx, input) =>
                            {
                                //遍历去重后的类型符号
                                foreach (var item in input.Distinct(SymbolEqualityComparer.Default))
                                {
                                    //为每个去重后的类型生成 `GetAwaiter` 拓展方法
                                }
                            });

接下来使用之前写的拓展方法模板生成每个类型的 GetAwaiter 拓展方法即可：

//获取类型符号的完整访问类型名
var fullyClassName = item!.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat);
//获取不包含无效符号的类名
var className = NormalizeClassName(fullyClassName);
//替换模板中的类型占位符为当前处理的目标类型
var code = templateCode.Replace("TargetTypeName", className)
                       .Replace("TargetType", fullyClassName);

//如果目标类型不是公开类型，那么拓展方法也应该不公开
if (item.DeclaredAccessibility != Accessibility.Public)
{
    code = code.Replace("public static class", "internal static class");
}

//将生成的代码添加到编译中
ctx.AddSource($"GetAwaiterFor_{className}.g.cs", code);

//将类型名称中不能作为类名的符号替换为_
private static string NormalizeClassName(string value)
{
    return value.Replace('.', '_')
                .Replace('<', '_')
                .Replace('>', '_')
                .Replace(' ', '_')
                .Replace(',', '_')
                .Replace(':', '_');
}

到这里我们就实现了所有的功能点，新建项目并引用分析器就可以 await 任何对象了，效果大概如下：

• 代码 - AwaitAnyObject.zip
• 也可以直接安装 NuGet 包 AwaitAnyObject 进行游玩