ShrinkNetwork.Integration.EventBus/CodeGen/Editor/NetworkEventILPostProcessor.cs

#nullable enable

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Mono.Cecil.Pdb;
using ShrinkEventBus.CodeGen;
using Unity.CompilationPipeline.Common.Diagnostics;
using Unity.CompilationPipeline.Common.ILPostProcessing;

namespace ShrinkNetwork.Integration.CodeGen
{
    public sealed class NetworkEventILPostProcessor : ILPostProcessor
    {
        private const string IntegrationAssembly = "ShrinkNetwork.Integration.EventBus";

        public override ILPostProcessor GetInstance() => this;

        public override bool WillProcess(ICompiledAssembly compiledAssembly) =>
            compiledAssembly.References.Any(reference =>
                string.Equals(Path.GetFileNameWithoutExtension(reference), IntegrationAssembly,
                    StringComparison.Ordinal));

        public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
        {
            var diagnostics = new List<DiagnosticMessage>();
            if (!WillProcess(compiledAssembly))
                return new ILPostProcessResult(compiledAssembly.InMemoryAssembly, diagnostics);

            var assembly = ReadAssembly(compiledAssembly);
            try
            {
                InjectRegistrations(assembly.MainModule);
            }
            catch (Exception exception)
            {
                diagnostics.Add(new DiagnosticMessage
                {
                    DiagnosticType = DiagnosticType.Error,
                    MessageData = $"[ShrinkNetwork.EventBus.CodeGen] {exception.Message}"
                });
            }
            return WriteAssembly(assembly, diagnostics);
        }

        private static void InjectRegistrations(ModuleDefinition module)
        {
            var networkEventAttribute = ResolveType(module,
                "ShrinkNetwork.Integration.ShrinkNetworkEventAttribute", IntegrationAssembly);
            var messageAttribute = ResolveType(module,
                "ShrinkNetwork.ShrinkNetworkMessageAttribute", "ShrinkNetwork.Runtime");
            var registryType = ResolveType(module,
                "ShrinkNetwork.Integration.ShrinkNetworkEventRegistry", IntegrationAssembly);
            if (networkEventAttribute == null || messageAttribute == null || registryType == null)
                return;

            var registrations = GetAllTypes(module.Types)
                .Where(type => !type.IsAbstract && HasAttribute(type, networkEventAttribute))
                .Select(type => new
                {
                    Type = type,
                    Message = type.CustomAttributes.FirstOrDefault(attribute =>
                        attribute.AttributeType.FullName == messageAttribute.FullName)
                })
                .Where(item => item.Message != null)
                .ToArray();
            if (registrations.Length == 0)
                return;

            var registerDefinition = registryType.Resolve()?.Methods.Single(method =>
                method.Name == "Register" && method.HasGenericParameters)
                ?? throw new InvalidOperationException("ShrinkNetworkEventRegistry.Register<T> was not found.");
            var registerOpen = module.ImportReference(registerDefinition);

            var bootstrap = new TypeDefinition("ShrinkNetwork.Integration.Generated",
                "ShrinkGeneratedNetworkEventBindings",
                TypeAttributes.Abstract | TypeAttributes.Sealed | TypeAttributes.NotPublic,
                module.TypeSystem.Object);
            module.Types.Add(bootstrap);
            var initialize = new MethodDefinition("Register",
                MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig,
                module.TypeSystem.Void);
            bootstrap.Methods.Add(initialize);
            var il = initialize.Body.GetILProcessor();

            foreach (var item in registrations)
            {
                var attribute = item.Message!;
                if (attribute.ConstructorArguments.Count == 0)
                    throw new InvalidOperationException(
                        $"[ShrinkNetworkMessage] on {item.Type.FullName} has no opcode.");

                var opcode = Convert.ToInt32(attribute.ConstructorArguments[0].Value);
                var route = attribute.ConstructorArguments.Count > 1
                    ? attribute.ConstructorArguments[1].Value as string
                    : null;
                var registerClosed = new GenericInstanceMethod(registerOpen);
                registerClosed.GenericArguments.Add(module.ImportReference(item.Type));
                il.Emit(OpCodes.Ldc_I4, opcode);
                if (route == null)
                    il.Emit(OpCodes.Ldnull);
                else
                    il.Emit(OpCodes.Ldstr, route);
                il.Emit(OpCodes.Call, registerClosed);
            }
            il.Emit(OpCodes.Ret);
            InjectModuleInitializer(module, initialize);
        }

        private static void InjectModuleInitializer(ModuleDefinition module, MethodReference register)
        {
            var moduleType = module.Types.First(type => type.Name == "<Module>");
            var initializer = moduleType.Methods.FirstOrDefault(method => method.Name == ".cctor");
            if (initializer == null)
            {
                initializer = new MethodDefinition(".cctor",
                    MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig |
                    MethodAttributes.SpecialName | MethodAttributes.RTSpecialName,
                    module.TypeSystem.Void);
                initializer.Body.GetILProcessor().Emit(OpCodes.Ret);
                moduleType.Methods.Add(initializer);
            }
            var il = initializer.Body.GetILProcessor();
            il.InsertBefore(initializer.Body.Instructions[0], il.Create(OpCodes.Call, register));
        }

        private static bool HasAttribute(ICustomAttributeProvider provider, TypeReference attributeType) =>
            provider.CustomAttributes.Any(attribute =>
                attribute.AttributeType.FullName == attributeType.FullName);

        private static IEnumerable<TypeDefinition> GetAllTypes(IEnumerable<TypeDefinition> roots)
        {
            foreach (var type in roots)
            {
                yield return type;
                foreach (var nested in GetAllTypes(type.NestedTypes))
                    yield return nested;
            }
        }

        private static TypeReference? ResolveType(ModuleDefinition module, string fullName,
            string assemblyName)
        {
            var reference = module.AssemblyReferences.FirstOrDefault(item => item.Name == assemblyName);
            var definition = reference == null ? null : module.AssemblyResolver.Resolve(reference);
            var type = definition?.MainModule.GetType(fullName);
            return type == null ? null : module.ImportReference(type);
        }

        private static AssemblyDefinition ReadAssembly(ICompiledAssembly compiledAssembly)
        {
            var resolver = new PostProcessorAssemblyResolver(compiledAssembly);
            var parameters = new ReaderParameters
            {
                SymbolStream = new MemoryStream(compiledAssembly.InMemoryAssembly.PdbData.ToArray()),
                SymbolReaderProvider = new PdbReaderProvider(),
                AssemblyResolver = resolver,
                ReflectionImporterProvider = new PostProcessorReflectionImporterProvider(),
                ReadingMode = ReadingMode.Immediate
            };
            var assembly = AssemblyDefinition.ReadAssembly(
                new MemoryStream(compiledAssembly.InMemoryAssembly.PeData.ToArray()), parameters);
            resolver.AddAssemblyDefinitionBeingOperatedOn(assembly);
            return assembly;
        }

        private static ILPostProcessResult WriteAssembly(AssemblyDefinition assembly,
            List<DiagnosticMessage> diagnostics)
        {
            var pe = new MemoryStream();
            var pdb = new MemoryStream();
            assembly.Write(pe, new WriterParameters
            {
                SymbolWriterProvider = new PdbWriterProvider(),
                SymbolStream = pdb,
                WriteSymbols = true
            });
            return new ILPostProcessResult(new InMemoryAssembly(pe.ToArray(), pdb.ToArray()), diagnostics);
        }
    }
}