Godot.Sharp.Extended/NodeBindGenerator.cs

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

namespace Godot.Sharp.Extended.Generators;

[Generator]
public class NodeBindGenerator : IIncrementalGenerator
{
    private const string AttributeSourceCode = """
                                               // <auto-generated />
                                               namespace Godot.Sharp.Extended;
                                               
                                               [System.AttributeUsage(System.AttributeTargets.Property | System.AttributeTargets.Field)]
                                               public class NodeBindAttribute : Attribute
                                               {
                                               }
                                               """;
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        context.RegisterPostInitializationOutput(ctx => ctx.AddSource(
            "NodeBindAttribute.g.cs",
            SourceText.From(AttributeSourceCode, Encoding.UTF8)));

        var propertyProvider = context.SyntaxProvider
            .CreateSyntaxProvider(
                predicate: (s, _) => s is ClassDeclarationSyntax,
                transform: (ctx, _) => GetDeclarationForSourceGen(ctx))
            .Where(t => t.attributeFound)
            .SelectMany((t, _) => t.Item1);

        context.RegisterSourceOutput(context.CompilationProvider.Combine(propertyProvider.Collect()),
            (ctx, value) => GenerateCode(ctx, value.Left, value.Right));
    }

    private (IEnumerable<MemberDeclarationSyntax>, bool attributeFound) GetDeclarationForSourceGen(
        GeneratorSyntaxContext context)
    {
        var classSyntax = (ClassDeclarationSyntax)context.Node;
        var members = classSyntax.Members;
        var foundMembers = members.Select(member => IsNodeBindAttributeOnMember(context.SemanticModel, member))
            .Where(syntax => syntax is not null)
            .ToList();
        return (foundMembers!, foundMembers.Count > 0);
    }

    private MemberDeclarationSyntax? IsNodeBindAttributeOnMember(SemanticModel model,
        MemberDeclarationSyntax declarationSyntax)
    {
        foreach (var attributeSyntax in declarationSyntax
                     .AttributeLists
                     .SelectMany(attributeList => attributeList.Attributes))
        {
            if (model.GetSymbolInfo(attributeSyntax).Symbol is not IMethodSymbol symbol)
                continue;

            var attributeName = symbol.ContainingType.ToDisplayString();
            if (attributeName == $"Godot.Sharp.Extended.NodeBind")
                return declarationSyntax;
        }

        return null;
    }

    private void GenerateCode(SourceProductionContext ctx, Compilation compilation,
        ImmutableArray<MemberDeclarationSyntax> members)
    {
        Dictionary<string, ClassDefinition> classDefinitions = [];

        foreach (var member in members)
        {
            if (member.Parent is ClassDeclarationSyntax classDeclaration)
            {
                var semanticModel = compilation.GetSemanticModel(classDeclaration.SyntaxTree);
                if (semanticModel.GetDeclaredSymbol(classDeclaration) is not INamedTypeSymbol classSymbol)
                    continue;

                var namespaceName = classSymbol.ContainingNamespace.ToDisplayString();
                var className = classDeclaration.Identifier.Text;
                var key = $"{namespaceName}.{className}";
                if (!classDefinitions.ContainsKey(key)) classDefinitions[key] = new ClassDefinition();

                var classDefinition = classDefinitions[key];
                classDefinition.ClassName = className;
                classDefinition.Namespace = namespaceName;
                if (member is FieldDeclarationSyntax field)
                {
                    var variable = field.Declaration.Variables.First();

                    if (semanticModel.GetDeclaredSymbol(variable) is not IFieldSymbol fieldSymbol) continue;

                    var memberDefinition = new MemberDefinition
                    {
                        Name = fieldSymbol.Name,
                        Type = fieldSymbol.Type.Name
                    };
                    classDefinition.MemberDefinitions.Add(memberDefinition);
                    classDefinition.UsingNamespaceName.Add(fieldSymbol.Type.ContainingNamespace.ToDisplayString());
                }
                else if (member is PropertyDeclarationSyntax property)
                {
                    if (semanticModel.GetDeclaredSymbol(property) is not IPropertySymbol propertySymbol) continue;

                    var memberDefinition = new MemberDefinition
                    {
                        Name = propertySymbol.Name,
                        Type = propertySymbol.Type.Name
                    };
                    classDefinition.MemberDefinitions.Add(memberDefinition);
                    classDefinition.UsingNamespaceName.Add(propertySymbol.Type.ContainingNamespace.ToDisplayString());
                }
            }
        }

        foreach (var definition in classDefinitions.Values) GenerateClass(ctx, definition);
    }

    private void GenerateClass(SourceProductionContext ctx, ClassDefinition definition)
    {
        var memberCode = new StringBuilder();
        foreach (var memberDefinition in definition.MemberDefinitions)
        {
            var nodeName = memberDefinition.Name.TrimStart('_').ToPascalCase();

            memberCode.AppendLine($$"""
                                            {{memberDefinition.Name}} = GetNode<{{memberDefinition.Type}}>("%{{nodeName}}");
                                    """);
        }

        var usingCode = new StringBuilder();
        foreach (var namespaceName in definition.UsingNamespaceName) usingCode.AppendLine($"using {namespaceName};");

        var code = $$"""
                     // <auto-generated/>

                     using System;
                     {{usingCode}}

                     namespace {{definition.Namespace}};

                     partial class {{definition.ClassName}}
                     {
                        private void BindNodes() {
                            {{memberCode}}
                        }
                     }
                     """;
        ctx.AddSource($"Godot.Sharp.Extended.{definition.ClassName}.NodeBind.g.cs", SourceText.From(code, Encoding.UTF8));
    }

    private class ClassDefinition
    {
        public string Namespace { get; set; } = "";
        public string ClassName { get; set; } = "";
        public HashSet<string> UsingNamespaceName { get; set; } = [];
        public List<MemberDefinition> MemberDefinitions { get; } = [];
    }

    private class MemberDefinition
    {
        public string Type { get; set; } = "";
        public string Name { get; set; } = "";
    }
}