Nerfed/Nerfed.Runtime.Generator/HookSourceGenerator.cs

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

namespace Nerfed.Runtime.Generator
{
    [Generator]
    public class HookSourceGenerator : ISourceGenerator
    {
        public void Execute(GeneratorExecutionContext context)
        {
            // Ensure the syntax receiver is not null and is of the expected type
            if (context.SyntaxReceiver is not HookSyntaxReceiver syntaxReceiver)
                return;

            // Check if we have collected any hook methods
            List<MethodDeclarationSyntax> hookMethods = syntaxReceiver.HookMethods;
            if (hookMethods == null || !hookMethods.Any())
                return;

            StringBuilder codeBuilder = new StringBuilder();

            codeBuilder.AppendLine("using System;");
            codeBuilder.AppendLine("");
            codeBuilder.AppendLine("namespace Nerfed.Runtime.Generator;");
            codeBuilder.AppendLine("");
            codeBuilder.AppendLine($"// Generated by {typeof(HookSourceGenerator)}");
            codeBuilder.AppendLine("public static class Hook");
            codeBuilder.AppendLine("{");
            codeBuilder.AppendLine("    public static void InvokeHooks()");
            codeBuilder.AppendLine("    {");

            foreach (MethodDeclarationSyntax method in hookMethods)
            {
                SemanticModel model = context.Compilation.GetSemanticModel(method.SyntaxTree);

                if (model.GetDeclaredSymbol(method) is not IMethodSymbol methodSymbol)
                {
                    continue;
                }

                if (methodSymbol.DeclaredAccessibility != Accessibility.Public || !methodSymbol.IsStatic)
                {
                    continue;
                }

                codeBuilder.AppendLine($"        {methodSymbol.ContainingType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)}.{methodSymbol.Name}();");
            }

            codeBuilder.AppendLine("    }");
            codeBuilder.AppendLine("}");

            // Add the generated code to the compilation
            context.AddSource("Hook.g.cs", codeBuilder.ToString());
        }


        public void Initialize(GeneratorInitializationContext context)
        {
            context.RegisterForSyntaxNotifications(() => new HookSyntaxReceiver());
        }

        public class HookSyntaxReceiver : ISyntaxReceiver
        {
            public List<MethodDeclarationSyntax> HookMethods { get; } = new List<MethodDeclarationSyntax>();

            public void OnVisitSyntaxNode(SyntaxNode syntaxNode)
            {
                // Check if the node is a method declaration
                if (syntaxNode is MethodDeclarationSyntax methodDeclaration)
                {
                    // Ensure the method declaration has attribute lists
                    if (methodDeclaration.AttributeLists.Count == 0)
                        return;

                    // Check if the method has the Hook attribute
                    bool hasHookAttribute = methodDeclaration.AttributeLists
                        .SelectMany(attrList => attrList.Attributes)
                        .Any(attr => attr.Name.ToString() == "Hook");

                    if (hasHookAttribute)
                    {
                        HookMethods.Add(methodDeclaration);
                    }
                }
            }
        }
    }
}