EventExtensionGenerator.cs

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Text;
using System.Threading;
using Lagrange.Milky.Generator.Extensions;
using Lagrange.Milky.Generator.Utilities;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace Lagrange.Milky.Generator.Events;

[Generator(LanguageNames.CSharp)]
public class EventExtensionGenerator : IIncrementalGenerator
{
    private const string EventConverterAttributeTypeName = "Lagrange.Milky.Events.Attributes.EventConverterAttribute";
    private const string IEventConverterTypeName = "Lagrange.Milky.Events.Converters.IEventConverter`2";

    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        var converterInfos = context.SyntaxProvider.ForAttributeWithMetadataName(
            EventConverterAttributeTypeName,
            predicate: static (node, _) => node is ClassDeclarationSyntax,
            transform: GetConverterInfo
        ).Collect();

        context.RegisterSourceOutput(
            converterInfos,
            static (context, converterInfos) =>
            {
                var validConverterInfos = FilterConverterInfoAndReportDiagnostics(context, converterInfos);

                string generatedCode = GenerateSourceCode(validConverterInfos);
                context.AddSource("Lagrange.Milky.Events.Extensions.EventExtension.g.cs", generatedCode);
            }
        );
    }

    private (string, string?, long, int, Location?) GetConverterInfo(GeneratorAttributeSyntaxContext context, CancellationToken ct)
    {
        var compilation = context.SemanticModel.Compilation;
        var converterSymbol = (INamedTypeSymbol)context.TargetSymbol;

        string converterTypeName = converterSymbol.ToDisplayString();

        var namedArguments = context.Attributes.First().NamedArguments;
        long priority = (long?)namedArguments
            .FirstOrDefault(a => a.Key == "Priority")
            .Value
            .Value
            ?? 0;
        int lifetime = (int?)namedArguments
            .FirstOrDefault(a => a.Key == "Lifetime")
            .Value
            .Value
            ?? 0;

        var converterInterfaceSymbol = compilation.GetTypeByMetadataName(IEventConverterTypeName);
        string? eventTypeName = converterSymbol.AllInterfaces
            .FirstOrDefault(s => s.OriginalDefinition.DefaultEquals(converterInterfaceSymbol))?
            .TypeArguments
            .FirstOrDefault()?
            .ToDisplayString();

        var location = converterSymbol.Locations.FirstOrDefault();

        return (converterTypeName, eventTypeName, priority, lifetime, location);
    }

    private static IEnumerable<(string, string, long, int)> FilterConverterInfoAndReportDiagnostics(SourceProductionContext context, ImmutableArray<(string, string?, long, int, Location?)> converterInfos)
    {
        var eventTypeNames = new HashSet<(string, long)>(StringLongTupleComparer.Default);
        var result = new List<(string, string, long, int)>();

        foreach ((string? converterTypeName, string? eventTypeName, long priority, int lifetime, var location) in converterInfos)
        {
            if (eventTypeName is null)
            {
                context.ReportDiagnostic(Diagnostic.Create(
                    EventDiagnosticDescriptors.MustImplementInterfaceError,
                    location,
                    converterTypeName
                ));
                continue;
            }

            if (!eventTypeNames.Add((eventTypeName, priority)))
            {
                context.ReportDiagnostic(Diagnostic.Create(
                    EventDiagnosticDescriptors.DuplicateEventSerializerError,
                    location,
                    eventTypeName,
                    converterTypeName,
                    priority
                ));
                continue;
            }

            result.Add((converterTypeName, eventTypeName, priority, lifetime));
        }

        return result;
    }

    private static string GenerateSourceCode(IEnumerable<(string, string, long Priority, int)> converterInfos)
    {
        converterInfos = [.. converterInfos.OrderBy(i => i.Priority)];

        var builder = new StringBuilder();

        builder.AppendLine("// <auto-generated/>");
        builder.AppendLine();
        builder.AppendLine("namespace Lagrange.Milky.Events.Extensions;");
        builder.AppendLine();
        builder.AppendLine("public static partial class EventExtension");
        builder.AppendLine("{");
        builder.AppendLine("    public static partial void RegisterConvertibleEvents(this Lagrange.Core.BotContext lagrange, Lagrange.Milky.Events.IGenericEventHandler handler)");
        builder.AppendLine("    {");
        foreach ((_, string? eventTypeName, _, _) in converterInfos)
        {
            builder.AppendLine($"        lagrange.EventInvoker.RegisterEvent<{eventTypeName}>(handler.OnEvent);");
        }
        builder.AppendLine("    }");
        builder.AppendLine();
        builder.AppendLine("    public static partial void UnregisterConvertibleEvents(this Lagrange.Core.BotContext lagrange, Lagrange.Milky.Events.IGenericEventHandler handler)");
        builder.AppendLine("    {");
        foreach ((_, string? eventTypeName, _, _) in converterInfos)
        {
            builder.AppendLine($"        lagrange.EventInvoker.UnregisterEvent<{eventTypeName}>(handler.OnEvent);");
        }
        builder.AppendLine("    }");
        builder.AppendLine();
        builder.AppendLine("    public static partial Microsoft.Extensions.DependencyInjection.IServiceCollection AddEventConverters(this Microsoft.Extensions.DependencyInjection.IServiceCollection services)");
        builder.AppendLine("    {");
        foreach (var (converterName, eventName, _, lifetime) in converterInfos)
        {
            builder.AppendLine($"        services.Add(new Microsoft.Extensions.DependencyInjection.ServiceDescriptor(typeof(Lagrange.Milky.Events.Converters.IEventConverter<{eventName}>), typeof({converterName}), (Microsoft.Extensions.DependencyInjection.ServiceLifetime){lifetime}));");
        }
        builder.AppendLine("        return services;");
        builder.AppendLine("    }");
        builder.AppendLine("}");

        return builder.ToString();
    }
}