ShrinkShared.CodeGen/Editor/ShrinkRegistryILPostProcessor.cs

#nullable enable

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

namespace ShrinkShared.CodeGen
{
    public sealed class ShrinkRegistryILPostProcessor : ILPostProcessor
    {
        public override ILPostProcessor GetInstance() => this;

        public override bool WillProcess(ICompiledAssembly compiledAssembly)
        {
            // 注意：不能因为"引用了 ShrinkEventBus.Runtime"就处理该程序集。
            // 一旦把 ShrinkApp.Core.Runtime 等核心程序集卷入 Cecil 读写，写回的 dll
            // 会被 Unity 判定为 "references itself" 而整条依赖链拒绝加载。
            // 只处理直接承载 Command/Network/App 注册表的程序集。
            return compiledAssembly.References.Any(path =>
                Path.GetFileNameWithoutExtension(path) is "ShrinkCommand.Runtime" or "ShrinkNetwork.Runtime"
                    or "ShrinkApp.Core.Runtime");
        }

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

            var assemblyDefinition = AssemblyDefinitionFor(compiledAssembly);
            var module = assemblyDefinition.MainModule;

            try
            {
                InjectShrinkCommandRegistry(module);
                InjectShrinkNetworkRegistry(module);
                InjectShrinkAppRegistry(module);
            }
            catch (Exception ex)
            {
                diagnostics.Add(new DiagnosticMessage
                {
                    DiagnosticType = DiagnosticType.Error,
                    MessageData = $"[ShrinkShared.CodeGen] {ex.Message}"
                });
            }

            return GetResult(assemblyDefinition, diagnostics);
        }

        private void InjectShrinkCommandRegistry(ModuleDefinition module)
        {
            var subscriberType = FindType(module, "ShrinkCommand.ShrinkCommandSubscriberAttribute", "ShrinkCommand.Runtime");
            var commandAttributeType = FindType(module, "ShrinkCommand.ShrinkCommandAttribute", "ShrinkCommand.Runtime");
            var registryCtor = FindTypeArrayConstructor(module,
                "ShrinkCommand.ShrinkCommandStaticRegistryAttribute",
                "ShrinkCommand.Runtime");

            if (subscriberType == null || commandAttributeType == null || registryCtor == null)
                return;

            var subscriberTypes = module.Types
                .Where(type => HasAttribute(type, subscriberType))
                .Where(type => type.Methods.Any(method => method.IsStatic && HasAttribute(method, commandAttributeType)))
                .Select(type => module.ImportReference(type))
                .ToArray();

            if (subscriberTypes.Length > 0)
                AddAssemblyTypeArrayAttribute(module, registryCtor, subscriberTypes);
        }

        private void InjectShrinkNetworkRegistry(ModuleDefinition module)
        {
            var messageAttributeType = FindType(module, "ShrinkNetwork.ShrinkNetworkMessageAttribute", "ShrinkNetwork.Runtime");
            var subscriberAttributeType = FindType(module, "ShrinkNetwork.ShrinkNetworkSubscriberAttribute", "ShrinkNetwork.Runtime");
            var subscribeAttributeType = FindType(module, "ShrinkNetwork.ShrinkNetworkSubscribeAttribute", "ShrinkNetwork.Runtime");
            var messageCtor = FindTypeArrayConstructor(module, "ShrinkNetwork.ShrinkNetworkMessageRegistryAttribute", "ShrinkNetwork.Runtime");
            var subscriberCtor = FindTypeArrayConstructor(module, "ShrinkNetwork.ShrinkNetworkStaticSubscriberRegistryAttribute", "ShrinkNetwork.Runtime");

            if (messageAttributeType != null && messageCtor != null)
            {
                var messageTypes = module.Types
                    .Where(type => HasAttribute(type, messageAttributeType))
                    .Select(type => module.ImportReference(type))
                    .ToArray();
                if (messageTypes.Length > 0)
                    AddAssemblyTypeArrayAttribute(module, messageCtor, messageTypes);
            }

            if (subscriberAttributeType != null && subscribeAttributeType != null && subscriberCtor != null)
            {
                var subscriberTypes = module.Types
                    .Where(type => HasAttribute(type, subscriberAttributeType))
                    .Where(type => type.Methods.Any(method => method.IsStatic && HasAttribute(method, subscribeAttributeType)))
                    .Select(type => module.ImportReference(type))
                    .ToArray();
                if (subscriberTypes.Length > 0)
                    AddAssemblyTypeArrayAttribute(module, subscriberCtor, subscriberTypes);
            }
        }

        private void InjectShrinkAppRegistry(ModuleDefinition module)
        {
            var installerAttributeType = FindType(module, "ShrinkApp.ShrinkAppModuleInstallerAttribute", "ShrinkApp.Core.Runtime");
            var installerInterfaceType = FindType(module, "ShrinkApp.IShrinkAppModuleInstaller", "ShrinkApp.Core.Runtime");
            var registryCtor = FindTypeArrayConstructor(module, "ShrinkApp.ShrinkAppInstallerRegistryAttribute", "ShrinkApp.Core.Runtime");

            if (installerAttributeType == null || installerInterfaceType == null || registryCtor == null)
                return;

            var installerTypes = module.Types
                .Where(type => !type.IsAbstract)
                .Where(type => HasAttribute(type, installerAttributeType))
                .Where(type => type.Interfaces.Any(item => item.InterfaceType.FullName == installerInterfaceType.FullName))
                .Select(type => module.ImportReference(type))
                .ToArray();

            if (installerTypes.Length > 0)
                AddAssemblyTypeArrayAttribute(module, registryCtor, installerTypes);
        }

        private static bool HasAttribute(ICustomAttributeProvider provider, TypeReference expectedAttributeType)
        {
            return provider.CustomAttributes.Any(attribute => attribute.AttributeType.FullName == expectedAttributeType.FullName);
        }

        private static void AddAssemblyTypeArrayAttribute(ModuleDefinition module, MethodReference ctor, TypeReference[] types)
        {
            var attribute = new CustomAttribute(ctor);
            var typeTypeRef = module.ImportReference(typeof(Type));
            attribute.ConstructorArguments.Add(new CustomAttributeArgument(
                module.ImportReference(typeof(Type[])),
                types.Select(type => new CustomAttributeArgument(typeTypeRef, type)).ToArray()));
            module.Assembly.CustomAttributes.Add(attribute);
        }

        private static TypeReference? FindType(ModuleDefinition module, string fullName, string assemblyName)
        {
            var resolved = Type.GetType($"{fullName}, {assemblyName}", false);
            return resolved == null ? null : module.ImportReference(resolved);
        }

        private static MethodReference? FindTypeArrayConstructor(ModuleDefinition module, string fullName, string assemblyName)
        {
            var typeRef = FindType(module, fullName, assemblyName);
            var typeDef = typeRef?.Resolve();
            var ctor = typeDef?.Methods.FirstOrDefault(method =>
                method.IsConstructor &&
                method.Parameters.Count == 1 &&
                method.Parameters[0].ParameterType.IsArray &&
                method.Parameters[0].ParameterType.GetElementType().FullName == module.ImportReference(typeof(Type)).FullName);
            return ctor == null ? null : module.ImportReference(ctor);
        }

        private static AssemblyDefinition AssemblyDefinitionFor(ICompiledAssembly compiledAssembly)
        {
            var assemblyResolver = new PostProcessorAssemblyResolver(compiledAssembly);
            var readerParameters = new ReaderParameters
            {
                SymbolStream = new MemoryStream(compiledAssembly.InMemoryAssembly.PdbData.ToArray()),
                SymbolReaderProvider = new PdbReaderProvider(),
                AssemblyResolver = assemblyResolver,
                ReflectionImporterProvider = new PostProcessorReflectionImporterProvider(),
                ReadingMode = ReadingMode.Immediate
            };

            var assemblyDefinition = AssemblyDefinition.ReadAssembly(
                new MemoryStream(compiledAssembly.InMemoryAssembly.PeData.ToArray()),
                readerParameters);
            assemblyResolver.AddAssemblyDefinitionBeingOperatedOn(assemblyDefinition);
            return assemblyDefinition;
        }

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