GFramework/GFramework.SourceGenerators/enums/EnumExtensionsGenerator.cs

using System.Linq;
using System.Text;
using GFramework.SourceGenerators.Common.constants;
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 : AttributeEnumGeneratorBase
{
    private static string AttributeMetadataName =>
        $"{PathContests.SourceGeneratorsAbstractionsPath}.enums.GenerateEnumExtensionsAttribute";

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

    protected override AttributeData? ResolveAttribute(Compilation compilation, INamedTypeSymbol symbol)
    {
        var attrSymbol = compilation.GetTypeByMetadataName(AttributeMetadataName);

        if (attrSymbol is null)
            return null;

        return symbol.GetAttributes()
            .FirstOrDefault(a =>
                SymbolEqualityComparer.Default.Equals(a.AttributeClass, attrSymbol));
    }

    protected override bool ValidateSymbol(SourceProductionContext context, Compilation compilation,
        EnumDeclarationSyntax syntax,
        INamedTypeSymbol symbol, AttributeData attr)
    {
        if (symbol.TypeKind == TypeKind.Enum) return true;
        var loc = syntax.Identifier.GetLocation();
        context.ReportDiagnostic(Diagnostic.Create(
            CommonDiagnostics.ClassMustBePartial,
            loc,
            symbol.Name
        ));
        return false;
    }

    /// <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;");

        sb.AppendLine(!string.IsNullOrEmpty(ns) ? $"namespace {ns}" : "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>是否为 {memberName}</summary>");
            sb.AppendLine(
                $"        public static bool Is{memberName}(this {fullEnumName} value) => value == {fullEnumName}.{memberName};");
            sb.AppendLine();
        }

        // 生成 IsIn 方法
        sb.AppendLine("        /// <summary>判断是否属于指定集合</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";
}