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

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

前言

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

关键技术：

• SourceGenerator

  • 在编译时动态生成代码的技术。官方文档：源生成器
  • 关于 IncrementalGenerator 的基本使用可以参考 dotnet 用 SourceGenerator 源代码生成技术实现中文编程语言
  • 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 进行游玩