oecs-sharp/OECS.SourceGen/ComponentDiscoveryGenerator.cs

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using System.Collections.Immutable;
using System.Text;

namespace OECS.SourceGen;

/// <summary>
/// Incremental source generator that discovers all component types used with
/// <see cref="OECS.World"/> and generates a strongly-typed registry so that
/// <see cref="OECS.WorldSerializer"/> can save/load without reflection.
/// </summary>
[Generator]
public class ComponentDiscoveryGenerator : IIncrementalGenerator
{
    private static readonly HashSet<string> _componentApiMethods = new()
    {
        // Component management
        "AddComponent",
        "RemoveComponent",
        "HasComponent",
        "GetComponent",
        "TryGetComponent",
        "ReadComponent",
        "MarkModified",

        // Singletons
        "SetSingleton",
        "GetSingleton",
        "ReadSingleton",
        "HasSingleton",
        "RemoveSingleton",

        // Observability
        "ObserveComponentChanges",

        // Query building
        "With",
        "Without",

        // Iteration and relationships
        "Select",
        "GetSources",
    };

    private static readonly HashSet<string> _singletonMethods = new()
    {
        "SetSingleton",
        "GetSingleton",
        "ReadSingleton",
        "HasSingleton",
        "RemoveSingleton",
    };

    private static readonly HashSet<string> _forEachNames = new()
    {
        "ForEach",
    };

    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        // Collect all generic method invocations that reference component types.
        var invocations = context.SyntaxProvider
            .CreateSyntaxProvider(
                predicate: IsCandidateInvocation,
                transform: ExtractComponentType)
            .Where(t => t.FullyQualifiedName != null)
            .Select((t, _) => (t.FullyQualifiedName!, t.AssemblyQualifiedName!, t.IsRelationship, t.IsSingleton));

        // Also collect ForEach type arguments from the World class.
        var forEachCalls = context.SyntaxProvider
            .CreateSyntaxProvider(
                predicate: IsForEachCandidate,
                transform: ExtractForEachTypes)
            .Where(list => list.Length > 0)
            .SelectMany((list, _) => list);

        var allTypes = invocations.Collect().Combine(forEachCalls.Collect());

        context.RegisterSourceOutput(allTypes, GenerateRegistry);
    }

    private static bool IsCandidateInvocation(SyntaxNode node, CancellationToken _)
    {
        if (node is not InvocationExpressionSyntax invocation)
            return false;

        if (invocation.Expression is MemberAccessExpressionSyntax memberAccess)
        {
            if (memberAccess.Name is GenericNameSyntax genericName)
            {
                return _componentApiMethods.Contains(genericName.Identifier.Text);
            }
        }

        return false;
    }

    private static bool IsForEachCandidate(SyntaxNode node, CancellationToken _)
    {
        if (node is not InvocationExpressionSyntax invocation)
            return false;

        if (invocation.Expression is MemberAccessExpressionSyntax memberAccess)
        {
            var name = memberAccess.Name is GenericNameSyntax gn
                ? gn.Identifier.Text
                : memberAccess.Name.Identifier.Text;
            return _forEachNames.Contains(name);
        }

        return false;
    }

    private static (string? FullyQualifiedName, string? AssemblyQualifiedName, bool IsRelationship, bool IsSingleton) ExtractComponentType(
        GeneratorSyntaxContext ctx, CancellationToken _)
    {
        if (ctx.Node is not InvocationExpressionSyntax invocation)
            return default;

        if (invocation.Expression is not MemberAccessExpressionSyntax memberAccess)
            return default;

        if (memberAccess.Name is not GenericNameSyntax genericName)
            return default;

        var typeArg = genericName.TypeArgumentList.Arguments.FirstOrDefault();
        if (typeArg == null)
            return default;

        var symbolInfo = ctx.SemanticModel.GetSymbolInfo(typeArg);
        if (symbolInfo.Symbol is not INamedTypeSymbol typeSymbol)
            return default;

        // Only public struct types are valid components.
        if (!typeSymbol.IsValueType || typeSymbol.IsAbstract)
            return default;
        if (typeSymbol.DeclaredAccessibility != Accessibility.Public)
            return default;

        var fqn = typeSymbol.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat);

        // Assembly-qualified name for stable serialization (without assembly version).
        var aqn = typeSymbol.ContainingAssembly != null
            ? $"{typeSymbol.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat.WithGlobalNamespaceStyle(SymbolDisplayGlobalNamespaceStyle.Omitted))}, {typeSymbol.ContainingAssembly.Name}"
            : fqn;

        bool isRelationship = typeSymbol.AllInterfaces.Any(i =>
            i.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::OECS.IRelationship");

        bool isSingleton = _singletonMethods.Contains(genericName.Identifier.Text);

        return (fqn, aqn, isRelationship, isSingleton);
    }

    private static ImmutableArray<(string FullyQualifiedName, string AssemblyQualifiedName, bool IsRelationship)> ExtractForEachTypes(
        GeneratorSyntaxContext ctx, CancellationToken _)
    {
        if (ctx.Node is not InvocationExpressionSyntax invocation)
            return ImmutableArray<(string, string, bool)>.Empty;

        if (invocation.Expression is not MemberAccessExpressionSyntax memberAccess)
            return ImmutableArray<(string, string, bool)>.Empty;

        if (memberAccess.Name is not GenericNameSyntax genericName)
            return ImmutableArray<(string, string, bool)>.Empty;

        var types = ImmutableArray.CreateBuilder<(string, string, bool)>();
        foreach (var typeArg in genericName.TypeArgumentList.Arguments)
        {
            var symbolInfo = ctx.SemanticModel.GetSymbolInfo(typeArg);
            if (symbolInfo.Symbol is INamedTypeSymbol typeSymbol &&
                typeSymbol.IsValueType && !typeSymbol.IsAbstract &&
                typeSymbol.DeclaredAccessibility == Accessibility.Public)
            {
                var fqn = typeSymbol.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat);
                var aqn = typeSymbol.ContainingAssembly != null
                    ? $"{typeSymbol.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat.WithGlobalNamespaceStyle(SymbolDisplayGlobalNamespaceStyle.Omitted))}, {typeSymbol.ContainingAssembly.Name}"
                    : fqn;
                bool isRelationship = typeSymbol.AllInterfaces.Any(i =>
                    i.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::OECS.IRelationship");
                types.Add((fqn, aqn, isRelationship));
            }
        }

        return types.ToImmutable();
    }

    private static void GenerateRegistry(SourceProductionContext context,
        (ImmutableArray<(string FullyQualifiedName, string AssemblyQualifiedName, bool IsRelationship, bool IsSingleton)> Left,
         ImmutableArray<(string FullyQualifiedName, string AssemblyQualifiedName, bool IsRelationship)> Right) data)
    {
        var (invocations, forEachTypes) = data;

        // Collect unique types by fully-qualified name, deduplicating.
        var typeMap = new Dictionary<string, (string Fqn, string Aqn, bool IsRel, bool IsSingleton)>();
        foreach (var t in invocations)
        {
            if (typeMap.TryGetValue(t.FullyQualifiedName, out var existing))
                typeMap[t.FullyQualifiedName] = (existing.Fqn, existing.Aqn, existing.IsRel || t.IsRelationship, existing.IsSingleton || t.IsSingleton);
            else
                typeMap[t.FullyQualifiedName] = (t.FullyQualifiedName, t.AssemblyQualifiedName, t.IsRelationship, t.IsSingleton);
        }
        foreach (var t in forEachTypes)
        {
            if (!typeMap.ContainsKey(t.FullyQualifiedName))
                typeMap[t.FullyQualifiedName] = (t.FullyQualifiedName, t.AssemblyQualifiedName, t.IsRelationship, false);
            else if (t.IsRelationship)
            {
                var existing = typeMap[t.FullyQualifiedName];
                typeMap[t.FullyQualifiedName] = (existing.Fqn, existing.Aqn, true, existing.IsSingleton);
            }
        }

        if (typeMap.Count == 0)
            return;

        var sb = new StringBuilder();
        sb.AppendLine("// <auto-generated/>");
        sb.AppendLine("using System.Runtime.CompilerServices;");
        sb.AppendLine("using MessagePack;");
        sb.AppendLine();
        sb.AppendLine("namespace OECS");
        sb.AppendLine("{");
        sb.AppendLine("    file static class ComponentRegistryInitializer");
        sb.AppendLine("    {");
        sb.AppendLine("        [ModuleInitializer]");
        sb.AppendLine("        public static void Initialize()");
        sb.AppendLine("        {");
        sb.AppendLine("            ComponentRegistry.Register(new ComponentDescriptor[]");
        sb.AppendLine("            {");

        bool first = true;
        foreach (var (fqn, aqn, isRel, isSingleton) in typeMap.Values.OrderBy(t => t.Fqn))
        {
            if (!first)
                sb.AppendLine(",");
            first = false;

            sb.AppendLine($"                new ComponentDescriptor(");
            sb.AppendLine($"                    typeName: \"{aqn}\",");
            sb.AppendLine($"                    type: typeof({fqn}),");
            sb.AppendLine($"                    serialize: obj => MessagePackSerializer.Serialize(({fqn})obj),");
            sb.AppendLine($"                    deserializeAndAdd: (world, entity, data) =>");
            sb.AppendLine($"                        world.AddComponent(entity, MessagePackSerializer.Deserialize<{fqn}>(data)),");
            sb.AppendLine($"                    deserialize: data => MessagePackSerializer.Deserialize<{fqn}>(data),");
            sb.AppendLine($"                    isSingleton: {(isSingleton ? "true" : "false")}");

            sb.Append("                )");
        }

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

        context.AddSource("ComponentRegistry.g.cs", sb.ToString());
    }
}