feat: add opt-in MonoBehaviour lifecycle weaving

CodeGen/Editor/EventBusILPostProcessor.cs

@@ -43,12 +43,25 @@ namespace ShrinkEventBus.CodeGen
                     return GetResult(assemblyDefinition, diagnostics);
 
                 foreach (var type in GetAllTypes(module.Types)
-                             .Where(type => !type.IsInterface && !type.IsAbstract)
+                             .Where(type => !type.IsInterface)
                              .Where(type => HasAttribute(type, generatedSubscriberType))
                              .Where(type => type.Methods.Any(method =>
                                  !method.IsStatic && HasAttribute(method, generatedSubscribeType))))
                 {
+                    var subscriberAttribute = type.CustomAttributes.First(attribute =>
+                        attribute.AttributeType.FullName == generatedSubscriberType.FullName);
                     InjectGeneratedBinding(type, module, generatedSubscribeType);
+                    switch (ReadIntProperty(subscriberAttribute, "Lifetime", 0))
+                    {
+                        case 0:
+                            break;
+                        case 1:
+                            InjectAwakeToDestroyLifetime(type, module);
+                            break;
+                        default:
+                            throw new InvalidOperationException(
+                                $"Unsupported ShrinkSubscriberLifetime on {type.FullName}.");
+                    }
                 }
 
                 var staticSubscriberTypes = GetAllTypes(module.Types)
@@ -176,6 +189,244 @@ namespace ShrinkEventBus.CodeGen
             type.Methods.Add(generatedMethod);
         }
 
+        private static void InjectAwakeToDestroyLifetime(TypeDefinition type, ModuleDefinition module)
+        {
+            if (!InheritsFrom(type, "UnityEngine.MonoBehaviour"))
+            {
+                throw new InvalidOperationException(
+                    $"[ShrinkEventSubscriber(Lifetime = AwakeToDestroy)] requires MonoBehaviour: {type.FullName}.");
+            }
+
+            const string bindingFieldName = "__shrinkEventBusAwakeToDestroyBinding";
+            if (type.Fields.Any(field => field.Name == bindingFieldName))
+            {
+                throw new InvalidOperationException(
+                    $"Reserved generated field already exists on {type.FullName}: {bindingFieldName}.");
+            }
+
+            var disposableType = module.ImportReference(typeof(IDisposable));
+            var bindingField = new FieldDefinition(bindingFieldName,
+                FieldAttributes.Private, disposableType);
+            type.Fields.Add(bindingField);
+
+            var eventBusType = FindType(module, "ShrinkEventBus.EventBus", RuntimeAssemblyName)
+                ?? throw new InvalidOperationException("EventBus was not found.");
+            var eventBusDefinition = eventBusType.Resolve()
+                ?? throw new InvalidOperationException("EventBus could not be resolved.");
+            var attachMethod = module.ImportReference(eventBusDefinition.Methods.Single(method =>
+                method.Name == "Attach" && method.IsStatic && method.Parameters.Count == 2));
+            var disposeMethod = module.ImportReference(typeof(IDisposable).GetMethod(nameof(IDisposable.Dispose))
+                ?? throw new InvalidOperationException("IDisposable.Dispose was not found."));
+
+            InjectAwake(type, module, bindingField, attachMethod);
+            InjectOnDestroy(type, module, bindingField, disposeMethod);
+        }
+
+        private static void InjectAwake(TypeDefinition type, ModuleDefinition module,
+            FieldDefinition bindingField, MethodReference attachMethod)
+        {
+            var awake = type.Methods.FirstOrDefault(method =>
+                method.Name == "Awake" && !method.IsStatic && method.Parameters.Count == 0);
+            if (awake != null)
+            {
+                InsertAttachAtStart(awake, module, bindingField, attachMethod);
+                return;
+            }
+
+            var baseAwake = FindBaseMethodReference(type, "Awake", module);
+            awake = new MethodDefinition("Awake",
+                baseAwake != null
+                    ? MethodAttributes.Family | MethodAttributes.HideBySig | MethodAttributes.Virtual
+                    : MethodAttributes.Family | MethodAttributes.HideBySig | MethodAttributes.Virtual |
+                      MethodAttributes.NewSlot,
+                module.TypeSystem.Void);
+            awake.Body.InitLocals = true;
+            var il = awake.Body.GetILProcessor();
+            if (baseAwake != null)
+            {
+                il.Emit(OpCodes.Ldarg_0);
+                il.Emit(OpCodes.Call, baseAwake);
+            }
+            EmitAttach(il, awake, module, bindingField, attachMethod);
+            il.Emit(OpCodes.Ret);
+            type.Methods.Add(awake);
+        }
+
+        private static void InjectOnDestroy(TypeDefinition type, ModuleDefinition module,
+            FieldDefinition bindingField, MethodReference disposeMethod)
+        {
+            var onDestroy = type.Methods.FirstOrDefault(method =>
+                method.Name == "OnDestroy" && !method.IsStatic && method.Parameters.Count == 0);
+            if (onDestroy != null)
+            {
+                InsertDisposeAtStart(onDestroy, bindingField, disposeMethod);
+                return;
+            }
+
+            var baseOnDestroy = FindBaseMethodReference(type, "OnDestroy", module);
+            onDestroy = new MethodDefinition("OnDestroy",
+                baseOnDestroy != null
+                    ? MethodAttributes.Family | MethodAttributes.HideBySig | MethodAttributes.Virtual
+                    : MethodAttributes.Family | MethodAttributes.HideBySig | MethodAttributes.Virtual |
+                      MethodAttributes.NewSlot,
+                module.TypeSystem.Void);
+            var il = onDestroy.Body.GetILProcessor();
+            EmitDispose(il, bindingField, disposeMethod);
+            if (baseOnDestroy != null)
+            {
+                il.Emit(OpCodes.Ldarg_0);
+                il.Emit(OpCodes.Call, baseOnDestroy);
+            }
+            il.Emit(OpCodes.Ret);
+            type.Methods.Add(onDestroy);
+        }
+
+        private static void InsertAttachAtStart(MethodDefinition method, ModuleDefinition module,
+            FieldDefinition bindingField, MethodReference attachMethod)
+        {
+            if (!method.HasBody || method.Body.Instructions.Count == 0)
+                throw new InvalidOperationException($"Awake has no body: {method.FullName}.");
+
+            method.Body.InitLocals = true;
+            var processor = method.Body.GetILProcessor();
+            var instructions = BuildAttachInstructions(processor, method, module, bindingField, attachMethod);
+            var first = method.Body.Instructions[0];
+            foreach (var instruction in instructions)
+                processor.InsertBefore(first, instruction);
+        }
+
+        private static void EmitAttach(ILProcessor il, MethodDefinition method, ModuleDefinition module,
+            FieldDefinition bindingField, MethodReference attachMethod)
+        {
+            foreach (var instruction in BuildAttachInstructions(il, method, module, bindingField, attachMethod))
+                il.Append(instruction);
+        }
+
+        private static IReadOnlyList<Instruction> BuildAttachInstructions(ILProcessor il,
+            MethodDefinition method, ModuleDefinition module, FieldDefinition bindingField,
+            MethodReference attachMethod)
+        {
+            var defaultBusType = module.ImportReference(attachMethod.Parameters[1].ParameterType);
+            var defaultBus = new VariableDefinition(defaultBusType);
+            method.Body.Variables.Add(defaultBus);
+            var attached = il.Create(OpCodes.Nop);
+            return new[]
+            {
+                il.Create(OpCodes.Ldarg_0),
+                il.Create(OpCodes.Ldfld, bindingField),
+                il.Create(OpCodes.Brtrue_S, attached),
+                il.Create(OpCodes.Ldarg_0),
+                il.Create(OpCodes.Ldarg_0),
+                il.Create(OpCodes.Ldloca_S, defaultBus),
+                il.Create(OpCodes.Initobj, defaultBusType),
+                il.Create(OpCodes.Ldloc, defaultBus),
+                il.Create(OpCodes.Call, attachMethod),
+                il.Create(OpCodes.Stfld, bindingField),
+                attached
+            };
+        }
+
+        private static void InsertDisposeAtStart(MethodDefinition method,
+            FieldDefinition bindingField, MethodReference disposeMethod)
+        {
+            if (!method.HasBody || method.Body.Instructions.Count == 0)
+                throw new InvalidOperationException($"OnDestroy has no body: {method.FullName}.");
+
+            var processor = method.Body.GetILProcessor();
+            var instructions = BuildDisposeInstructions(processor, bindingField, disposeMethod);
+            var first = method.Body.Instructions[0];
+            foreach (var instruction in instructions)
+                processor.InsertBefore(first, instruction);
+        }
+
+        private static void EmitDispose(ILProcessor il, FieldDefinition bindingField,
+            MethodReference disposeMethod)
+        {
+            foreach (var instruction in BuildDisposeInstructions(il, bindingField, disposeMethod))
+                il.Append(instruction);
+        }
+
+        private static IReadOnlyList<Instruction> BuildDisposeInstructions(ILProcessor il,
+            FieldDefinition bindingField, MethodReference disposeMethod)
+        {
+            var disposed = il.Create(OpCodes.Nop);
+            return new[]
+            {
+                il.Create(OpCodes.Ldarg_0),
+                il.Create(OpCodes.Ldfld, bindingField),
+                il.Create(OpCodes.Brfalse_S, disposed),
+                il.Create(OpCodes.Ldarg_0),
+                il.Create(OpCodes.Ldfld, bindingField),
+                il.Create(OpCodes.Callvirt, disposeMethod),
+                il.Create(OpCodes.Ldarg_0),
+                il.Create(OpCodes.Ldnull),
+                il.Create(OpCodes.Stfld, bindingField),
+                disposed
+            };
+        }
+
+        private static MethodReference? FindBaseMethodReference(TypeDefinition type,
+            string methodName, ModuleDefinition module)
+        {
+            try
+            {
+                var baseTypeReference = type.BaseType;
+                while (baseTypeReference != null)
+                {
+                    var baseType = baseTypeReference.Resolve();
+                    if (baseType == null)
+                        break;
+
+                    var method = baseType.Methods.FirstOrDefault(candidate =>
+                        candidate.Name == methodName && !candidate.IsStatic && candidate.IsVirtual &&
+                        candidate.Parameters.Count == 0);
+                    if (method != null)
+                    {
+                        if (baseTypeReference is GenericInstanceType genericBase)
+                        {
+                            var methodReference = new MethodReference(method.Name,
+                                module.ImportReference(method.ReturnType), module.ImportReference(genericBase))
+                            {
+                                HasThis = method.HasThis,
+                                ExplicitThis = method.ExplicitThis,
+                                CallingConvention = method.CallingConvention
+                            };
+                            return methodReference;
+                        }
+
+                        return module.ImportReference(method);
+                    }
+
+                    baseTypeReference = baseType.BaseType;
+                }
+            }
+            catch
+            {
+            }
+
+            return null;
+        }
+
+        private static bool InheritsFrom(TypeDefinition type, string expectedFullName)
+        {
+            var current = type.BaseType;
+            while (current != null)
+            {
+                if (current.FullName == expectedFullName)
+                    return true;
+                try
+                {
+                    current = current.Resolve()?.BaseType;
+                }
+                catch
+                {
+                    return false;
+                }
+            }
+
+            return false;
+        }
+
         private static void EmitGeneratedSubscription(ILProcessor il, ModuleDefinition module,
             TypeDefinition ownerType, MethodDefinition handler, CustomAttribute attribute,
             string classDefaultBus, VariableDefinition bindingLocal, MethodReference bindingAdd,