Godot.Sharp.Extended/Generators/ResourceGenerator.cs

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

namespace Godot.Sharp.Extended.Generators;

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

        var classPaths = context.SyntaxProvider
            .CreateSyntaxProvider(
                predicate: (s, _) => s is ClassDeclarationSyntax syntax,
                transform: (ctx, _) => GetClassDeclaration(ctx))
            .Where(t => t.reportAttributeFound)
            .Select((t, _) => (t.Item1, t.path))
            .Collect();

        var provider2 = context.AdditionalTextsProvider.Combine(classPaths)
            .Where(CheckResourcePath)
            .Select((t1, _) =>
            {
                var classDeclarationForPath = GetClassDeclarationForPath(t1);
                return new Storage(t1.Left, classDeclarationForPath.Item1, classDeclarationForPath.path);
            });

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

    private (ClassDeclarationSyntax, string path) GetClassDeclarationForPath((AdditionalText file,
        ImmutableArray<(ClassDeclarationSyntax, string path)> ClassDeclarationSyntax) tuple)
    {
        var realFilePath = tuple.file.Path;
        var relativeFilePath = Path.GetDirectoryName(realFilePath)?.Replace('\\', '/');
        return tuple.ClassDeclarationSyntax.First(data => relativeFilePath!.EndsWith(data.path));
    }

    private bool CheckResourcePath(
        (AdditionalText file, ImmutableArray<(ClassDeclarationSyntax, string path)> Right) tuple)
    {
        var realFilePath = tuple.file.Path;
        var filePathToCheck = tuple.Right.Select(data => data.path);
        var relativeFilePath = Path.GetDirectoryName(realFilePath)?.Replace('\\', '/');
        return filePathToCheck.Any(path => relativeFilePath?.EndsWith(path) ?? false);
    }

    private (ClassDeclarationSyntax classDeclarationSyntax, bool reportAttributeFound, string path) GetClassDeclaration(
        GeneratorSyntaxContext ctx)
    {
        var classDeclarationSyntax = (ClassDeclarationSyntax)ctx.Node;

        foreach (var attributeSyntax in classDeclarationSyntax
                     .AttributeLists
                     .SelectMany(attributeListSyntax => attributeListSyntax.Attributes))
        {
            var symbolInfo = ctx.SemanticModel.GetSymbolInfo(attributeSyntax);
            if (symbolInfo.Symbol is not IMethodSymbol attributeSymbol) continue;
            if (attributeSymbol.ContainingType.ToDisplayString() != "Godot.Sharp.Extended.Attributes.ResourcePath") continue;
            var parameter = attributeSyntax.ArgumentList;
            var pathArgument = parameter?.Arguments[0];

            var path = "";
            if (pathArgument?.Expression is LiteralExpressionSyntax literal) path = literal.Token.ValueText;

            return (classDeclarationSyntax, true, path);
        }
        
        return (classDeclarationSyntax, false, "");
    }

    private void GenerateCode(
        SourceProductionContext context,
        Compilation compilation,
        IEnumerable<Storage> storages)
    {
        Dictionary<string, ClassDefinition> classDefinitions = new();

        foreach (var storage in storages)
        {
            var path = storage.Path;
            var text = storage.Text;
            var classDeclarationSyntax = storage.ClassDeclarationSyntax;
            
            var firstLine = text.GetText()?.Lines[0].ToString();
            if (firstLine == null) continue;
            var scriptNameMatch = Regex.Match(firstLine, @"script_class=""(.+?)""");
            if (!scriptNameMatch.Success) scriptNameMatch = Regex.Match(firstLine, @"type=""(.+?)""");

            var scriptName = scriptNameMatch.Groups[1].Value;
            var fileName = Path.GetFileNameWithoutExtension(path);
            
            var scriptNamespaceType = FindClassSymbol(compilation.GlobalNamespace, scriptName);
            var scriptNamespace = scriptNamespaceType?.ContainingNamespace?.ToDisplayString();
            
            var semanticModel = compilation.GetSemanticModel(classDeclarationSyntax.SyntaxTree);
            if (semanticModel.GetDeclaredSymbol(classDeclarationSyntax) is not INamedTypeSymbol classSymbol) continue;
            
            var namespaceName = classSymbol.ContainingNamespace.ToDisplayString();
            var className = classDeclarationSyntax.Identifier.Text;

            var key = $"{namespaceName}.{className}";
            if (!classDefinitions.ContainsKey(key))
                classDefinitions[key] = new ClassDefinition
                {
                    ClassName = className,
                    NamespaceName = namespaceName,
                };
            
            var def = classDefinitions[key];
            def.Resources.Add(new ResourceDefinition
            {
                Type = scriptName,
                FieldName = fileName,
                FilePath = path + "/" + fileName + ".tres"
            });
            
            if (scriptNamespace != null) def.UsingDefinitionSet.Add(scriptNamespace);
        }

        foreach (var classDefinition in classDefinitions.Values)
        {
            var resources = new StringBuilder();
            var allDefinition = new StringBuilder();
            var usingDefinition = new StringBuilder();
            
            foreach(var def in classDefinition.UsingDefinitionSet) usingDefinition.AppendLine($"using {def};");
            
            if (!classDefinition.UsingDefinitionSet.Contains("Godot")) usingDefinition.AppendLine("using Godot;");
            
            foreach(var resource in classDefinition.Resources)
                resources.AppendLine(
                    $"""
                         public static readonly {resource.Type} {resource.FieldName} =
                             GD.Load<{resource.Type}>(ProjectSettings.GlobalizePath("res://{resource.FilePath}"));
                     """);

            var resourceType = "";
            var generateAll = true;
            foreach (var resource in classDefinition.Resources)
            {
                if (resourceType != "" && resourceType != resource.Type) generateAll = false;

                resourceType = resource.Type;
                allDefinition.AppendLine(
                    $"      {resource.FieldName},");
            }

            var allMethod = "";
            if (generateAll)
                allMethod =
                    $"""
                         public static {resourceType}[] All => [
                     {allDefinition}];
                     """;
            else
                allMethod =
                    $"""
                         public static Resource[] All => [
                     {allDefinition}];
                     """;

            var code =
                $$"""
                  // <auto-generated />
                  using System;
                  {{usingDefinition}}

                  namespace {{classDefinition.NamespaceName}};

                  public partial class {{classDefinition.ClassName}}
                  {
                  {{resources}}
                  {{allMethod}}
                  }
                  """;
            context.AddSource($"Godot.Sharp.Extended.Generators.{classDefinition.ClassName}.Res.g.cs", SourceText.From(code, Encoding.UTF8));
        }
    }
    
    private INamedTypeSymbol? FindClassSymbol(INamespaceSymbol namespaceSymbol, string className)
    {
        var classSymbol = namespaceSymbol.GetTypeMembers().FirstOrDefault(m => m.Name == className);
        if (classSymbol != null) return classSymbol;

        foreach (var childNamespace in namespaceSymbol.GetNamespaceMembers())
        {
            classSymbol = FindClassSymbol(childNamespace, className);
            if (classSymbol != null) return classSymbol;
        }
        
        return null;
    }
    
    private class ClassDefinition
    {
        public HashSet<string> UsingDefinitionSet { get; private set; } = [];
        public string NamespaceName { get; set; } = "";
        public string ClassName { get; set; } = "";
        public List<ResourceDefinition> Resources { get; private set; } = [];
    }
    
    private class ResourceDefinition
    {
        public string Type { get; set; } = "";
        public string FilePath { get; set; } = "";
        public string FieldName { get; set; } = "";
    }
    
    private class Storage(AdditionalText text, ClassDeclarationSyntax classDeclarationSyntax, string path)
    {
        public AdditionalText Text => text;
        public ClassDeclarationSyntax ClassDeclarationSyntax => classDeclarationSyntax;
        public string Path => path;
    }
}