Disambiguate candidates in DynFixSplitMethod

definition/src/main/java/org/sinytra/adapter/patch/transformer/dynfix/DynFixSplitMethod.java

@@ -3,6 +3,8 @@
 import com.google.common.collect.Multimap;
 import org.jetbrains.annotations.Nullable;
 import org.objectweb.asm.tree.*;
+import org.sinytra.adapter.patch.analysis.InsnComparator;
+import org.sinytra.adapter.patch.analysis.InstructionMatcher;
 import org.sinytra.adapter.patch.analysis.MethodCallAnalyzer;
 import org.sinytra.adapter.patch.api.MethodContext;
 import org.sinytra.adapter.patch.api.PatchAuditTrail;
@@ -13,6 +15,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Handle cases where a single method is split into multiple smaller pieces.
@@ -21,7 +24,8 @@
 public class DynFixSplitMethod implements DynamicFixer<DynFixSplitMethod.Data> {
     private static final String DEPRECATED = "Ljava/lang/Deprecated;";
 
-    public record Data() {}
+    public record Data() {
+    }
 
     @Nullable
     @Override
@@ -35,10 +39,10 @@ public DynFixSplitMethod.Data prepare(MethodContext methodContext) {
     @Override
     @Nullable
     public FixResult apply(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, PatchAuditTrail auditTrail, Data data) {
-        List<MethodNode> candidates = locateCandidates(methodContext);
+        List<CandidateMethod> candidates = disambiguate(locateCandidates(methodContext), methodContext);
 
         if (candidates.size() == 1) {
-            MethodNode method = candidates.getFirst();
+            MethodNode method = candidates.getFirst().method();
             String newTarget = method.name + method.desc;
             methodContext.recordAudit(this, "Adjusting split method target to %s", newTarget);
             return FixResult.of(new ModifyInjectionTarget(List.of(newTarget)).apply(methodContext), PatchAuditTrail.Match.FULL);
@@ -47,7 +51,7 @@ public FixResult apply(ClassNode classNode, MethodNode methodNode, MethodContext
         return null;
     }
 
-    private static List<MethodNode> locateCandidates(MethodContext methodContext) {
+    private static List<CandidateMethod> locateCandidates(MethodContext methodContext) {
         MethodNode cleanTargetMethod = methodContext.findCleanInjectionTarget().methodNode();
         ClassNode dirtyTargetClass = methodContext.findDirtyInjectionTarget().classNode();
         MethodNode dirtyTargetMethod = methodContext.findDirtyInjectionTarget().methodNode();
@@ -74,7 +78,7 @@ private static List<MethodNode> locateCandidates(MethodContext methodContext) {
             }
         }
 
-        List<MethodNode> candidates = findInsnsCalls(invocations, methodContext);
+        List<CandidateMethod> candidates = findInsnsCalls(invocations, methodContext);
 
         // Attempt to find matching insns in lambdas
         if (candidates.isEmpty()) {
@@ -89,7 +93,7 @@ private static List<MethodNode> locateCandidates(MethodContext methodContext) {
     }
 
     // Handle cases where only part of the method is moved away
-    private static List<MethodNode> tryFindPartialCandidates(MethodNode cleanTargetMethod, ClassNode dirtyTargetClass, MethodNode dirtyTargetMethod, MethodContext methodContext) {
+    private static List<CandidateMethod> tryFindPartialCandidates(MethodNode cleanTargetMethod, ClassNode dirtyTargetClass, MethodNode dirtyTargetMethod, MethodContext methodContext) {
         Multimap<String, MethodInsnNode> cleanMethodCalls = MethodCallAnalyzer.getMethodCalls(cleanTargetMethod, new ArrayList<>());
         Multimap<String, MethodInsnNode> dirtyMethodCalls = MethodCallAnalyzer.getMethodCalls(dirtyTargetMethod, new ArrayList<>());
 
@@ -102,8 +106,44 @@ private static List<MethodNode> tryFindPartialCandidates(MethodNode cleanTargetM
         return findInsnsCalls(dirtyOnlyCalls, methodContext);
     }
 
-    private static List<MethodNode> findInsnsCalls(List<MethodNode> methods, MethodContext methodContext) {
+    // If multiple candidates have been found, try comparing the surrounding method instructions to find one match
+    private static List<CandidateMethod> disambiguate(List<CandidateMethod> candidates, MethodContext methodContext) {
+        if (candidates.size() <= 1) {
+            return candidates;
+        }
+
+        List<AbstractInsnNode> cleanInsns = methodContext.findInjectionTargetInsns(methodContext.findCleanInjectionTarget());
+        if (cleanInsns.size() != 1) {
+            return candidates;
+        }
+
+        InstructionMatcher cleanMatcher = MethodCallAnalyzer.findSurroundingInstructions(cleanInsns.getFirst(), 5);
+
+        List<CandidateMethod> matchingCandidates = candidates.stream()
+            .filter(method -> method.insns().size() == 1)
+            .filter(method -> {
+                InstructionMatcher matcher = MethodCallAnalyzer.findSurroundingInstructions(method.insns().getFirst(), 5);
+                return cleanMatcher.test(matcher, InsnComparator.IGNORE_VAR_INDEX);
+            })
+            .toList();
+
+        if (matchingCandidates.size() == 1) {
+            return matchingCandidates;
+        }
+
+        return candidates;
+    }
+
+    private static List<CandidateMethod> findInsnsCalls(List<MethodNode> methods, MethodContext methodContext) {
         ClassNode dirtyTargetClass = methodContext.findDirtyInjectionTarget().classNode();
-        return methods.stream().filter(method -> !methodContext.findInjectionTargetInsns(new MethodContext.TargetPair(dirtyTargetClass, method)).isEmpty()).toList();
+        return methods.stream()
+            .map(method -> {
+                List<AbstractInsnNode> insns = methodContext.findInjectionTargetInsns(new MethodContext.TargetPair(dirtyTargetClass, method));
+                return !insns.isEmpty() ? new CandidateMethod(method, insns) : null;
+            })
+            .filter(Objects::nonNull)
+            .toList();
     }
+
+    private record CandidateMethod(MethodNode method, List<AbstractInsnNode> insns) {}
 }

test/src/test/java/org/sinytra/adapter/patch/test/mixin/DynamicMixinPatchTest.java

@@ -178,6 +178,12 @@ void testSplitMethodInjectionTarget() throws Exception {
             assertTargetMethod(),
             assertInjectionPoint()
         );
+        assertSameCode(
+            "org/sinytra/adapter/test/mixin/GuiMixin",
+            "moveHealthDown",
+            assertTargetMethod(),
+            assertInjectionPoint()
+        );
     }
 
     @Test

test/src/testClasses/java/org/sinytra/adapter/test/mixin/GuiMixin.java

@@ -1,5 +1,6 @@
 package org.sinytra.adapter.test.mixin;
 
+import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
 import com.llamalad7.mixinextras.sugar.Local;
 import net.minecraft.client.DeltaTracker;
 import net.minecraft.client.gui.Gui;
@@ -80,4 +81,21 @@ private void cozyBackgroundExpected(GuiGraphics context, DeltaTracker tickCounte
         // Use only a single captured local
         boolean something = registryEntry == null;
     }
+
+    // https://github.com/SkyblockerMod/Skyblocker/blob/fa67224da0cdcd2325b24f41732e820ff7ef04e8/src/main/java/de/hysky/skyblocker/mixins/InGameHudMixin.java#L97
+    @ModifyExpressionValue(
+        method = "renderPlayerHealth(Lnet/minecraft/client/gui/GuiGraphics;)V",
+        at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/GuiGraphics;guiHeight()I")
+    )
+    private int moveHealthDown(int original) {
+        return original;
+    }
+
+    @ModifyExpressionValue(
+        method = "renderHealthLevel(Lnet/minecraft/client/gui/GuiGraphics;)V",
+        at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/GuiGraphics;guiHeight()I")
+    )
+    private int moveHealthDownExpected(int original) {
+        return original;
+    }
 }