GFramework/GFramework.SourceGenerators/enums/EnumExtensionsGenerator.cs

using System;
using System.Linq;
using System.Text;
using GFramework.SourceGenerators.Abstractions.enums;
using GFramework.SourceGenerators.Common.diagnostics;
using GFramework.SourceGenerators.Common.generator;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace GFramework.SourceGenerators.enums;

/// <summary>
/// 枚举扩展方法生成器，用于自动生成枚举相关的扩展方法
/// </summary>
[Generator]
public sealed class EnumExtensionsGenerator : AttributeClassGeneratorBase
{
    /// <summary>
    /// 使用强类型 Attribute，替代字符串
    /// </summary>
    protected override Type AttributeType =>
        typeof(GenerateEnumExtensionsAttribute);

    /// <summary>
    /// 仅用于 Syntax 粗筛选
    /// </summary>
    protected override string AttributeShortNameWithoutSuffix => "GenerateEnumExtensions";

    /// <summary>
    /// 验证符号是否为有效的枚举类型
    /// </summary>
    /// <param name="context">源生产上下文</param>
    /// <param name="syntax">类声明语法节点</param>
    /// <param name="symbol">命名类型符号</param>
    /// <param name="attr">属性数据</param>
    /// <returns>验证是否通过</returns>
    protected override bool ValidateSymbol(
        SourceProductionContext context,
        ClassDeclarationSyntax syntax,
        INamedTypeSymbol symbol,
        AttributeData attr)
    {
        if (symbol.TypeKind != TypeKind.Enum)
        {
            var loc = syntax.Identifier.GetLocation();
            context.ReportDiagnostic(Diagnostic.Create(
                CommonDiagnostics.ClassMustBePartial, // 可以定义一个新的 Enum 专用 Diagnostic
                loc,
                symbol.Name
            ));
            return false;
        }

        return true;
    }

    /// <summary>
    /// 生成枚举扩展方法的源代码
    /// </summary>
    /// <param name="symbol">枚举类型符号</param>
    /// <param name="attr">属性数据</param>
    /// <returns>生成的源代码字符串</returns>
    protected override string Generate(INamedTypeSymbol symbol, AttributeData attr)
    {
        var ns = symbol.ContainingNamespace.IsGlobalNamespace
            ? null
            : symbol.ContainingNamespace.ToDisplayString();

        var enumName = symbol.Name;
        var fullEnumName = symbol.ToDisplayString();
        var members = symbol.GetMembers()
            .OfType<IFieldSymbol>()
            .Where(f => f.ConstantValue != null)
            .ToArray();

        var sb = new StringBuilder();
        sb.AppendLine("// <auto-generated />");
        sb.AppendLine("using System;");

        if (!string.IsNullOrEmpty(ns))
        {
            sb.AppendLine($"namespace {ns}");
            sb.AppendLine("{");
        }
        else
        {
            sb.AppendLine("namespace EnumExtensionsGenerated");
            sb.AppendLine("{");
        }

        sb.AppendLine($"    public static partial class {enumName}Extensions");
        sb.AppendLine("    {");

        // 生成 IsX 方法
        foreach (var memberName in members.Select(m => m.Name))
        {
            sb.AppendLine($"        /// <summary>Auto-generated: 是否为 {memberName}</summary>");
            sb.AppendLine(
                $"        public static bool Is{memberName}(this {fullEnumName} value) => value == {fullEnumName}.{memberName};");
            sb.AppendLine();
        }

        // 生成 IsIn 方法
        sb.AppendLine("        /// <summary>Auto-generated: 判断是否属于指定集合</summary>");
        sb.AppendLine($"        public static bool IsIn(this {fullEnumName} value, params {fullEnumName}[] values)");
        sb.AppendLine("        {");
        sb.AppendLine("            if (values == null) return false;");
        sb.AppendLine("            foreach (var v in values) if (value == v) return true;");
        sb.AppendLine("            return false;");
        sb.AppendLine("        }");

        sb.AppendLine("    }");
        sb.AppendLine("}"); // namespace

        return sb.ToString();
    }

    /// <summary>
    /// 获取生成文件的提示名称
    /// </summary>
    /// <param name="symbol">命名类型符号</param>
    /// <returns>生成文件的提示名称</returns>
    protected override string GetHintName(INamedTypeSymbol symbol)
        => $"{symbol.Name}.EnumExtensions.g.cs";
}