Godot.Sharp.Extended/Generators/NodePropBindGenerator.cs

#nullable enable
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Net.Mime;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Text;

namespace Godot.Sharp.Extended.Generators;

[Generator]
public class NodePropBindGenerator : IIncrementalGenerator
{
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {

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

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

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

    private MemberDeclarationSyntax? IsNodePropBindAttributeOnMember(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.Attributes.NodePropBind")
                return declarationSyntax;
        }

        return null;
    }

    private void GenerateCode(SourceProductionContext context, 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(namespaceName, className, new HashSet<string>(),
                        new List<MemberDefinition>());
                
                var classDefinition = classDefinitions[key];
                if (member is FieldDeclarationSyntax field)
                {
                    var variable = field.Declaration.Variables.First();

                    if (semanticModel.GetDeclaredSymbol(variable) is not IFieldSymbol fieldSymbol) continue;
                    
                    var attributeData = fieldSymbol.GetAttributes()
                        .FirstOrDefault(ad => ad.AttributeClass?.ToDisplayString() == "Godot.Sharp.Extended.Generators.NodePropBindAttribute");

                    if (attributeData is null) continue;

                    var nodeProp = attributeData.ConstructorArguments[0].Value!.ToString();
                    var godotProp = attributeData.ConstructorArguments[1].Value!.ToString();

                    var memberDefinition = new MemberDefinition(fieldSymbol.Type.Name, fieldSymbol.Name, nodeProp, godotProp);
                    classDefinition.MemberDefinitions.Add(memberDefinition);
                    classDefinition.UsingNamespaceName.Add(fieldSymbol.ContainingNamespace.ToDisplayString());
                }
            }
        }
        
        foreach (var definition in classDefinitions.Values) GenerateClass(context, definition);
    }

    private void GenerateClass(SourceProductionContext context, ClassDefinition definition)
    {
        var propCode = new StringBuilder();
        foreach (var memberDefinition in definition.MemberDefinitions)
        {
            var propName = memberDefinition.Name.TrimStart('_').ToPascalCase();
            var godotProp = memberDefinition.GodotProp.ToPascalCase();
            
            propCode.AppendLine($$"""
                                      public {{memberDefinition.Type}} {{propName}}
                                      {
                                          get {
                                              if ({{memberDefinition.NodeProp}} == null)
                                                  return {{memberDefinition.Name}};
                                              return {{memberDefinition.NodeProp}}.{{godotProp}};
                                          }
                                          set {
                                              {{memberDefinition.Name}} = value;
                                              if ({{memberDefinition.NodeProp}} != null)
                                                  {{memberDefinition.NodeProp}}.{{godotProp}} = value;
                                          }
                                      }
                                  """);
        }

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

        var code = $$"""
                    // <auto-generated />
                    
                    using System;
                    using Godot;
                    {{usingCode}}
                    
                    namespace {{definition.Namespace}};
                    
                    partial class {{definition.ClassName}}
                    {
                    {{propCode}}
                    }
                   """;
        context.AddSource($"Godot.Sharp.Extended.Generators.{definition.ClassName}.NodeBind.g.cs", SourceText.From(code, Encoding.UTF8));
    }

    private class ClassDefinition(string theNamespace, string className, HashSet<string> usingNamespaceName, List<MemberDefinition> memberDefinitions)
    {
        public string Namespace => theNamespace;
        public string ClassName => className;
        public HashSet<string> UsingNamespaceName => usingNamespaceName;
        public List<MemberDefinition> MemberDefinitions => memberDefinitions;
    }

    private class MemberDefinition(string type, string name, string nodeProp, string godotProp)
    {
        public string Type => type;
        public string Name => name;
        public string NodeProp => nodeProp;
        public string GodotProp => godotProp;
    }
}