refactor: Moved INetworkMessage discovery to ILPP (#1276)

ShadauxCat · web-flow · commit 388d775a7a51 · 2021-10-11T13:16:18.000-07:00
diff --git a/com.unity.netcode.gameobjects/Editor/CodeGen/CodeGenHelpers.cs b/com.unity.netcode.gameobjects/Editor/CodeGen/CodeGenHelpers.cs
@@ -24,6 +24,7 @@ internal static class CodeGenHelpers
         public static readonly string ClientRpcParams_FullName = typeof(ClientRpcParams).FullName;
         public static readonly string INetworkSerializable_FullName = typeof(INetworkSerializable).FullName;
         public static readonly string INetworkSerializable_NetworkSerialize_Name = nameof(INetworkSerializable.NetworkSerialize);
+        public static readonly string IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute_FullName = typeof(IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute).FullName;
         public static readonly string UnityColor_FullName = typeof(Color).FullName;
         public static readonly string UnityColor32_FullName = typeof(Color32).FullName;
         public static readonly string UnityVector2_FullName = typeof(Vector2).FullName;
@@ -265,6 +266,42 @@ public static void AddError(this List<DiagnosticMessage> diagnostics, SequencePo
             });
         }
 
+        public static void RemoveRecursiveReferences(this ModuleDefinition moduleDefinition)
+        {
+            // Weird behavior from Cecil: When importing a reference to a specific implementation of a generic
+            // method, it's importing the main module as a reference into itself. This causes Unity to have issues
+            // when attempting to iterate the assemblies to discover unit tests, as it goes into infinite recursion
+            // and eventually hits a stack overflow. I wasn't able to find any way to stop Cecil from importing the module
+            // into itself, so at the end of it all, we're just going to go back and remove it again.
+            var moduleName = moduleDefinition.Name;
+            if (moduleName.EndsWith(".dll") || moduleName.EndsWith(".exe"))
+            {
+                moduleName = moduleName.Substring(0, moduleName.Length - 4);
+            }
+
+            foreach (var reference in moduleDefinition.AssemblyReferences)
+            {
+                var referenceName = reference.Name.Split(',')[0];
+                if (referenceName.EndsWith(".dll") || referenceName.EndsWith(".exe"))
+                {
+                    referenceName = referenceName.Substring(0, referenceName.Length - 4);
+                }
+
+                if (moduleName == referenceName)
+                {
+                    try
+                    {
+                        moduleDefinition.AssemblyReferences.Remove(reference);
+                        break;
+                    }
+                    catch (Exception)
+                    {
+                        //
+                    }
+                }
+            }
+        }
+
         public static AssemblyDefinition AssemblyDefinitionFor(ICompiledAssembly compiledAssembly, out PostProcessorAssemblyResolver assemblyResolver)
         {
             assemblyResolver = new PostProcessorAssemblyResolver(compiledAssembly);
diff --git a/com.unity.netcode.gameobjects/Editor/CodeGen/INetworkMessageILPP.cs b/com.unity.netcode.gameobjects/Editor/CodeGen/INetworkMessageILPP.cs
@@ -1,11 +1,15 @@
+using System;
 using System.IO;
 using System.Linq;
 using System.Collections.Generic;
+using System.Reflection;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
+using Mono.Cecil.Rocks;
 using Unity.CompilationPipeline.Common.Diagnostics;
 using Unity.CompilationPipeline.Common.ILPostProcessing;
 using ILPPInterface = Unity.CompilationPipeline.Common.ILPostProcessing.ILPostProcessor;
+using MethodAttributes = Mono.Cecil.MethodAttributes;
 
 namespace Unity.Netcode.Editor.CodeGen
 {
@@ -14,7 +18,9 @@ internal sealed class INetworkMessageILPP : ILPPInterface
     {
         public override ILPPInterface GetInstance() => this;
 
-        public override bool WillProcess(ICompiledAssembly compiledAssembly) => compiledAssembly.References.Any(filePath => Path.GetFileNameWithoutExtension(filePath) == CodeGenHelpers.RuntimeAssemblyName);
+        public override bool WillProcess(ICompiledAssembly compiledAssembly) =>
+            compiledAssembly.Name == CodeGenHelpers.RuntimeAssemblyName ||
+            compiledAssembly.References.Any(filePath => Path.GetFileNameWithoutExtension(filePath) == CodeGenHelpers.RuntimeAssemblyName);
 
         private readonly List<DiagnosticMessage> m_Diagnostics = new List<DiagnosticMessage>();
 
@@ -42,11 +48,24 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
             {
                 if (ImportReferences(mainModule))
                 {
+                    var types = mainModule.GetTypes()
+                        .Where(t => t.Resolve().HasInterface(CodeGenHelpers.INetworkMessage_FullName) && !t.Resolve().IsAbstract)
+                        .ToList();
                     // process `INetworkMessage` types
-                    mainModule.GetTypes()
-                        .Where(t => t.HasInterface(CodeGenHelpers.INetworkMessage_FullName))
-                        .ToList()
-                        .ForEach(b => ProcessINetworkMessage(b));
+                    if (types.Count == 0)
+                    {
+                        return null;
+                    }
+
+                    try
+                    {
+                        types.ForEach(b => ProcessINetworkMessage(b));
+                        CreateModuleInitializer(assemblyDefinition, types);
+                    }
+                    catch (Exception e)
+                    {
+                        m_Diagnostics.AddError((e.ToString() + e.StackTrace.ToString()).Replace("\n", "|").Replace("\r", "|"));
+                    }
                 }
                 else
                 {
@@ -58,6 +77,8 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
                 m_Diagnostics.AddError($"Cannot get main module from assembly definition: {compiledAssembly.Name}");
             }
 
+            mainModule.RemoveRecursiveReferences();
+
             // write
             var pe = new MemoryStream();
             var pdb = new MemoryStream();
@@ -77,12 +98,50 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
 
         private TypeReference m_FastBufferReader_TypeRef;
         private TypeReference m_NetworkContext_TypeRef;
+        private FieldReference m_MessagingSystem___network_message_types_FieldRef;
+        private MethodReference m_Type_GetTypeFromHandle_MethodRef;
+
+        private MethodReference m_List_Add_MethodRef;
 
         private bool ImportReferences(ModuleDefinition moduleDefinition)
         {
             m_FastBufferReader_TypeRef = moduleDefinition.ImportReference(typeof(FastBufferReader));
             m_NetworkContext_TypeRef = moduleDefinition.ImportReference(typeof(NetworkContext));
 
+            var typeType = typeof(Type);
+            foreach (var methodInfo in typeType.GetMethods())
+            {
+                switch (methodInfo.Name)
+                {
+                    case nameof(Type.GetTypeFromHandle):
+                        m_Type_GetTypeFromHandle_MethodRef = moduleDefinition.ImportReference(methodInfo);
+                        break;
+                }
+            }
+
+            var messagingSystemType = typeof(MessagingSystem);
+            foreach (var fieldInfo in messagingSystemType.GetFields(BindingFlags.Static | BindingFlags.NonPublic))
+            {
+                switch (fieldInfo.Name)
+                {
+                    case nameof(MessagingSystem.__network_message_types):
+                        m_MessagingSystem___network_message_types_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                        break;
+                }
+            }
+
+            var listType = typeof(List<Type>);
+            foreach (var methodInfo in listType.GetMethods())
+            {
+                switch (methodInfo.Name)
+                {
+                    case nameof(List<Type>.Add):
+                        m_List_Add_MethodRef = moduleDefinition.ImportReference(methodInfo);
+                        break;
+                }
+            }
+
+
             return true;
         }
 
@@ -98,6 +157,7 @@ private void ProcessINetworkMessage(TypeDefinition typeDefinition)
                 {
                     typeSequence = methodSequence;
                 }
+
                 if (resolved.IsStatic && resolved.IsPublic && resolved.Name == "Receive" && resolved.Parameters.Count == 2
                     && !resolved.Parameters[0].IsIn
                     && !resolved.Parameters[0].ParameterType.IsByReference
@@ -118,5 +178,62 @@ private void ProcessINetworkMessage(TypeDefinition typeDefinition)
                 m_Diagnostics.AddError(typeSequence, $"Class {typeDefinition.FullName} does not implement required function: `public static void Receive(FastBufferReader, in NetworkContext)`");
             }
         }
+
+        private MethodDefinition GetOrCreateStaticConstructor(TypeDefinition typeDefinition)
+        {
+            var staticCtorMethodDef = typeDefinition.GetStaticConstructor();
+            if (staticCtorMethodDef == null)
+            {
+                staticCtorMethodDef = new MethodDefinition(
+                    ".cctor", // Static Constructor (constant-constructor)
+                    MethodAttributes.HideBySig |
+                    MethodAttributes.SpecialName |
+                    MethodAttributes.RTSpecialName |
+                    MethodAttributes.Static,
+                    typeDefinition.Module.TypeSystem.Void);
+                staticCtorMethodDef.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
+                typeDefinition.Methods.Add(staticCtorMethodDef);
+            }
+
+            return staticCtorMethodDef;
+        }
+
+        private void CreateInstructionsToRegisterType(ILProcessor processor, List<Instruction> instructions, TypeReference type)
+        {
+            // MessagingSystem.__network_message_types.Add(typeof(type));
+            instructions.Add(processor.Create(OpCodes.Ldsfld, m_MessagingSystem___network_message_types_FieldRef));
+            instructions.Add(processor.Create(OpCodes.Ldtoken, type));
+            instructions.Add(processor.Create(OpCodes.Call, m_Type_GetTypeFromHandle_MethodRef));
+            instructions.Add(processor.Create(OpCodes.Callvirt, m_List_Add_MethodRef));
+        }
+
+        // Creates a static module constructor (which is executed when the module is loaded) that registers all the
+        // message types in the assembly with MessagingSystem.
+        // This is the same behavior as annotating a static method with [ModuleInitializer] in standardized
+        // C# (that attribute doesn't exist in Unity, but the static module constructor still works)
+        // https://docs.microsoft.com/en-us/dotnet/api/system.runtime.compilerservices.moduleinitializerattribute?view=net-5.0
+        // https://web.archive.org/web/20100212140402/http://blogs.msdn.com/junfeng/archive/2005/11/19/494914.aspx
+        private void CreateModuleInitializer(AssemblyDefinition assembly, List<TypeDefinition> networkMessageTypes)
+        {
+            foreach (var typeDefinition in assembly.MainModule.Types)
+            {
+                if (typeDefinition.FullName == "<Module>")
+                {
+                    var staticCtorMethodDef = GetOrCreateStaticConstructor(typeDefinition);
+
+                    var processor = staticCtorMethodDef.Body.GetILProcessor();
+
+                    var instructions = new List<Instruction>();
+
+                    foreach (var type in networkMessageTypes)
+                    {
+                        CreateInstructionsToRegisterType(processor, instructions, type);
+                    }
+
+                    instructions.ForEach(instruction => processor.Body.Instructions.Insert(processor.Body.Instructions.Count - 1, instruction));
+                    break;
+                }
+            }
+        }
     }
 }
diff --git a/com.unity.netcode.gameobjects/Editor/CodeGen/NetworkBehaviourILPP.cs b/com.unity.netcode.gameobjects/Editor/CodeGen/NetworkBehaviourILPP.cs
@@ -467,38 +467,7 @@ private void ProcessNetworkBehaviour(TypeDefinition typeDefinition, string[] ass
                 typeDefinition.Methods.Add(newGetTypeNameMethod);
             }
 
-            // Weird behavior from Cecil: When importing a reference to a specific implementation of a generic
-            // method, it's importing the main module as a reference into itself. This causes Unity to have issues
-            // when attempting to iterate the assemblies to discover unit tests, as it goes into infinite recursion
-            // and eventually hits a stack overflow. I wasn't able to find any way to stop Cecil from importing the module
-            // into itself, so at the end of it all, we're just going to go back and remove it again.
-            var moduleName = m_MainModule.Name;
-            if (moduleName.EndsWith(".dll") || moduleName.EndsWith(".exe"))
-            {
-                moduleName = moduleName.Substring(0, moduleName.Length - 4);
-            }
-
-            foreach (var reference in m_MainModule.AssemblyReferences)
-            {
-                var referenceName = reference.Name.Split(',')[0];
-                if (referenceName.EndsWith(".dll") || referenceName.EndsWith(".exe"))
-                {
-                    referenceName = referenceName.Substring(0, referenceName.Length - 4);
-                }
-
-                if (moduleName == referenceName)
-                {
-                    try
-                    {
-                        m_MainModule.AssemblyReferences.Remove(reference);
-                        break;
-                    }
-                    catch (Exception)
-                    {
-                        //
-                    }
-                }
-            }
+            m_MainModule.RemoveRecursiveReferences();
         }
 
         private CustomAttribute CheckAndGetRpcAttribute(MethodDefinition methodDefinition)
diff --git a/com.unity.netcode.gameobjects/Runtime/Messaging/MessagingSystem.cs b/com.unity.netcode.gameobjects/Runtime/Messaging/MessagingSystem.cs
@@ -17,6 +17,13 @@ public InvalidMessageStructureException(string issue) : base(issue) { }
 
     internal class MessagingSystem : IDisposable
     {
+
+
+#pragma warning disable IDE1006 // disable naming rule violation check
+        // This is NOT modified by RuntimeAccessModifiersILPP right now, but is populated by ILPP.
+        internal static readonly List<Type> __network_message_types = new List<Type>();
+#pragma warning restore IDE1006 // restore naming rule violation check
+
         private struct ReceiveQueueItem
         {
             public FastBufferReader Reader;
@@ -77,54 +84,43 @@ public MessagingSystem(IMessageSender messageSender, object owner, ulong localCl
                 m_MessageSender = messageSender;
                 m_Owner = owner;
 
-                var interfaceType = typeof(INetworkMessage);
-                var implementationTypes = new List<Type>();
-                foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies())
+                var allowedTypes = new List<Type>();
+                foreach (var type in __network_message_types)
                 {
-                    foreach (var type in assembly.GetTypes())
+                    var attributes = type.GetCustomAttributes(typeof(IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute), false);
+                    // If [IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute(ownerType)] isn't provided, it defaults
+                    // to being bound to NetworkManager. This is technically a breach of domain by having
+                    // MessagingSystem know about the existence of NetworkManager... but ultimately,
+                    // IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute is provided to support testing, not to support
+                    // general use of MessagingSystem outside of Netcode for GameObjects, so having MessagingSystem
+                    // know about NetworkManager isn't so bad. Especially since it's just a default value.
+                    // This is just a convenience to keep us and our users from having to use
+                    // [Bind(typeof(NetworkManager))] on every message - only tests that don't want to use
+                    // the full NetworkManager need to worry about it.
+                    var shouldSkip = attributes.Length != 0 || !(m_Owner is NetworkManager);
+                    for (var i = 0; i < attributes.Length; ++i)
                     {
-                        if (type.IsInterface || type.IsAbstract)
+                        var bindAttribute = (IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute)attributes[i];
+                        if (
+                            (bindAttribute.BoundType != null &&
+                             bindAttribute.BoundType.IsInstanceOfType(m_Owner)) ||
+                            (m_Owner == null && bindAttribute.BoundType == null))
                         {
-                            continue;
+                            shouldSkip = false;
+                            break;
                         }
+                    }
 
-                        if (interfaceType.IsAssignableFrom(type))
-                        {
-                            var attributes = type.GetCustomAttributes(typeof(IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute), false);
-                            // If [Bind(ownerType)] isn't provided, it defaults to being bound to NetworkManager
-                            // This is technically a breach of domain by having MessagingSystem know about the existence
-                            // of NetworkManager... but ultimately, Bind is provided to support testing, not to support
-                            // general use of MessagingSystem outside of Netcode for GameObjects, so having MessagingSystem
-                            // know about NetworkManager isn't so bad. Especially since it's just a default value.
-                            // This is just a convenience to keep us and our users from having to use
-                            // [Bind(typeof(NetworkManager))] on every message - only tests that don't want to use
-                            // the full NetworkManager need to worry about it.
-                            var allowedToBind = attributes.Length == 0 && m_Owner is NetworkManager;
-                            for (var i = 0; i < attributes.Length; ++i)
-                            {
-                                var bindAttribute = (IgnoreMessageIfSystemOwnerIsNotOfTypeAttribute)attributes[i];
-                                if (
-                                    (bindAttribute.BoundType != null &&
-                                     bindAttribute.BoundType.IsInstanceOfType(m_Owner)) ||
-                                    (m_Owner == null && bindAttribute.BoundType == null))
-                                {
-                                    allowedToBind = true;
-                                    break;
-                                }
-                            }
-
-                            if (!allowedToBind)
-                            {
-                                continue;
-                            }
-
-                            implementationTypes.Add(type);
-                        }
+                    if (shouldSkip)
+                    {
+                        continue;
                     }
+
+                    allowedTypes.Add(type);
                 }
 
-                implementationTypes.Sort((a, b) => string.CompareOrdinal(a.FullName, b.FullName));
-                foreach (var type in implementationTypes)
+                allowedTypes.Sort((a, b) => string.CompareOrdinal(a.FullName, b.FullName));
+                foreach (var type in allowedTypes)
                 {
                     RegisterMessageType(type);
                 }
