Files
Workspace/Assets/Modules/ShrinkShared.CodeGen/Editor/ShrinkRegistryILPostProcessor.cs
T

402 lines
18 KiB
C#

#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
{
InjectEventBusAutoRegister(module);
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 InjectEventBusAutoRegister(ModuleDefinition module)
{
var subscriberType = FindType(module, "ShrinkEventBus.EventBusSubscriberAttribute", "ShrinkEventBus.Runtime");
var eventBusType = FindType(module, "ShrinkEventBus.EventBus", "ShrinkEventBus.Runtime");
var subscribeAttributeType = FindType(module, "ShrinkEventBus.EventSubscribeAttribute", "ShrinkEventBus.Runtime");
var staticRegistryCtor = FindTypeArrayConstructor(module,
"ShrinkEventBus.EventBusStaticRegistryAttribute",
"ShrinkEventBus.Runtime");
if (subscriberType == null || eventBusType == null)
return;
var eventBusTypeDef = eventBusType.Resolve();
var autoRegisterMethod = eventBusTypeDef?.Methods.FirstOrDefault(method =>
method.Name == "AutoRegister" &&
method.IsStatic &&
method.Parameters.Count == 1);
var unregisterMethod = eventBusTypeDef?.Methods.FirstOrDefault(method =>
method.Name == "UnregisterInstance" &&
method.IsStatic &&
method.Parameters.Count == 1);
if (autoRegisterMethod == null || unregisterMethod == null)
return;
var autoRegisterMethodRef = module.ImportReference(autoRegisterMethod);
var unregisterMethodRef = module.ImportReference(unregisterMethod);
foreach (var type in module.Types)
{
if (!HasAttribute(type, subscriberType) || !InheritsFromMonoBehaviour(type))
continue;
if (subscribeAttributeType != null && !HasInstanceSubscribeMethod(type, subscribeAttributeType))
continue;
InjectAutoRegister(type, module, autoRegisterMethodRef);
InjectAutoUnregister(type, module, unregisterMethodRef);
}
if (subscribeAttributeType != null && staticRegistryCtor != null)
{
var staticSubscriberTypes = module.Types
.Where(type => HasAttribute(type, subscriberType))
.Where(type => type.Methods.Any(method => method.IsStatic && HasAttribute(method, subscribeAttributeType)))
.Select(type => module.ImportReference(type))
.ToArray();
if (staticSubscriberTypes.Length > 0)
AddAssemblyTypeArrayAttribute(module, staticRegistryCtor, staticSubscriberTypes);
}
}
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 bool HasInstanceSubscribeMethod(TypeDefinition type, TypeReference subscribeAttributeType)
{
var current = type;
while (current != null && current.Name != "MonoBehaviour")
{
if (current.Methods.Any(method => !method.IsStatic && HasAttribute(method, subscribeAttributeType)))
return true;
var baseTypeRef = current.BaseType;
if (baseTypeRef == null)
return false;
TypeDefinition? resolvedBase;
try
{
resolvedBase = baseTypeRef.Resolve();
}
catch
{
return true;
}
if (resolvedBase == null)
return true;
current = resolvedBase;
}
return false;
}
private static bool InheritsFromMonoBehaviour(TypeDefinition type)
{
var current = type.BaseType;
while (current != null)
{
if (current.Name == "MonoBehaviour")
return true;
try
{
current = current.Resolve()?.BaseType;
}
catch
{
break;
}
}
return false;
}
private static void InjectAutoRegister(TypeDefinition type, ModuleDefinition module, MethodReference autoRegisterMethodRef)
{
var awake = type.Methods.FirstOrDefault(method => method.Name == "Awake" && !method.IsStatic);
if (awake == null)
{
var baseAwakeRef = FindBaseMethodReference(type, "Awake", module);
if (baseAwakeRef != null)
{
awake = new MethodDefinition("Awake",
MethodAttributes.Family | MethodAttributes.HideBySig | MethodAttributes.Virtual,
module.TypeSystem.Void);
var il = awake.Body.GetILProcessor();
il.Emit(Mono.Cecil.Cil.OpCodes.Ldarg_0);
il.Emit(Mono.Cecil.Cil.OpCodes.Call, baseAwakeRef);
il.Emit(Mono.Cecil.Cil.OpCodes.Ldarg_0);
il.Emit(Mono.Cecil.Cil.OpCodes.Call, autoRegisterMethodRef);
il.Emit(Mono.Cecil.Cil.OpCodes.Ret);
type.Methods.Add(awake);
return;
}
awake = new MethodDefinition("Awake",
MethodAttributes.Private | MethodAttributes.HideBySig,
module.TypeSystem.Void);
var retIl = awake.Body.GetILProcessor();
retIl.Emit(Mono.Cecil.Cil.OpCodes.Ret);
type.Methods.Add(awake);
}
var processor = awake.Body.GetILProcessor();
var instructions = new List<Mono.Cecil.Cil.Instruction>
{
processor.Create(Mono.Cecil.Cil.OpCodes.Ldarg_0),
processor.Create(Mono.Cecil.Cil.OpCodes.Call, autoRegisterMethodRef),
processor.Create(Mono.Cecil.Cil.OpCodes.Nop)
};
instructions.Reverse();
instructions.ForEach(instruction => processor.Body.Instructions.Insert(0, instruction));
}
private static void InjectAutoUnregister(TypeDefinition type, ModuleDefinition module, MethodReference unregisterMethodRef)
{
var onDestroy = type.Methods.FirstOrDefault(method => method.Name == "OnDestroy" && !method.IsStatic);
if (onDestroy == null)
{
var baseOnDestroyRef = FindBaseMethodReference(type, "OnDestroy", module);
if (baseOnDestroyRef != null)
{
onDestroy = new MethodDefinition("OnDestroy",
MethodAttributes.Family | MethodAttributes.HideBySig | MethodAttributes.Virtual,
module.TypeSystem.Void);
var il = onDestroy.Body.GetILProcessor();
il.Emit(Mono.Cecil.Cil.OpCodes.Ldarg_0);
il.Emit(Mono.Cecil.Cil.OpCodes.Call, unregisterMethodRef);
il.Emit(Mono.Cecil.Cil.OpCodes.Ldarg_0);
il.Emit(Mono.Cecil.Cil.OpCodes.Call, baseOnDestroyRef);
il.Emit(Mono.Cecil.Cil.OpCodes.Ret);
type.Methods.Add(onDestroy);
return;
}
onDestroy = new MethodDefinition("OnDestroy",
MethodAttributes.Private | MethodAttributes.HideBySig,
module.TypeSystem.Void);
var retIl = onDestroy.Body.GetILProcessor();
retIl.Emit(Mono.Cecil.Cil.OpCodes.Ret);
type.Methods.Add(onDestroy);
}
var processor = onDestroy.Body.GetILProcessor();
var instructions = new List<Mono.Cecil.Cil.Instruction>
{
processor.Create(Mono.Cecil.Cil.OpCodes.Ldarg_0),
processor.Create(Mono.Cecil.Cil.OpCodes.Call, unregisterMethodRef),
processor.Create(Mono.Cecil.Cil.OpCodes.Nop)
};
instructions.Reverse();
instructions.ForEach(instruction => processor.Body.Instructions.Insert(0, instruction));
}
private static MethodReference? FindBaseMethodReference(TypeDefinition type, string methodName, ModuleDefinition module)
{
try
{
var baseTypeRef = type.BaseType;
while (baseTypeRef != null)
{
var baseTypeDef = baseTypeRef.Resolve();
if (baseTypeDef == null)
break;
var method = baseTypeDef.Methods.FirstOrDefault(candidate =>
candidate.Name == methodName && !candidate.IsStatic && candidate.IsVirtual);
if (method != null)
return module.ImportReference(method);
baseTypeRef = baseTypeDef.BaseType;
}
}
catch
{
}
return null;
}
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);
}
}
}