Improve relinking for callsites

Author: SignatureBeef

ModFramework/Modder.cs

@@ -239,6 +239,15 @@ public void RelinkAssembly(string fromAssemblyName, ModuleDefinition? toModule =
         this.RelinkModuleMap[fromAssemblyName] = toModule ?? this.Module;
     }
 
+    public override IMetadataTokenProvider MainRelinker(IMetadataTokenProvider mtp, IGenericParameterProvider context)
+    {
+        // TODO: remove this when monomod is updated - this is missing from the utilised version.
+        if(mtp is CallSite)
+            return Module.ImportReference(mtp);
+
+        return base.MainRelinker(mtp, context);
+    }
+
     public override void Write(Stream? output = null, string? outputPath = null)
     {
         ModContext.Apply(ModType.PreWrite, this);

ModFramework/Relinker/TypeRelinker.cs

@@ -45,6 +45,7 @@ public override void PreWrite()
     bool CheckType<TRef>(TRef type, Action<TRef>? update = null)
         where TRef : TypeReference
     {
+        if (type is null) return false;
         if (NoChangeCache.Contains(type)) return false;
 
         if (ChangeCache.TryGetValue(type, out TypeReference? tref))
@@ -179,6 +180,16 @@ public void Relink(Instruction instr)
             foreach (var ins in instructions)
                 Relink(ins);
         }
+        else if (instr.Operand is CallSite callSite)
+        {
+            CheckType(callSite.ReturnType, nt => callSite.ReturnType = nt);
+
+            foreach (var prm in callSite.Parameters)
+            {
+                CheckType(prm.ParameterType, nt => prm.ParameterType = nt);
+                FixAttributes(prm.CustomAttributes);
+            }
+        }
         else if (!(
             instr.Operand is null
             || instr.Operand is Instruction