Improve split method dynfix

definition/src/main/java/org/sinytra/adapter/patch/analysis/MethodCallAnalyzer.java

@@ -50,6 +50,12 @@ public static Optional<MethodNode> findMethodByUniqueName(ClassNode cls, String
         return Optional.of(methods.getFirst());
     }
 
+    public static Optional<MethodNode> findMethodByNameOrThrow(ClassNode cls, String name, String desc) {
+        return cls.methods.stream()
+            .filter(m -> m.name.equals(name) && m.desc.equals(desc))
+            .findFirst();
+    }
+
     public static Multimap<String, MethodInsnNode> getMethodCalls(MethodNode node, List<String> callOrder) {
         ImmutableMultimap.Builder<String, MethodInsnNode> calls = ImmutableMultimap.builder();
         for (AbstractInsnNode insn : node.instructions) {

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

@@ -1,5 +1,6 @@
 package org.sinytra.adapter.patch.transformer.dynfix;
 
+import com.google.common.collect.Multimap;
 import org.jetbrains.annotations.Nullable;
 import org.objectweb.asm.tree.*;
 import org.sinytra.adapter.patch.analysis.MethodCallAnalyzer;
@@ -11,6 +12,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Handle cases where a single method is split into multiple smaller pieces.
@@ -33,13 +35,26 @@ 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);
+
+        if (candidates.size() == 1) {
+            MethodNode method = candidates.getFirst();
+            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);
+        }
+
+        return null;
+    }
+
+    private static List<MethodNode> locateCandidates(MethodContext methodContext) {
         MethodNode cleanTargetMethod = methodContext.findCleanInjectionTarget().methodNode();
         ClassNode dirtyTargetClass = methodContext.findDirtyInjectionTarget().classNode();
         MethodNode dirtyTargetMethod = methodContext.findDirtyInjectionTarget().methodNode();
 
         // Check that a Deprecated annotation was added to the dirty method 
         if (AdapterUtil.hasAnnotation(cleanTargetMethod.visibleAnnotations, DEPRECATED) || !AdapterUtil.hasAnnotation(dirtyTargetMethod.visibleAnnotations, DEPRECATED)) {
-            return null;
+            return tryFindPartialCandidates(cleanTargetMethod, dirtyTargetClass, dirtyTargetMethod, methodContext);
         }
 
         // Iterate over isns, leave out first and last elements
@@ -49,7 +64,7 @@ public FixResult apply(ClassNode classNode, MethodNode methodNode, MethodContext
         for (int i = 1; i < dirtyTargetMethod.instructions.size() - 1; i++) {
             AbstractInsnNode insn = dirtyTargetMethod.instructions.get(i);
             if (insn instanceof LabelNode) {
-                AbstractInsnNode previous = insn.getPrevious(); 
+                AbstractInsnNode previous = insn.getPrevious();
                 if (previous instanceof MethodInsnNode methodInsn && methodInsn.owner.equals(dirtyTargetClass.name)) {
                     MethodNode method = dirtyTargetClass.methods.stream().filter(m -> m.name.equals(methodInsn.name) && m.desc.equals(methodInsn.desc)).findFirst().orElseThrow();
                     invocations.add(method);
@@ -67,17 +82,24 @@ public FixResult apply(ClassNode classNode, MethodNode methodNode, MethodContext
                 .flatMap(m -> MethodCallAnalyzer.findLambdasInMethod(dirtyTargetClass, m, null).stream())
                 .flatMap(s -> MethodCallAnalyzer.findMethodByUniqueName(dirtyTargetClass, s).stream())
                 .toList();
-            candidates = findInsnsCalls(nestedLambdas, methodContext);
+            return findInsnsCalls(nestedLambdas, methodContext);
         }
 
-        if (candidates.size() == 1) {
-            MethodNode method = candidates.getFirst();
-            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);
-        }
+        return candidates;
+    }
 
-        return null;
+    // Handle cases where only part of the method is moved away
+    private static List<MethodNode> 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<>());
+
+        List<MethodNode> dirtyOnlyCalls = dirtyMethodCalls.entries().stream()
+            .filter(e -> !cleanMethodCalls.containsKey(e.getKey()) && e.getValue().owner.equals(dirtyTargetClass.name))
+            .map(Map.Entry::getValue)
+            .flatMap(i -> MethodCallAnalyzer.findMethodByNameOrThrow(dirtyTargetClass, i.name, i.desc).stream())
+            .toList();
+
+        return findInsnsCalls(dirtyOnlyCalls, methodContext);
     }
 
     private static List<MethodNode> findInsnsCalls(List<MethodNode> methods, MethodContext methodContext) {

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

@@ -112,6 +112,16 @@ void testMovedInjectionPointToMethod4() throws Exception {
         );
     }
 
+    @Test
+    void testMovedInjectionPointToMethod5() throws Exception {
+        assertSameCode(
+            "org/sinytra/adapter/test/mixin/LevelMixin",
+            "modifyLeastStatus",
+            assertTargetMethod(),
+            assertInjectionPoint()
+        );
+    }
+
     @Test
     void testUpdatedInjectionPointFieldToMethod() throws Exception {
         assertSameCode(

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

@@ -0,0 +1,33 @@
+package org.sinytra.adapter.test.mixin;
+
+import net.minecraft.server.level.FullChunkStatus;
+import net.minecraft.world.level.Level;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.ModifyArg;
+
+@Mixin(Level.class)
+public class LevelMixin {
+    // https://github.com/RelativityMC/C2ME-fabric/blob/93ffdce7b92051ab9d2d9f6ee7dd74b31a018533/c2me-notickvd/src/main/java/com/ishland/c2me/notickvd/mixin/MixinWorld.java
+    @ModifyArg(
+        method = "setBlock(Lnet/minecraft/core/BlockPos;Lnet/minecraft/world/level/block/state/BlockState;II)Z",
+        at = @At(
+            value = "INVOKE",
+            target = "Lnet/minecraft/server/level/FullChunkStatus;isOrAfter(Lnet/minecraft/server/level/FullChunkStatus;)Z"
+        )
+    )
+    private FullChunkStatus modifyLeastStatus(FullChunkStatus levelType) {
+        return levelType.ordinal() > FullChunkStatus.FULL.ordinal() ? FullChunkStatus.FULL : levelType;
+    }
+
+    @ModifyArg(
+        method = "markAndNotifyBlock(Lnet/minecraft/core/BlockPos;Lnet/minecraft/world/level/chunk/LevelChunk;Lnet/minecraft/world/level/block/state/BlockState;Lnet/minecraft/world/level/block/state/BlockState;II)V",
+        at = @At(
+            value = "INVOKE",
+            target = "Lnet/minecraft/server/level/FullChunkStatus;isOrAfter(Lnet/minecraft/server/level/FullChunkStatus;)Z"
+        )
+    )
+    private FullChunkStatus modifyLeastStatusExpected(FullChunkStatus levelType) {
+        return levelType.ordinal() > FullChunkStatus.FULL.ordinal() ? FullChunkStatus.FULL : levelType;
+    }
+}