rewrite assembly rewriting, merge Harmony rewriters (#711)

This reduces duplication, decouples it from the assembly loader, and makes it more flexible to handle Harmony rewriting.

6 files changed, 174 insertions, 137 deletions

diff --git a/src/SMAPI/Framework/ModLoading/Rewriters/FieldReplaceRewriter.cs b/src/SMAPI/Framework/ModLoading/Rewriters/FieldReplaceRewriter.cs
index ff86c6e2..8043b13a 100644
--- a/src/SMAPI/Framework/ModLoading/Rewriters/FieldReplaceRewriter.cs
+++ b/src/SMAPI/Framework/ModLoading/Rewriters/FieldReplaceRewriter.cs
@@ -2,16 +2,22 @@ using System;
 using System.Reflection;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
-using StardewModdingAPI.Framework.ModLoading.Finders;
+using StardewModdingAPI.Framework.ModLoading.Framework;
 
 namespace StardewModdingAPI.Framework.ModLoading.Rewriters
 {
     /// <summary>Rewrites references to one field with another.</summary>
-    internal class FieldReplaceRewriter : FieldFinder
+    internal class FieldReplaceRewriter : BaseInstructionHandler
     {
         /*********
         ** Fields
         *********/
+        /// <summary>The type containing the field to which references should be rewritten.</summary>
+        private readonly Type Type;
+
+        /// <summary>The field name to which references should be rewritten.</summary>
+        private readonly string FromFieldName;
+
         /// <summary>The new field to reference.</summary>
         private readonly FieldInfo ToField;
 
@@ -20,31 +26,36 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         ** Public methods
         *********/
         /// <summary>Construct an instance.</summary>
-        /// <param name="type">The type whose field to which references should be rewritten.</param>
+        /// <param name="type">The type whose field to rewrite.</param>
         /// <param name="fromFieldName">The field name to rewrite.</param>
         /// <param name="toFieldName">The new field name to reference.</param>
         public FieldReplaceRewriter(Type type, string fromFieldName, string toFieldName)
-            : base(type.FullName, fromFieldName, InstructionHandleResult.None)
+            : base(defaultPhrase: $"{type.FullName}.{fromFieldName} field")
         {
+            this.Type = type;
+            this.FromFieldName = fromFieldName;
             this.ToField = type.GetField(toFieldName);
             if (this.ToField == null)
                 throw new InvalidOperationException($"The {type.FullName} class doesn't have a {toFieldName} field.");
         }
 
-        /// <summary>Perform the predefined logic for an instruction if applicable.</summary>
+        /// <summary>Rewrite a CIL instruction reference if needed.</summary>
         /// <param name="module">The assembly module containing the instruction.</param>
         /// <param name="cil">The CIL processor.</param>
-        /// <param name="instruction">The instruction to handle.</param>
-        /// <param name="assemblyMap">Metadata for mapping assemblies to the current platform.</param>
-        /// <param name="platformChanged">Whether the mod was compiled on a different platform.</param>
-        public override InstructionHandleResult Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, PlatformAssemblyMap assemblyMap, bool platformChanged)
+        /// <param name="instruction">The CIL instruction to handle.</param>
+        /// <param name="replaceWith">Replaces the CIL instruction with a new one.</param>
+        /// <returns>Returns whether the instruction was changed.</returns>
+        public override bool Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, Action<Instruction> replaceWith)
         {
-            if (!this.IsMatch(instruction))
-                return InstructionHandleResult.None;
+            // get field reference
+            FieldReference fieldRef = RewriteHelper.AsFieldReference(instruction);
+            if (!RewriteHelper.IsFieldReferenceTo(fieldRef, this.Type.FullName, this.FromFieldName))
+                return false;
 
+            // replace with new field
             FieldReference newRef = module.ImportReference(this.ToField);
-            cil.Replace(instruction, cil.Create(instruction.OpCode, newRef));
-            return InstructionHandleResult.Rewritten;
+            replaceWith(cil.Create(instruction.OpCode, newRef));
+            return this.MarkRewritten();
         }
     }
 }
diff --git a/src/SMAPI/Framework/ModLoading/Rewriters/FieldToPropertyRewriter.cs b/src/SMAPI/Framework/ModLoading/Rewriters/FieldToPropertyRewriter.cs
index a43c5e9a..c3b5854e 100644
--- a/src/SMAPI/Framework/ModLoading/Rewriters/FieldToPropertyRewriter.cs
+++ b/src/SMAPI/Framework/ModLoading/Rewriters/FieldToPropertyRewriter.cs
@@ -1,21 +1,24 @@
 using System;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
-using StardewModdingAPI.Framework.ModLoading.Finders;
+using StardewModdingAPI.Framework.ModLoading.Framework;
 
 namespace StardewModdingAPI.Framework.ModLoading.Rewriters
 {
     /// <summary>Rewrites field references into property references.</summary>
-    internal class FieldToPropertyRewriter : FieldFinder
+    internal class FieldToPropertyRewriter : BaseInstructionHandler
     {
         /*********
         ** Fields
         *********/
-        /// <summary>The type whose field to which references should be rewritten.</summary>
+        /// <summary>The type containing the field to which references should be rewritten.</summary>
         private readonly Type Type;
 
-        /// <summary>The property name.</summary>
-        private readonly string PropertyName;
+        /// <summary>The field name to which references should be rewritten.</summary>
+        private readonly string FromFieldName;
+
+        /// <summary>The new property name.</summary>
+        private readonly string ToPropertyName;
 
 
         /*********
@@ -26,10 +29,11 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         /// <param name="fieldName">The field name to rewrite.</param>
         /// <param name="propertyName">The property name (if different).</param>
         public FieldToPropertyRewriter(Type type, string fieldName, string propertyName)
-            : base(type.FullName, fieldName, InstructionHandleResult.None)
+            : base(defaultPhrase: $"{type.FullName}.{fieldName} field")
         {
             this.Type = type;
-            this.PropertyName = propertyName;
+            this.FromFieldName = fieldName;
+            this.ToPropertyName = propertyName;
         }
 
         /// <summary>Construct an instance.</summary>
@@ -38,22 +42,24 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         public FieldToPropertyRewriter(Type type, string fieldName)
             : this(type, fieldName, fieldName) { }
 
-        /// <summary>Perform the predefined logic for an instruction if applicable.</summary>
+        /// <summary>Rewrite a CIL instruction reference if needed.</summary>
         /// <param name="module">The assembly module containing the instruction.</param>
         /// <param name="cil">The CIL processor.</param>
-        /// <param name="instruction">The instruction to handle.</param>
-        /// <param name="assemblyMap">Metadata for mapping assemblies to the current platform.</param>
-        /// <param name="platformChanged">Whether the mod was compiled on a different platform.</param>
-        public override InstructionHandleResult Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, PlatformAssemblyMap assemblyMap, bool platformChanged)
+        /// <param name="instruction">The CIL instruction to handle.</param>
+        /// <param name="replaceWith">Replaces the CIL instruction with a new one.</param>
+        /// <returns>Returns whether the instruction was changed.</returns>
+        public override bool Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, Action<Instruction> replaceWith)
         {
-            if (!this.IsMatch(instruction))
-                return InstructionHandleResult.None;
+            // get field ref
+            FieldReference fieldRef = RewriteHelper.AsFieldReference(instruction);
+            if (!RewriteHelper.IsFieldReferenceTo(fieldRef, this.Type.FullName, this.FromFieldName))
+                return false;
 
+            // replace with property
             string methodPrefix = instruction.OpCode == OpCodes.Ldsfld || instruction.OpCode == OpCodes.Ldfld ? "get" : "set";
-            MethodReference propertyRef = module.ImportReference(this.Type.GetMethod($"{methodPrefix}_{this.PropertyName}"));
-            cil.Replace(instruction, cil.Create(OpCodes.Call, propertyRef));
-
-            return InstructionHandleResult.Rewritten;
+            MethodReference propertyRef = module.ImportReference(this.Type.GetMethod($"{methodPrefix}_{this.ToPropertyName}"));
+            replaceWith(cil.Create(OpCodes.Call, propertyRef));
+            return this.MarkRewritten();
         }
     }
 }
diff --git a/src/SMAPI/Framework/ModLoading/Rewriters/Harmony1AssemblyRewriter.cs b/src/SMAPI/Framework/ModLoading/Rewriters/Harmony1AssemblyRewriter.cs
index 9faca235..a7a0b9c3 100644
--- a/src/SMAPI/Framework/ModLoading/Rewriters/Harmony1AssemblyRewriter.cs
+++ b/src/SMAPI/Framework/ModLoading/Rewriters/Harmony1AssemblyRewriter.cs
@@ -1,28 +1,20 @@
 using System;
+using HarmonyLib;
 using Mono.Cecil;
-using StardewModdingAPI.Framework.ModLoading.Finders;
+using Mono.Cecil.Cil;
 using StardewModdingAPI.Framework.ModLoading.Framework;
+using StardewModdingAPI.Framework.ModLoading.RewriteFacades;
 
 namespace StardewModdingAPI.Framework.ModLoading.Rewriters
 {
     /// <summary>Rewrites Harmony 1.x assembly references to work with Harmony 2.x.</summary>
-    internal class Harmony1AssemblyRewriter : BaseTypeReferenceRewriter
+    internal class Harmony1AssemblyRewriter : BaseInstructionHandler
     {
         /*********
         ** Fields
         *********/
-        /// <summary>The full assembly name to which to find references.</summary>
-        private const string FromAssemblyName = "0Harmony";
-
-        /// <summary>The main Harmony type.</summary>
-        private readonly Type HarmonyType = typeof(HarmonyLib.Harmony);
-
-
-        /*********
-        ** Accessors
-        *********/
-        /// <summary>A brief noun phrase indicating what the rewriter matches.</summary>
-        public const string DefaultNounPhrase = "Harmony 1.x";
+        /// <summary>Whether any Harmony 1.x types were replaced.</summary>
+        private bool ReplacedTypes;
 
 
         /*********
@@ -30,41 +22,80 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         *********/
         /// <summary>Construct an instance.</summary>
         public Harmony1AssemblyRewriter()
-            : base(new TypeAssemblyFinder(Harmony1AssemblyRewriter.FromAssemblyName, InstructionHandleResult.None), Harmony1AssemblyRewriter.DefaultNounPhrase)
-        { }
+            : base(defaultPhrase: "Harmony 1.x") { }
+
+        /// <summary>Rewrite a type reference if needed.</summary>
+        /// <param name="module">The assembly module containing the instruction.</param>
+        /// <param name="type">The type definition to handle.</param>
+        /// <param name="replaceWith">Replaces the type reference with a new one.</param>
+        /// <returns>Returns whether the type was changed.</returns>
+        public override bool Handle(ModuleDefinition module, TypeReference type, Action<TypeReference> replaceWith)
+        {
+            // rewrite Harmony 1.x type to Harmony 2.0 type
+            if (type.Scope is AssemblyNameReference scope && scope.Name == "0Harmony" && scope.Version.Major == 1)
+            {
+                Type targetType = this.GetMappedType(type);
+                replaceWith(module.ImportReference(targetType));
+                this.MarkRewritten();
+                this.ReplacedTypes = true;
+                return true;
+            }
+
+            return false;
+        }
+
+        /// <summary>Rewrite a CIL instruction reference if needed.</summary>
+        /// <param name="module">The assembly module containing the instruction.</param>
+        /// <param name="cil">The CIL processor.</param>
+        /// <param name="instruction">The CIL instruction to handle.</param>
+        /// <param name="replaceWith">Replaces the CIL instruction with a new one.</param>
+        /// <returns>Returns whether the instruction was changed.</returns>
+        public override bool Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, Action<Instruction> replaceWith)
+        {
+            // rewrite Harmony 1.x methods to Harmony 2.0
+            MethodReference methodRef = RewriteHelper.AsMethodReference(instruction);
+            if (this.TryRewriteMethodsToFacade(module, methodRef))
+                return true;
+
+            return false;
+        }
 
 
         /*********
         ** Private methods
         *********/
-        /// <summary>Change a type reference if needed.</summary>
-        /// <param name="module">The assembly module containing the instruction.</param>
-        /// <param name="type">The type to replace if it matches.</param>
-        /// <param name="set">Assign the new type reference.</param>
-        protected override bool RewriteIfNeeded(ModuleDefinition module, TypeReference type, Action<TypeReference> set)
+        /// <summary>Rewrite methods to use Harmony facades if needed.</summary>
+        /// <param name="module">The assembly module containing the method reference.</param>
+        /// <param name="methodRef">The method reference to map.</param>
+        private bool TryRewriteMethodsToFacade(ModuleDefinition module, MethodReference methodRef)
         {
-            bool rewritten = false;
+            if (!this.ReplacedTypes)
+                return false; // not Harmony (or already using Harmony 2.0)
 
-            // current type
-            if (type.Scope.Name == Harmony1AssemblyRewriter.FromAssemblyName && type.Scope is AssemblyNameReference assemblyScope && assemblyScope.Version.Major == 1)
+            // get facade type
+            Type toType;
+            switch (methodRef?.DeclaringType.FullName)
             {
-                Type targetType = this.GetMappedType(type);
-                set(module.ImportReference(targetType));
-                return true;
+                case "HarmonyLib.Harmony":
+                    toType = typeof(HarmonyInstanceMethods);
+                    break;
+
+                case "HarmonyLib.AccessTools":
+                    toType = typeof(AccessToolsMethods);
+                    break;
+
+                default:
+                    return false;
             }
 
-            // recurse into generic arguments
-            if (type is GenericInstanceType genericType)
+            // map if there's a matching method
+            if (RewriteHelper.HasMatchingSignature(toType, methodRef))
             {
-                for (int i = 0; i < genericType.GenericArguments.Count; i++)
-                    rewritten |= this.RewriteIfNeeded(module, genericType.GenericArguments[i], typeRef => genericType.GenericArguments[i] = typeRef);
+                methodRef.DeclaringType = module.ImportReference(toType);
+                return true;
             }
 
-            // recurse into generic parameters (e.g. constraints)
-            for (int i = 0; i < type.GenericParameters.Count; i++)
-                rewritten |= this.RewriteIfNeeded(module, type.GenericParameters[i], typeRef => type.GenericParameters[i] = new GenericParameter(typeRef));
-
-            return rewritten;
+            return false;
         }
 
         /// <summary>Get an equivalent Harmony 2.x type.</summary>
@@ -73,11 +104,11 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         {
             // main Harmony object
             if (type.FullName == "Harmony.HarmonyInstance")
-                return this.HarmonyType;
+                return typeof(Harmony);
 
             // other objects
             string fullName = type.FullName.Replace("Harmony.", "HarmonyLib.");
-            string targetName = this.HarmonyType.AssemblyQualifiedName.Replace(this.HarmonyType.FullName, fullName);
+            string targetName = typeof(Harmony).AssemblyQualifiedName.Replace(typeof(Harmony).FullName, fullName);
             return Type.GetType(targetName, throwOnError: true);
         }
     }
diff --git a/src/SMAPI/Framework/ModLoading/Rewriters/MethodParentRewriter.cs b/src/SMAPI/Framework/ModLoading/Rewriters/MethodParentRewriter.cs
index d0fe8b13..b8e53f40 100644
--- a/src/SMAPI/Framework/ModLoading/Rewriters/MethodParentRewriter.cs
+++ b/src/SMAPI/Framework/ModLoading/Rewriters/MethodParentRewriter.cs
@@ -18,9 +18,6 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         /// <summary>The type with methods to map to.</summary>
         private readonly Type ToType;
 
-        /// <summary>Whether to only rewrite references if loading the assembly on a different platform than it was compiled on.</summary>
-        private readonly bool OnlyIfPlatformChanged;
-
 
         /*********
         ** Public methods
@@ -28,54 +25,49 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         /// <summary>Construct an instance.</summary>
         /// <param name="fromType">The type whose methods to remap.</param>
         /// <param name="toType">The type with methods to map to.</param>
-        /// <param name="onlyIfPlatformChanged">Whether to only rewrite references if loading the assembly on a different platform than it was compiled on.</param>
         /// <param name="nounPhrase">A brief noun phrase indicating what the instruction finder matches (or <c>null</c> to generate one).</param>
-        public MethodParentRewriter(string fromType, Type toType, bool onlyIfPlatformChanged = false, string nounPhrase = null)
+        public MethodParentRewriter(string fromType, Type toType, string nounPhrase = null)
             : base(nounPhrase ?? $"{fromType.Split('.').Last()} methods")
         {
             this.FromType = fromType;
             this.ToType = toType;
-            this.OnlyIfPlatformChanged = onlyIfPlatformChanged;
         }
 
         /// <summary>Construct an instance.</summary>
         /// <param name="fromType">The type whose methods to remap.</param>
         /// <param name="toType">The type with methods to map to.</param>
-        /// <param name="onlyIfPlatformChanged">Whether to only rewrite references if loading the assembly on a different platform than it was compiled on.</param>
         /// <param name="nounPhrase">A brief noun phrase indicating what the instruction finder matches (or <c>null</c> to generate one).</param>
-        public MethodParentRewriter(Type fromType, Type toType, bool onlyIfPlatformChanged = false, string nounPhrase = null)
-            : this(fromType.FullName, toType, onlyIfPlatformChanged, nounPhrase) { }
-
+        public MethodParentRewriter(Type fromType, Type toType, string nounPhrase = null)
+            : this(fromType.FullName, toType, nounPhrase) { }
 
-        /// <summary>Perform the predefined logic for an instruction if applicable.</summary>
+        /// <summary>Rewrite a CIL instruction reference if needed.</summary>
         /// <param name="module">The assembly module containing the instruction.</param>
         /// <param name="cil">The CIL processor.</param>
         /// <param name="instruction">The CIL instruction to handle.</param>
-        /// <param name="assemblyMap">Metadata for mapping assemblies to the current platform.</param>
-        /// <param name="platformChanged">Whether the mod was compiled on a different platform.</param>
-        public override InstructionHandleResult Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, PlatformAssemblyMap assemblyMap, bool platformChanged)
+        /// <param name="replaceWith">Replaces the CIL instruction with a new one.</param>
+        /// <returns>Returns whether the instruction was changed.</returns>
+        public override bool Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, Action<Instruction> replaceWith)
         {
-            if (!this.IsMatch(instruction, platformChanged))
-                return InstructionHandleResult.None;
+            // get method ref
+            MethodReference methodRef = RewriteHelper.AsMethodReference(instruction);
+            if (!this.IsMatch(methodRef))
+                return false;
 
-            MethodReference methodRef = (MethodReference)instruction.Operand;
+            // rewrite
             methodRef.DeclaringType = module.ImportReference(this.ToType);
-            return InstructionHandleResult.Rewritten;
+            return this.MarkRewritten();
         }
 
 
         /*********
-        ** Protected methods
+        ** Private methods
         *********/
         /// <summary>Get whether a CIL instruction matches.</summary>
-        /// <param name="instruction">The IL instruction.</param>
-        /// <param name="platformChanged">Whether the mod was compiled on a different platform.</param>
-        protected bool IsMatch(Instruction instruction, bool platformChanged)
+        /// <param name="methodRef">The method reference.</param>
+        private bool IsMatch(MethodReference methodRef)
         {
-            MethodReference methodRef = RewriteHelper.AsMethodReference(instruction);
             return
                 methodRef != null
-                && (platformChanged || !this.OnlyIfPlatformChanged)
                 && methodRef.DeclaringType.FullName == this.FromType
                 && RewriteHelper.HasMatchingSignature(this.ToType, methodRef);
         }
diff --git a/src/SMAPI/Framework/ModLoading/Rewriters/StaticFieldToConstantRewriter.cs b/src/SMAPI/Framework/ModLoading/Rewriters/StaticFieldToConstantRewriter.cs
index 7e7c0efa..6ef18b26 100644
--- a/src/SMAPI/Framework/ModLoading/Rewriters/StaticFieldToConstantRewriter.cs
+++ b/src/SMAPI/Framework/ModLoading/Rewriters/StaticFieldToConstantRewriter.cs
@@ -1,17 +1,23 @@
 using System;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
-using StardewModdingAPI.Framework.ModLoading.Finders;
+using StardewModdingAPI.Framework.ModLoading.Framework;
 
 namespace StardewModdingAPI.Framework.ModLoading.Rewriters
 {
     /// <summary>Rewrites static field references into constant values.</summary>
     /// <typeparam name="TValue">The constant value type.</typeparam>
-    internal class StaticFieldToConstantRewriter<TValue> : FieldFinder
+    internal class StaticFieldToConstantRewriter<TValue> : BaseInstructionHandler
     {
         /*********
         ** Fields
         *********/
+        /// <summary>The type containing the field to which references should be rewritten.</summary>
+        private readonly Type Type;
+
+        /// <summary>The field name to which references should be rewritten.</summary>
+        private readonly string FromFieldName;
+
         /// <summary>The constant value to replace with.</summary>
         private readonly TValue Value;
 
@@ -24,24 +30,29 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         /// <param name="fieldName">The field name to rewrite.</param>
         /// <param name="value">The constant value to replace with.</param>
         public StaticFieldToConstantRewriter(Type type, string fieldName, TValue value)
-            : base(type.FullName, fieldName, InstructionHandleResult.None)
+            : base(defaultPhrase: $"{type.FullName}.{fieldName} field")
         {
+            this.Type = type;
+            this.FromFieldName = fieldName;
             this.Value = value;
         }
 
-        /// <summary>Perform the predefined logic for an instruction if applicable.</summary>
+        /// <summary>Rewrite a CIL instruction reference if needed.</summary>
         /// <param name="module">The assembly module containing the instruction.</param>
         /// <param name="cil">The CIL processor.</param>
-        /// <param name="instruction">The instruction to handle.</param>
-        /// <param name="assemblyMap">Metadata for mapping assemblies to the current platform.</param>
-        /// <param name="platformChanged">Whether the mod was compiled on a different platform.</param>
-        public override InstructionHandleResult Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, PlatformAssemblyMap assemblyMap, bool platformChanged)
+        /// <param name="instruction">The CIL instruction to handle.</param>
+        /// <param name="replaceWith">Replaces the CIL instruction with a new one.</param>
+        /// <returns>Returns whether the instruction was changed.</returns>
+        public override bool Handle(ModuleDefinition module, ILProcessor cil, Instruction instruction, Action<Instruction> replaceWith)
         {
-            if (!this.IsMatch(instruction))
-                return InstructionHandleResult.None;
+            // get field reference
+            FieldReference fieldRef = RewriteHelper.AsFieldReference(instruction);
+            if (!RewriteHelper.IsFieldReferenceTo(fieldRef, this.Type.FullName, this.FromFieldName))
+                return false;
 
-            cil.Replace(instruction, this.CreateConstantInstruction(cil, this.Value));
-            return InstructionHandleResult.Rewritten;
+            // rewrite to constant
+            replaceWith(this.CreateConstantInstruction(cil, this.Value));
+            return this.MarkRewritten();
         }
 
 
diff --git a/src/SMAPI/Framework/ModLoading/Rewriters/TypeReferenceRewriter.cs b/src/SMAPI/Framework/ModLoading/Rewriters/TypeReferenceRewriter.cs
index d95e5ac9..c2120444 100644
--- a/src/SMAPI/Framework/ModLoading/Rewriters/TypeReferenceRewriter.cs
+++ b/src/SMAPI/Framework/ModLoading/Rewriters/TypeReferenceRewriter.cs
@@ -1,12 +1,11 @@
 using System;
 using Mono.Cecil;
-using StardewModdingAPI.Framework.ModLoading.Finders;
 using StardewModdingAPI.Framework.ModLoading.Framework;
 
 namespace StardewModdingAPI.Framework.ModLoading.Rewriters
 {
     /// <summary>Rewrites all references to a type.</summary>
-    internal class TypeReferenceRewriter : BaseTypeReferenceRewriter
+    internal class TypeReferenceRewriter : BaseInstructionHandler
     {
         /*********
         ** Fields
@@ -17,6 +16,9 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         /// <summary>The new type to reference.</summary>
         private readonly Type ToType;
 
+        /// <summary>Get whether a matched type should be ignored.</summary>
+        private readonly Func<TypeReference, bool> ShouldIgnore;
+
 
         /*********
         ** Public methods
@@ -24,45 +26,29 @@ namespace StardewModdingAPI.Framework.ModLoading.Rewriters
         /// <summary>Construct an instance.</summary>
         /// <param name="fromTypeFullName">The full type name to which to find references.</param>
         /// <param name="toType">The new type to reference.</param>
-        /// <param name="shouldIgnore">A lambda which overrides a matched type.</param>
+        /// <param name="shouldIgnore">Get whether a matched type should be ignored.</param>
         public TypeReferenceRewriter(string fromTypeFullName, Type toType, Func<TypeReference, bool> shouldIgnore = null)
-            : base(new TypeFinder(fromTypeFullName, InstructionHandleResult.None, shouldIgnore), $"{fromTypeFullName} type")
+            : base($"{fromTypeFullName} type")
         {
             this.FromTypeName = fromTypeFullName;
             this.ToType = toType;
+            this.ShouldIgnore = shouldIgnore;
         }
 
-
-        /*********
-        ** Protected methods
-        *********/
-        /// <summary>Change a type reference if needed.</summary>
+        /// <summary>Rewrite a type reference if needed.</summary>
         /// <param name="module">The assembly module containing the instruction.</param>
-        /// <param name="type">The type to replace if it matches.</param>
-        /// <param name="set">Assign the new type reference.</param>
-        protected override bool RewriteIfNeeded(ModuleDefinition module, TypeReference type, Action<TypeReference> set)
+        /// <param name="type">The type definition to handle.</param>
+        /// <param name="replaceWith">Replaces the type reference with a new one.</param>
+        /// <returns>Returns whether the type was changed.</returns>
+        public override bool Handle(ModuleDefinition module, TypeReference type, Action<TypeReference> replaceWith)
         {
-            bool rewritten = false;
-
-            // current type
-            if (type.FullName == this.FromTypeName)
-            {
-                set(module.ImportReference(this.ToType));
-                return true;
-            }
-
-            // recurse into generic arguments
-            if (type is GenericInstanceType genericType)
-            {
-                for (int i = 0; i < genericType.GenericArguments.Count; i++)
-                    rewritten |= this.RewriteIfNeeded(module, genericType.GenericArguments[i], typeRef => genericType.GenericArguments[i] = typeRef);
-            }
-
-            // recurse into generic parameters (e.g. constraints)
-            for (int i = 0; i < type.GenericParameters.Count; i++)
-                rewritten |= this.RewriteIfNeeded(module, type.GenericParameters[i], typeRef => type.GenericParameters[i] = new GenericParameter(typeRef));
+            // check type reference
+            if (type.FullName != this.FromTypeName || this.ShouldIgnore?.Invoke(type) == true)
+                return false;
 
-            return rewritten;
+            // rewrite to new type
+            replaceWith(module.ImportReference(this.ToType));
+            return this.MarkRewritten();
         }
     }
 }