如何优化多阶段依赖Compilation的增量生成器缓存逻辑

需求描述 #

生成器需要在引用项目中找到符合特定条件的所有类型，并依据这些类型的方法返回值为每个类型生成对应的类。

示例输入（引用项目中） #

class SuitableClass
{
     public CustomType1 GetCustomType1() { ... }
     public CustomType2 GetCustomType2() { ... }
}

class SuitableClass2
{
     public CustomType3 GetCustomType4() { ... }
     public CustomType4 GetCustomType4() { ... }
}

示例输出 #

public class SuitableClassGenerated
{
    public void Handle(CustomType1) { ... }
    public void Handle(CustomType2) { ... }
}

public class SuitableClass2Generated
{
    public void Handle(CustomType3) { ... }
    public void Handle(CustomType4) { ... }
}

当前生成器代码 #

public void Initialize(IncrementalGeneratorInitializationContext context)
{
    IncrementalValuesProvider<string> typeNames = context
        .CompilationProvider
        .SelectMany((c, _) => GetAllSupportedTypes(c.SourceModule.ReferencedAssemblySymbols)
            .Select(x => x.ContainingNamespace.ToDisplayString() + "." + x.MetadataName));

    IncrementalValuesProvider<Info> generationInput = typeNames
        .Combine(context.CompilationProvider)
        .Select((c, _) =>
        {
            var clientType = c.Right.GetTypeByMetadataName(c.Left);
            if (clientType == null)
                return default;
            return new Info(clientType.Name,
                clientType.GetMembers().OfType<IMethodSymbol>()
                    .Where(m => m is { DeclaredAccessibility: Accessibility.Public, IsStatic: false, ReturnsVoid: false })
                    .Select(m => m.ReturnType.ToDisplayString())
                    .ToImmutableArray());
        })
        .Where(x => x.Methods != null);

    context.RegisterSourceOutput(generationInput,
        (ctx, input) =>
        {
            var code = GenerateCode(input);
            if (code != null)
                ctx.AddSource($"{input.Name}.g.cs", code);
        });
}

record struct Info(string Name, ImmutableArray<string> Methods);

问题点 #

当前实现中两个阶段均直接依赖Compilation，破坏了Roslyn增量生成器的缓存机制。希望实现缓存优化，例如当某类型添加带有新返回类型的方法时，仅重新处理该类型对应的第二阶段，而非全量重新运行整个流水线。

核心问题分析 #

当前流水线的核心问题在于：

1. 第一阶段仅提取类型名字符串，后续阶段需要再次通过Compilation.GetTypeByMetadataName查找符号，不仅增加开销，还导致两次依赖CompilationProvider。
2. CompilationProvider是全局级别的变更源，任何编译变动都会触发整个流水线全量重新处理，完全失去增量缓存能力。

优化方案：直接传递符号，避免重复依赖Compilation #

我们可以直接在第一阶段提取INamedTypeSymbol（而非仅类型名），利用Roslyn的符号缓存特性，仅当目标类型本身发生变化时才重新处理该类型。重写后的代码如下：

public void Initialize(IncrementalGeneratorInitializationContext context)
{
    // 第一阶段：提取符合条件的INamedTypeSymbol，提前过滤无效符号
    IncrementalValuesProvider<INamedTypeSymbol> suitableTypes = context
        .CompilationProvider
        .SelectMany((compilation, _) => 
            GetAllSupportedTypes(compilation.SourceModule.ReferencedAssemblySymbols)
                .Where(type => 
                    type != null && 
                    type.DeclaredAccessibility == Accessibility.Public));

    // 第二阶段：直接从符号提取方法返回类型，无需再依赖Compilation
    IncrementalValuesProvider<Info> generationInput = suitableTypes
        .Select((typeSymbol, _) => 
        {
            var returnTypes = typeSymbol.GetMembers()
                .OfType<IMethodSymbol>()
                .Where(m => 
                    m.DeclaredAccessibility == Accessibility.Public && 
                    !m.IsStatic && 
                    !m.ReturnsVoid)
                .Select(m => m.ReturnType.ToDisplayString())
                .ToImmutableArray();

            return new Info(typeSymbol.Name, returnTypes);
        })
        .Where(info => !info.Methods.IsEmpty);

    // 输出阶段保持不变
    context.RegisterSourceOutput(generationInput,
        (ctx, input) =>
        {
            var code = GenerateCode(input);
            if (code != null)
                ctx.AddSource($"{input.Name}.g.cs", code);
        });
}

record struct Info(string Name, ImmutableArray<string> Methods);

缓存友好性说明 #

1. 符号级增量缓存：INamedTypeSymbol是Roslyn的最小缓存单元，只有当目标类型的定义（如新增/删除方法、修改方法返回值）发生变化时，该符号才会被标记为变更，触发对应类型的第二阶段重新处理，其他类型直接复用缓存结果。
2. 消除重复查询：直接传递符号避免了GetTypeByMetadataName的开销，同时规避了因命名空间变化导致的查找失败风险。
3. 精准增量处理：当仅修改某一个SuitableClass的方法时，只有该类型的生成代码会重新生成，其余类型的生成逻辑完全复用缓存，大幅提升生成器性能。

额外优化建议 #

• 将GetAllSupportedTypes的筛选逻辑（如自定义特性检查）合并到第一阶段的Where中，减少后续需要处理的符号数量。
• 若生成代码需要精确的类型引用（如包含命名空间），可使用m.ReturnType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)来生成完整的类型名称。

内容的提问来源于stack exchange，提问作者Vlad