RobustToolbox/Robust.Shared.EntitySystemSubscriptionsGenerator/EntitySystemSubscriptionGenerator.cs

using System.Collections.Immutable;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Robust.Roslyn.Shared;
using Robust.Roslyn.Shared.Helpers;
using static Robust.Shared.EntitySystemSubscriptionsGenerator.KnownTypes;

namespace Robust.Shared.EntitySystemSubscriptionsGenerator;

/// <summary>
/// This generator implements <c>EntitySystem.AutoSubscriptions()</c> for all <c>EntitySystem</c>s with methods
/// annotated by auto-subscription attributes. In case any attributes are applied to methods incorrectly, this generator
/// just silently ignores them and expects <see cref="EntitySystemSubscriptionGeneratorErrorAnalyzer"/> will complain
/// on its behalf (except in the case of attempting to add generated code to a non-partial type -- we complain about
/// that here).
/// </summary>
/// <seealso cref="EntitySystemSubscriptionGeneratorErrorAnalyzer"/>
[Generator(LanguageNames.CSharp)]
public class EntitySystemSubscriptionGenerator : IIncrementalGenerator
{
    private static readonly DiagnosticDescriptor NotPartial = new(
        Diagnostics.IdNonPartialContainingTypeForGeneratedSubscription,
        "Containing class must be declared as Partial",
        "Method is declared in type \"{0}\" which is not Partial",
        "Usage",
        DiagnosticSeverity.Error,
        true
    );

    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        var annotatedEntitySystems = Aggregate(
                GetEntityTypeCandidatesContainingAnnotatedMethods(context, AllSubscriptionMemberAttributeName),
                GetEntityTypeCandidatesContainingAnnotatedMethods(context, NetworkSubscriptionMemberAttributeName),
                GetEntityTypeCandidatesContainingAnnotatedMethods(context, LocalSubscriptionMemberAttributeName)
            ) // Get all candidate types containing subscription annotated methods
            .SelectMany((array, _) =>
                array.ToImmutableHashSet(PartialTypeInfo.WithoutLocationEqualityComparer)) // Dedupe
            .Combine(context.CompilationProvider)
            .Select((inputs, cancel) =>
                {
                    // For each EntitySystem we've identified as containing subscriptions...

                    var (partialTypeInfo, compilation) = inputs;
                    if (compilation.GetTypeByMetadataName(partialTypeInfo.GetMetadataName()) is not
                        { } entitySystemType)
                        return new EntitySystemInfo(partialTypeInfo, []);

                    // ... check all methods in the type to see if it's a subscription, assembling subscriptions into an array.
                    var subs = entitySystemType.GetMembers()
                        .OfType<IMethodSymbol>()
                        .Select(method =>
                        {
                            cancel.ThrowIfCancellationRequested();
                            return TryParseSubscriptions(method);
                        })
                        .OfType<SubscriptionInfo>()
                        .ToImmutableArray();

                    return new EntitySystemInfo(partialTypeInfo, subs);
                }
            );

        context.RegisterImplementationSourceOutput(
            // Only deal with types that have subscriptions.
            annotatedEntitySystems.Where(it => !it.Subscriptions.IsEmpty),
            (productionContext, info) =>
            {
                var (partialTypeInfo, subscriptions) = info;
                if (partialTypeInfo.CheckPartialDiagnostic(productionContext, NotPartial))
                    return;

                var subscriptionsSyntax = new StringBuilder();
                foreach (var method in subscriptions)
                {
                    productionContext.CancellationToken.ThrowIfCancellationRequested();
                    var subscriptionMethod = method.Type.ToSubscriptionMethod();
                    var typeArgs = string.Join(", ", method.TypeArgs);

                    var before = method.Before.HasValue ? ("[" + string.Join(", ", method.Before.Value.Select(t => $"typeof({t})")) + "]") : "null";
                    var after = method.After.HasValue ? ("[" + string.Join(", ", method.After.Value.Select(t => $"typeof({t})")) + "]") : "null";

                    subscriptionsSyntax.AppendLine($"        {subscriptionMethod}<{typeArgs}>({method.MethodName}, {before}, {after});");
                }

                var builder = new StringBuilder(@"
// <auto-generated />

using Robust.Shared.GameObjects;
using JetBrains.Annotations;

#pragma warning disable CS0618 // Type or member is obsolete: event handlers will have their own obsoletion warnings, dont need to dupe them

");
                partialTypeInfo.WriteHeader(builder);
                builder.AppendLine($@"
{{
    /// <inheritdoc />
    [MustCallBase]
    public override void AutoSubscriptions()
    {{
        base.AutoSubscriptions();

{subscriptionsSyntax}
    }}
}}
");
                partialTypeInfo.WriteFooter(builder);

                productionContext.AddSource(partialTypeInfo.GetGeneratedFileName(), builder.ToString());
            }
        );
    }

    /// Returns <see cref="PartialTypeInfo"/>s for all types in the compilation that contain methods with the given
    /// <paramref name="attributeName">attribute</paramref>.
    private static IncrementalValuesProvider<PartialTypeInfo> GetEntityTypeCandidatesContainingAnnotatedMethods(
        IncrementalGeneratorInitializationContext context,
        string attributeName
    )
    {
        return context.SyntaxProvider.ForAttributeWithMetadataName(
                attributeName,
                (node, _) => node is MethodDeclarationSyntax,
                (ctx, _) =>
                {
                    if (ctx.TargetSymbol is not IMethodSymbol symbol ||
                        ctx.TargetNode is not MethodDeclarationSyntax { Parent: TypeDeclarationSyntax parentSyntax })
                        return null;

                    return PartialTypeInfo.FromSymbol(symbol.ContainingType, parentSyntax);
                })
            .Where(it => it is not null)
            .Select((it, _) => it ?? throw new("Unreachable"));
    }

    /// Tries to parse <paramref name="method"/>'s signature as an even subscription, returning the information required
    /// to make the subscription function call in the generated code. Returns <c>null</c> if the method is not a
    /// subscription, or is a subscription and its signature is invalid.
    private static SubscriptionInfo? TryParseSubscriptions(IMethodSymbol method)
    {
        return TryParseSubscription(
            method,
            AllSubscriptionMemberAttributeName,
            m => TryParseEntityEventHandler(m) ?? TryParseEntitySessionEventHandler(m)
        ) ?? TryParseSubscription(
            method,
            NetworkSubscriptionMemberAttributeName,
            m => TryParseEntityEventHandler(m) ?? TryParseEntitySessionEventHandler(m)
        ) ?? TryParseSubscription(
            method,
            LocalSubscriptionMemberAttributeName,
            m => TryParseEntityEventHandler(m) ??
                 TryParseEntitySessionEventHandler(m) ??
                 TryParseComponentEventHandler(m) ??
                 TryParseEntityEventRefHandler(m)
        );
    }

    /// Tries to parse <paramref name="method"/>'s signature as <c>Robust.Shared.GameObjects.EntityEventHandler</c>.
    /// <returns>The type argument syntax to include in the subscription function call.</returns>
    public static ImmutableArray<string>? TryParseEntityEventHandler(IMethodSymbol method)
    {
        if (method.Parameters.Length != 1 ||
            method.Parameters[0].Type is not INamedTypeSymbol eventType)
            return null;

        return [eventType.ToString()];
    }

    /// Tries to parse <paramref name="method"/>'s signature as <c>Robust.Shared.GameObjects.EntitySessionEventHandler</c>.
    /// <returns>The type argument syntax to include in the subscription function call.</returns>
    public static ImmutableArray<string>? TryParseEntitySessionEventHandler(IMethodSymbol method)
    {
        if (method.Parameters.Length != 2 ||
            method.Parameters[0].Type is not INamedTypeSymbol eventType ||
            !TypeSymbolHelper.ShittyTypeMatch(
                method.Parameters[1].Type,
                EntitySessionEventArgsTypeName
            ))
            return null;

        return [eventType.ToString()];
    }

    /// Tries to parse <paramref name="method"/>'s signature as <c>Robust.Shared.GameObjects.EntityEventRefHandler</c>.
    /// <returns>The type argument syntax to include in the subscription function call.</returns>
    public static ImmutableArray<string>? TryParseEntityEventRefHandler(IMethodSymbol method)
    {
        if (method.Parameters.Length != 2 ||
            method.Parameters[0].Type is not INamedTypeSymbol entityType ||
            method.Parameters[1].Type is not INamedTypeSymbol eventType ||
            method.Parameters[1].RefKind != RefKind.Ref)
            return null;

        if (entityType.OriginalDefinition.ToDisplayString() != EntityTypeName ||
            entityType.TypeArguments is not [INamedTypeSymbol componentType] ||
            componentType.NullableAnnotation == NullableAnnotation.Annotated ||
            !TypeSymbolHelper.ImplementsInterface(componentType, IComponentTypeName))
            return null;

        return [componentType.ToString(), eventType.ToString()];
    }

    /// Tries to parse <paramref name="method"/>'s signature as <c>Robust.Shared.GameObjects.ComponentEventHandler</c>.
    /// <returns>The type argument syntax to include in the subscription function call.</returns>
    public static ImmutableArray<string>? TryParseComponentEventHandler(IMethodSymbol method)
    {
        if (method.Parameters.Length != 3 ||
            method.Parameters[0].Type is not INamedTypeSymbol entityUidType ||
            method.Parameters[1].Type is not INamedTypeSymbol componentType ||
            method.Parameters[2].Type is not INamedTypeSymbol eventType ||
            !TypeSymbolHelper.ShittyTypeMatch(entityUidType, EntityUidTypeName) ||
            componentType.NullableAnnotation == NullableAnnotation.Annotated ||
            !TypeSymbolHelper.ImplementsInterface(componentType, IComponentTypeName))
            return null;

        return [componentType.ToString(), eventType.ToString()];
    }

    private static SubscriptionInfo? TryParseSubscription(
        IMethodSymbol method,
        string annotationName,
        Func<IMethodSymbol, ImmutableArray<string>?> parseFunc
    )
    {
        if (annotationName.ToSubscriptionType() is not { } subType ||
            !AttributeHelper.HasAttribute(method, annotationName, out var attribute) ||
            parseFunc(method) is not { } parameters)
            return null;

        var args = attribute.ConstructorArguments;
        return new SubscriptionInfo(method.Name, subType, parameters, GetTypes(args[0]), GetTypes(args[1]));
    }

    /// <summary>
    /// Gets an array of type names from the typed constant.
    /// </summary>
    private static ImmutableArray<string>? GetTypes(TypedConstant constant)
    {
        if (constant.IsNull || constant.Kind != TypedConstantKind.Array)
            return null;

        return [.. constant.Values.Select(v => (v.Value as ITypeSymbol)!.ToDisplayString())];
    }

    /// Aggregates all of the <typeparamref name="T"/>s across all the given providers into a single array value
    /// provided by the returned provider.
    private static IncrementalValueProvider<ImmutableArray<T>> Aggregate<T>(
        IncrementalValuesProvider<T> first,
        params IncrementalValuesProvider<T>[] more
    )
    {
        return more.Aggregate(
            first.Collect(),
            (acc, valuesProvider) =>
                acc.Combine(valuesProvider.Collect())
                    .Select((values, _) => values.Left.AddRange(values.Right))
        );
    }

    private record struct EntitySystemInfo(PartialTypeInfo Type, EquatableArray<SubscriptionInfo> Subscriptions);

    private record struct SubscriptionInfo(string MethodName, SubscriptionType Type, EquatableArray<string> TypeArgs, EquatableArray<string>? Before, EquatableArray<string>? After);
}