refactor: ILPP to import references from companion modules (#2199)

Author: 0xFA11

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -11,7 +11,8 @@ Additional documentation and release notes are available at [Multiplayer Documen
 
 ### Fixed
 
-- Fixed issue where `NetworkTransform` was not ending extrapolation for the previous state causing non-authoritative instances to become out of synch. (#2170)
+- Fixed ILPP `TypeLoadException` on WebGL on MacOS Editor and potentially other platforms. (#2199)
+- Fixed issue where `NetworkTransform` was not ending extrapolation for the previous state causing non-authoritative instances to become out of sync. (#2170)
 - Fixed issue where `NetworkTransform` was not continuing to interpolate for the remainder of the associated tick period. (#2170)
 - Fixed issue during `NetworkTransform.OnNetworkSpawn` for non-authoritative instances where it was initializing interpolators with the replicated network state which now only contains the transform deltas that occurred during a network tick and not the entire transform state. (#2170)
 - Fixed issue where `NetworkTransform` was not honoring the InLocalSpace property on the authority side during OnNetworkSpawn. (#2170)

com.unity.netcode.gameobjects/Editor/CodeGen/CodeGenHelpers.cs

@@ -15,6 +15,10 @@ namespace Unity.Netcode.Editor.CodeGen
 {
     internal static class CodeGenHelpers
     {
+        public const string DotnetModuleName = "netstandard.dll";
+        public const string UnityModuleName = "UnityEngine.CoreModule.dll";
+        public const string NetcodeModuleName = "Unity.Netcode.Runtime.dll";
+
         public const string RuntimeAssemblyName = "Unity.Netcode.Runtime";
 
         public static readonly string NetworkBehaviour_FullName = typeof(NetworkBehaviour).FullName;
@@ -380,5 +384,68 @@ public static AssemblyDefinition AssemblyDefinitionFor(ICompiledAssembly compile
 
             return assemblyDefinition;
         }
+
+        public static (ModuleDefinition DotnetModule, ModuleDefinition UnityModule, ModuleDefinition NetcodeModule) FindBaseModules(AssemblyDefinition assemblyDefinition, PostProcessorAssemblyResolver assemblyResolver)
+        {
+            ModuleDefinition dotnetModule = null;
+            ModuleDefinition unityModule = null;
+            ModuleDefinition netcodeModule = null;
+
+            foreach (var module in assemblyDefinition.Modules)
+            {
+                if (dotnetModule == null && module.Name == DotnetModuleName)
+                {
+                    dotnetModule = module;
+                    continue;
+                }
+
+                if (unityModule == null && module.Name == UnityModuleName)
+                {
+                    unityModule = module;
+                    continue;
+                }
+
+                if (netcodeModule == null && module.Name == NetcodeModuleName)
+                {
+                    netcodeModule = module;
+                    continue;
+                }
+            }
+
+            if (dotnetModule != null && unityModule != null && netcodeModule != null)
+            {
+                return (dotnetModule, unityModule, netcodeModule);
+            }
+
+            foreach (var assemblyNameReference in assemblyDefinition.MainModule.AssemblyReferences)
+            {
+                foreach (var module in assemblyResolver.Resolve(assemblyNameReference).Modules)
+                {
+                    if (dotnetModule == null && module.Name == DotnetModuleName)
+                    {
+                        dotnetModule = module;
+                        continue;
+                    }
+                    if (unityModule == null && module.Name == UnityModuleName)
+                    {
+                        unityModule = module;
+                        continue;
+                    }
+
+                    if (netcodeModule == null && module.Name == NetcodeModuleName)
+                    {
+                        netcodeModule = module;
+                        continue;
+                    }
+                }
+
+                if (dotnetModule != null && unityModule != null && netcodeModule != null)
+                {
+                    return (dotnetModule, unityModule, netcodeModule);
+                }
+            }
+
+            return (dotnetModule, unityModule, netcodeModule);
+        }
     }
 }

com.unity.netcode.gameobjects/Editor/CodeGen/INetworkMessageILPP.cs

@@ -2,7 +2,6 @@
 using System.IO;
 using System.Linq;
 using System.Collections.Generic;
-using System.Reflection;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
 using Mono.Cecil.Rocks;
@@ -17,8 +16,7 @@ internal sealed class INetworkMessageILPP : ILPPInterface
     {
         public override ILPPInterface GetInstance() => this;
 
-        public override bool WillProcess(ICompiledAssembly compiledAssembly) =>
-            compiledAssembly.Name == CodeGenHelpers.RuntimeAssemblyName;
+        public override bool WillProcess(ICompiledAssembly compiledAssembly) => compiledAssembly.Name == CodeGenHelpers.RuntimeAssemblyName;
 
         private readonly List<DiagnosticMessage> m_Diagnostics = new List<DiagnosticMessage>();
 
@@ -32,13 +30,28 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
             m_Diagnostics.Clear();
 
             // read
-            var assemblyDefinition = CodeGenHelpers.AssemblyDefinitionFor(compiledAssembly, out var resolver);
+            var assemblyDefinition = CodeGenHelpers.AssemblyDefinitionFor(compiledAssembly, out m_AssemblyResolver);
             if (assemblyDefinition == null)
             {
                 m_Diagnostics.AddError($"Cannot read assembly definition: {compiledAssembly.Name}");
                 return null;
             }
 
+            // modules
+            (m_DotnetModule, _, m_NetcodeModule) = CodeGenHelpers.FindBaseModules(assemblyDefinition, m_AssemblyResolver);
+
+            if (m_DotnetModule == null)
+            {
+                m_Diagnostics.AddError($"Cannot find .NET module: {CodeGenHelpers.DotnetModuleName}");
+                return null;
+            }
+
+            if (m_NetcodeModule == null)
+            {
+                m_Diagnostics.AddError($"Cannot find Netcode module: {CodeGenHelpers.NetcodeModuleName}");
+                return null;
+            }
+
             // process
             var mainModule = assemblyDefinition.MainModule;
             if (mainModule != null)
@@ -91,6 +104,9 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
             return new ILPostProcessResult(new InMemoryAssembly(pe.ToArray(), pdb.ToArray()), m_Diagnostics);
         }
 
+        private ModuleDefinition m_DotnetModule;
+        private ModuleDefinition m_NetcodeModule;
+        private PostProcessorAssemblyResolver m_AssemblyResolver;
 
         private MethodReference m_MessagingSystem_ReceiveMessage_MethodRef;
         private TypeReference m_MessagingSystem_MessageWithHandler_TypeRef;
@@ -105,63 +121,108 @@ public override ILPostProcessResult Process(ICompiledAssembly compiledAssembly)
 
         private bool ImportReferences(ModuleDefinition moduleDefinition)
         {
-            m_MessagingSystem_MessageHandler_Constructor_TypeRef = moduleDefinition.ImportReference(typeof(MessagingSystem.MessageHandler).GetConstructors()[0]);
+            TypeDefinition typeTypeDef = null;
+            TypeDefinition listTypeDef = null;
+            foreach (var dotnetTypeDef in m_DotnetModule.GetAllTypes())
+            {
+                if (typeTypeDef == null && dotnetTypeDef.Name == typeof(Type).Name)
+                {
+                    typeTypeDef = dotnetTypeDef;
+                    continue;
+                }
+
+                if (listTypeDef == null && dotnetTypeDef.Name == typeof(List<>).Name)
+                {
+                    listTypeDef = dotnetTypeDef;
+                    continue;
+                }
+            }
+
+            TypeDefinition messageHandlerTypeDef = null;
+            TypeDefinition messageWithHandlerTypeDef = null;
+            TypeDefinition ilppMessageProviderTypeDef = null;
+            TypeDefinition messagingSystemTypeDef = null;
+            foreach (var netcodeTypeDef in m_NetcodeModule.GetAllTypes())
+            {
+                if (messageHandlerTypeDef == null && netcodeTypeDef.Name == nameof(MessagingSystem.MessageHandler))
+                {
+                    messageHandlerTypeDef = netcodeTypeDef;
+                    continue;
+                }
+
+                if (messageWithHandlerTypeDef == null && netcodeTypeDef.Name == nameof(MessagingSystem.MessageWithHandler))
+                {
+                    messageWithHandlerTypeDef = netcodeTypeDef;
+                    continue;
+                }
+
+                if (ilppMessageProviderTypeDef == null && netcodeTypeDef.Name == nameof(ILPPMessageProvider))
+                {
+                    ilppMessageProviderTypeDef = netcodeTypeDef;
+                    continue;
+                }
+
+                if (messagingSystemTypeDef == null && netcodeTypeDef.Name == nameof(MessagingSystem))
+                {
+                    messagingSystemTypeDef = netcodeTypeDef;
+                    continue;
+                }
+            }
+
+            m_MessagingSystem_MessageHandler_Constructor_TypeRef = moduleDefinition.ImportReference(messageHandlerTypeDef.GetConstructors().First());
 
-            var messageWithHandlerType = typeof(MessagingSystem.MessageWithHandler);
-            m_MessagingSystem_MessageWithHandler_TypeRef = moduleDefinition.ImportReference(messageWithHandlerType);
-            foreach (var fieldInfo in messageWithHandlerType.GetFields())
+            m_MessagingSystem_MessageWithHandler_TypeRef = moduleDefinition.ImportReference(messageWithHandlerTypeDef);
+            foreach (var fieldDef in messageWithHandlerTypeDef.Fields)
             {
-                switch (fieldInfo.Name)
+                switch (fieldDef.Name)
                 {
                     case nameof(MessagingSystem.MessageWithHandler.MessageType):
-                        m_MessagingSystem_MessageWithHandler_MessageType_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                        m_MessagingSystem_MessageWithHandler_MessageType_FieldRef = moduleDefinition.ImportReference(fieldDef);
                         break;
                     case nameof(MessagingSystem.MessageWithHandler.Handler):
-                        m_MessagingSystem_MessageWithHandler_Handler_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                        m_MessagingSystem_MessageWithHandler_Handler_FieldRef = moduleDefinition.ImportReference(fieldDef);
                         break;
                 }
             }
 
-            var typeType = typeof(Type);
-            foreach (var methodInfo in typeType.GetMethods())
+            foreach (var methodDef in typeTypeDef.Methods)
             {
-                switch (methodInfo.Name)
+                switch (methodDef.Name)
                 {
                     case nameof(Type.GetTypeFromHandle):
-                        m_Type_GetTypeFromHandle_MethodRef = moduleDefinition.ImportReference(methodInfo);
+                        m_Type_GetTypeFromHandle_MethodRef = moduleDefinition.ImportReference(methodDef);
                         break;
                 }
             }
 
-            var ilppMessageProviderType = typeof(ILPPMessageProvider);
-            foreach (var fieldInfo in ilppMessageProviderType.GetFields(BindingFlags.Static | BindingFlags.NonPublic))
+            foreach (var fieldDef in ilppMessageProviderTypeDef.Fields)
             {
-                switch (fieldInfo.Name)
+                switch (fieldDef.Name)
                 {
                     case nameof(ILPPMessageProvider.__network_message_types):
-                        m_ILPPMessageProvider___network_message_types_FieldRef = moduleDefinition.ImportReference(fieldInfo);
+                        m_ILPPMessageProvider___network_message_types_FieldRef = moduleDefinition.ImportReference(fieldDef);
                         break;
                 }
             }
 
-            var listType = typeof(List<MessagingSystem.MessageWithHandler>);
-            foreach (var methodInfo in listType.GetMethods())
+            foreach (var methodDef in listTypeDef.Methods)
             {
-                switch (methodInfo.Name)
+                switch (methodDef.Name)
                 {
-                    case nameof(List<MessagingSystem.MessageWithHandler>.Add):
-                        m_List_Add_MethodRef = moduleDefinition.ImportReference(methodInfo);
+                    case "Add":
+                        m_List_Add_MethodRef = methodDef;
+                        m_List_Add_MethodRef.DeclaringType = listTypeDef.MakeGenericInstanceType(messageWithHandlerTypeDef);
+                        m_List_Add_MethodRef = moduleDefinition.ImportReference(m_List_Add_MethodRef);
                         break;
                 }
             }
 
-            var messagingSystemType = typeof(MessagingSystem);
-            foreach (var methodInfo in messagingSystemType.GetMethods(BindingFlags.Static | BindingFlags.NonPublic | BindingFlags.Public))
+            foreach (var methodDef in messagingSystemTypeDef.Methods)
             {
-                switch (methodInfo.Name)
+                switch (methodDef.Name)
                 {
                     case k_ReceiveMessageName:
-                        m_MessagingSystem_ReceiveMessage_MethodRef = moduleDefinition.ImportReference(methodInfo);
+                        m_MessagingSystem_ReceiveMessage_MethodRef = moduleDefinition.ImportReference(methodDef);
                         break;
                 }
             }
@@ -217,10 +278,8 @@ private void CreateInstructionsToRegisterType(ILProcessor processor, List<Instru
             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)
+        // 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)

com.unity.netcode.gameobjects/Editor/CodeGen/INetworkSerializableILPP.cs

@@ -10,7 +10,6 @@
 
 namespace Unity.Netcode.Editor.CodeGen
 {
-
     internal sealed class INetworkSerializableILPP : ILPPInterface
     {
         public override ILPPInterface GetInstance() => this;