Improve mixin extraction with generated injectors

definition/src/main/java/org/sinytra/adapter/patch/MethodContextImpl.java

@@ -172,7 +172,8 @@ public List<LocalVariable> getTargetMethodLocals(TargetPair target, int startPos
         return AdapterUtil.summariseLocals(locals, startPos);
     }
 
-    private List<AbstractInsnNode> computeInjectionTargetInsns(@Nullable TargetPair target) {
+    @Nullable
+    public List<AbstractInsnNode> computeInjectionTargetInsns(@Nullable TargetPair target) {
         if (target == null) {
             return List.of();
         }

definition/src/main/java/org/sinytra/adapter/patch/PatchInstance.java

@@ -9,8 +9,8 @@
 import org.sinytra.adapter.patch.api.*;
 import org.sinytra.adapter.patch.selector.AnnotationHandle;
 import org.sinytra.adapter.patch.selector.AnnotationValueHandle;
-import org.sinytra.adapter.patch.transformer.*;
-import org.sinytra.adapter.patch.transformer.param.TransformParameters;
+import org.sinytra.adapter.patch.transformer.ModifyTargetClasses;
+import org.sinytra.adapter.patch.util.MethodTransformBuilderImpl;
 import org.slf4j.Marker;
 import org.slf4j.MarkerFactory;
 
@@ -143,12 +143,11 @@ public static <T> Optional<AnnotationValueHandle<T>> findAnnotationValue(@Nullab
 
     private record ClassTarget(@Nullable AnnotationValueHandle<?> handle, List<Type> targetTypes) {}
 
-    protected abstract static class BaseBuilder<T extends Builder<T>> implements Builder<T> {
+    protected abstract static class BaseBuilder<T extends Builder<T>> extends MethodTransformBuilderImpl<T> implements Builder<T> {
         protected final Set<String> targetClasses = new HashSet<>();
         protected final Set<String> targetAnnotations = new HashSet<>();
         protected Predicate<AnnotationHandle> targetAnnotationValues;
         protected final List<ClassTransform> classTransforms = new ArrayList<>();
-        protected final List<MethodTransform> transforms = new ArrayList<>();
 
         @Override
         public T targetClass(String... targets) {
@@ -173,60 +172,6 @@ public T modifyTargetClasses(Consumer<List<Type>> consumer) {
             return transform(new ModifyTargetClasses(consumer));
         }
 
-        @Override
-        public T modifyParams(Consumer<ModifyMethodParams.Builder> consumer) {
-            ModifyMethodParams.Builder builder = ModifyMethodParams.builder();
-            consumer.accept(builder);
-            return transform(builder.build());
-        }
-
-        @Override
-        public T transformParams(Consumer<TransformParameters.Builder> consumer) {
-            final var builder = new TransformParameters.Builder();
-            consumer.accept(builder);
-            return transform(builder.build());
-        }
-
-        @Override
-        public T modifyTarget(String... methods) {
-            return transform(new ModifyInjectionTarget(List.of(methods)));
-        }
-
-        @Override
-        public T modifyTarget(ModifyInjectionTarget.Action action, String... methods) {
-            return transform(new ModifyInjectionTarget(List.of(methods), action));
-        }
-
-        @Override
-        public T modifyVariableIndex(int start, int offset) {
-            return transform(new ChangeModifiedVariableIndex(start, offset));
-        }
-
-        @Override
-        public T modifyMethodAccess(ModifyMethodAccess.AccessChange... changes) {
-            return transform(new ModifyMethodAccess(List.of(changes)));
-        }
-
-        @Override
-        public T extractMixin(String targetClass) {
-            return transform(ModifyVarUpgradeToModifyExprVal.INSTANCE).transform(new ExtractMixin(targetClass));
-        }
-
-        @Override
-        public T splitMixin(String targetClass) {
-            return transform(new SplitMixinTransform(targetClass));
-        }
-
-        @Override
-        public T improveModifyVar() {
-            return transform(ModifyVarUpgradeToModifyExprVal.INSTANCE);
-        }
-
-        @Override
-        public T modifyMixinType(String newType, Consumer<ModifyMixinType.Builder> consumer) {
-            return transform(new ModifyMixinType(newType, consumer));
-        }
-
         @Override
         public T transform(List<ClassTransform> classTransforms) {
             this.classTransforms.addAll(classTransforms);
@@ -239,24 +184,6 @@ public T transform(ClassTransform transformer) {
             return coerce();
         }
 
-        @Override
-        public T transform(MethodTransform transformer) {
-            this.transforms.add(transformer);
-            return coerce();
-        }
-
-        @Override
-        public T transformMethods(List<MethodTransform> transformers) {
-            transformers.forEach(this::transform);
-            return coerce();
-        }
-
-        @Override
-        public T chain(Consumer<T> consumer) {
-            consumer.accept(coerce());
-            return coerce();
-        }
-
         @SuppressWarnings("unchecked")
         private T coerce() {
             return (T) this;

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

@@ -92,6 +92,20 @@ public static boolean test(InsnList first, InsnList second, int flags) {
         return false;
     }
 
+    public static boolean test(List<AbstractInsnNode> first, List<AbstractInsnNode> second, int flags) {
+        if (first.size() == second.size()) {
+            for (int i = 0; i < first.size(); i++) {
+                AbstractInsnNode insn = first.get(i);
+                AbstractInsnNode otherInsn = second.get(i);
+                if (!InsnComparator.instructionsEqual(insn, otherInsn, flags)) {
+                    return false;
+                }
+            }
+            return true;
+        }
+        return false;
+    }
+
     public static <T> int count(List<T> list, T item) {
         return (int) list.stream().filter(item::equals).count();
     }

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

@@ -125,6 +125,12 @@ public static List<List<AbstractInsnNode>> getInvocationInsns(MethodNode methodN
         });
     }
 
+    @Nullable
+    public static List<AbstractInsnNode> findMethodCallParamInsns(MethodNode methodNode, MethodInsnNode insn) {
+        MethodCallInterpreter interpreter = MethodCallAnalyzer.analyzeInterpretMethod(methodNode, new MethodCallInterpreter(insn));
+        return interpreter.getTargetArgs();
+    }
+
     public static <T> List<T> analyzeMethod(MethodNode methodNode, NaryOperationHandler<T> handler) {
         return analyzeMethod(methodNode, (insn, values) -> true, handler);
     }
@@ -171,6 +177,34 @@ public static MethodNode findUniqueMethod(Multimap<String, MethodNode> methods,
         throw new NullPointerException("Method " + name + " not found");
     }
 
+    private static class MethodCallInterpreter extends SourceInterpreter {
+        private final MethodInsnNode targetInsn;
+        private List<AbstractInsnNode> targetArgs;
+
+        public MethodCallInterpreter(MethodInsnNode targetInsn) {
+            super(Opcodes.ASM9);
+            this.targetInsn = targetInsn;
+        }
+
+        @javax.annotation.Nullable
+        public List<AbstractInsnNode> getTargetArgs() {
+            return this.targetArgs;
+        }
+
+        @Override
+        public SourceValue naryOperation(AbstractInsnNode insn, List<? extends SourceValue> values) {
+            if (insn == this.targetInsn && this.targetArgs == null) {
+                List<AbstractInsnNode> targetArgs = values.stream()
+                    .map(v -> v.insns.size() == 1 ? v.insns.iterator().next() : null)
+                    .toList();
+                if (!targetArgs.contains(null)) {
+                    this.targetArgs = targetArgs;
+                }
+            }
+            return super.naryOperation(insn, values);
+        }
+    }
+
     private static class AnalysingSourceInterpreter<T> extends SourceInterpreter {
         private final BiPredicate<MethodInsnNode, List<? extends SourceValue>> filter;
         private final NaryOperationHandler<T> handler;

definition/src/main/java/org/sinytra/adapter/patch/api/MethodContext.java

@@ -49,6 +49,11 @@ public interface MethodContext {
 
     List<AbstractInsnNode> findInjectionTargetInsns(@Nullable TargetPair target);
 
+    /**
+     * Uncached variant of {@link #findInjectionTargetInsns(TargetPair)}
+     */
+    List<AbstractInsnNode> computeInjectionTargetInsns(@Nullable TargetPair target);
+
     @Nullable
     Pair<ClassNode, List<MethodNode>> findInjectionTargetCandidates(ClassLookup lookup);
 

definition/src/main/java/org/sinytra/adapter/patch/api/MethodTransform.java

@@ -16,8 +16,8 @@ default Collection<String> getAcceptedAnnotations() {
         return Set.of();
     }
 
-    default Patch.Result apply(ClassNode classNode, MethodNode methodNode, MethodContext methodContext) {
-        return apply(classNode, methodNode, methodContext, methodContext.patchContext());
+    default Patch.Result apply(MethodContext methodContext) {
+        return apply(methodContext.getMixinClass(), methodContext.getMixinMethod(), methodContext, methodContext.patchContext());
     }
 
     Patch.Result apply(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, PatchContext context);

definition/src/main/java/org/sinytra/adapter/patch/api/MethodTransformBuilder.java

@@ -0,0 +1,41 @@
+package org.sinytra.adapter.patch.api;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.sinytra.adapter.patch.transformer.ModifyInjectionTarget;
+import org.sinytra.adapter.patch.transformer.ModifyMethodAccess;
+import org.sinytra.adapter.patch.transformer.ModifyMethodParams;
+import org.sinytra.adapter.patch.transformer.ModifyMixinType;
+import org.sinytra.adapter.patch.transformer.param.TransformParameters;
+
+import java.util.List;
+import java.util.function.Consumer;
+
+public interface MethodTransformBuilder<T extends MethodTransformBuilder<T>> {
+    @Deprecated
+    T modifyParams(Consumer<ModifyMethodParams.Builder> consumer);
+
+    @ApiStatus.Experimental
+    T transformParams(Consumer<TransformParameters.Builder> consumer);
+
+    T modifyTarget(String... methods);
+
+    T modifyTarget(ModifyInjectionTarget.Action action, String... methods);
+
+    T modifyVariableIndex(int start, int offset);
+
+    T modifyMethodAccess(ModifyMethodAccess.AccessChange... changes);
+
+    T extractMixin(String targetClass);
+
+    T splitMixin(String targetClass);
+
+    T improveModifyVar();
+
+    T modifyMixinType(String newType, Consumer<ModifyMixinType.Builder> consumer);
+
+    T transform(MethodTransform transformer);
+
+    T transformMethods(List<MethodTransform> transformers);
+
+    T chain(Consumer<T> consumer);
+}

definition/src/main/java/org/sinytra/adapter/patch/api/Patch.java

@@ -1,7 +1,6 @@
 package org.sinytra.adapter.patch.api;
 
 import com.mojang.serialization.Codec;
-import org.jetbrains.annotations.ApiStatus;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.commons.InstructionAdapter;
 import org.objectweb.asm.tree.ClassNode;
@@ -12,11 +11,6 @@
 import org.sinytra.adapter.patch.PatchInstance;
 import org.sinytra.adapter.patch.selector.AnnotationHandle;
 import org.sinytra.adapter.patch.selector.AnnotationValueHandle;
-import org.sinytra.adapter.patch.transformer.ModifyInjectionTarget;
-import org.sinytra.adapter.patch.transformer.ModifyMethodAccess;
-import org.sinytra.adapter.patch.transformer.ModifyMethodParams;
-import org.sinytra.adapter.patch.transformer.ModifyMixinType;
-import org.sinytra.adapter.patch.transformer.param.TransformParameters;
 
 import java.util.List;
 import java.util.function.BiConsumer;
@@ -52,7 +46,7 @@ public Result or(Result other) {
         }
     }
 
-    interface Builder<T extends Builder<T>> {
+    interface Builder<T extends Builder<T>> extends MethodTransformBuilder<T> {
         T targetClass(String... targets);
 
         T targetMixinType(String... annotationDescs);
@@ -61,38 +55,10 @@ interface Builder<T extends Builder<T>> {
 
         T modifyTargetClasses(Consumer<List<Type>> consumer);
 
-        @Deprecated
-        T modifyParams(Consumer<ModifyMethodParams.Builder> consumer);
-
-        @ApiStatus.Experimental
-        T transformParams(Consumer<TransformParameters.Builder> consumer);
-
-        T modifyTarget(String... methods);
-
-        T modifyTarget(ModifyInjectionTarget.Action action, String... methods);
-
-        T modifyVariableIndex(int start, int offset);
-
-        T modifyMethodAccess(ModifyMethodAccess.AccessChange... changes);
-
-        T extractMixin(String targetClass);
-
-        T splitMixin(String targetClass);
-
-        T improveModifyVar();
-
-        T modifyMixinType(String newType, Consumer<ModifyMixinType.Builder> consumer);
-
         T transform(List<ClassTransform> classTransforms);
 
         T transform(ClassTransform transformer);
 
-        T transform(MethodTransform transformer);
-
-        T transformMethods(List<MethodTransform> transformers);
-
-        T chain(Consumer<T> consumer);
-
         PatchInstance build();
     }
 
@@ -112,7 +78,7 @@ default ClassPatchBuilder targetConstant(double doubleValue) {
                 .orElseGet(() -> handle.getNested("at")
                     .flatMap(at -> at.<String>getValue("value").map(s -> s.get().equals("CONSTANT") &&
                         at.<List<String>>getValue("args").map(AnnotationValueHandle::get).map(t -> t.size() == 1
-                            && (t.getFirst().equals("doubleValue=" + doubleValue + "D") || t.getFirst().equals("doubleValue=" + doubleValue)))
+                                && (t.getFirst().equals("doubleValue=" + doubleValue + "D") || t.getFirst().equals("doubleValue=" + doubleValue)))
                             .orElse(false)))
                     .orElse(false)));
         }

definition/src/main/java/org/sinytra/adapter/patch/fixes/MethodUpgrader.java

@@ -2,7 +2,6 @@
 
 import com.google.common.collect.ImmutableList;
 import org.objectweb.asm.Type;
-import org.objectweb.asm.tree.ClassNode;
 import org.objectweb.asm.tree.MethodNode;
 import org.sinytra.adapter.patch.analysis.LocalVarAnalyzer;
 import org.sinytra.adapter.patch.analysis.params.EnhancedParamsDiff;
@@ -22,7 +21,7 @@
 
 public final class MethodUpgrader {
 
-    public static void upgradeMethod(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, String originalDesc, String modifiedDesc) {
+    public static void upgradeMethod(MethodNode methodNode, MethodContext methodContext, String originalDesc, String modifiedDesc) {
         MethodQualifier cleanQualifier = MethodQualifier.create(originalDesc).orElse(null);
         if (cleanQualifier == null) {
             return;
@@ -34,11 +33,11 @@ public static void upgradeMethod(ClassNode classNode, MethodNode methodNode, Met
         if (methodContext.methodAnnotation().matchesDesc(MixinConstants.MODIFY_ARGS)) {
             ModifyArgsOffsetTransformer.handleModifiedDesc(methodNode, cleanQualifier.desc(), dirtyQualifier.desc());
         } else if (methodContext.methodAnnotation().matchesDesc(MixinConstants.WRAP_OPERATION)) {
-            upgradeWrapOperation(classNode, methodNode, methodContext, cleanQualifier, dirtyQualifier);
+            upgradeWrapOperation(methodNode, methodContext, cleanQualifier, dirtyQualifier);
         }
     }
 
-    public static void upgradeCapturedLocals(ClassNode classNode, MethodNode methodNode, MethodContext methodContext) {
+    public static void upgradeCapturedLocals(MethodNode methodNode, MethodContext methodContext) {
         AdapterUtil.CapturedLocals capturedLocals = AdapterUtil.getCapturedLocals(methodNode, methodContext);
         if (capturedLocals == null) {
             return;
@@ -51,7 +50,7 @@ public static void upgradeCapturedLocals(ClassNode classNode, MethodNode methodN
         }
 
         LocalVarAnalyzer.CapturedLocalsTransform transform = LocalVarAnalyzer.analyzeCapturedLocals(capturedLocals, methodNode);
-        transform.remover().apply(classNode, methodNode, methodContext);
+        transform.remover().apply(methodContext);
 
         List<Type> expected = List.of(Type.getArgumentTypes(methodNode.desc));
         List<Type> required = ImmutableList.<Type>builder()
@@ -64,11 +63,11 @@ public static void upgradeCapturedLocals(ClassNode classNode, MethodNode methodN
                 .map(LayeredParamsDiffSnapshot.ParamModification::asParameterTransformer)
                 .toList();
             MethodTransform patch = TransformParameters.builder().transform(transformers).withOffset().targetType(ParamTransformTarget.METHOD).build();
-            patch.apply(classNode, methodNode, methodContext);
+            patch.apply(methodContext);
         }
     }
 
-    private static void upgradeWrapOperation(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, MethodQualifier cleanQualifier, MethodQualifier dirtyQualifier) {
+    private static void upgradeWrapOperation(MethodNode methodNode, MethodContext methodContext, MethodQualifier cleanQualifier, MethodQualifier dirtyQualifier) {
         if (dirtyQualifier.owner() == null || cleanQualifier.desc() == null) {
             return;
         }
@@ -87,7 +86,7 @@ private static void upgradeWrapOperation(ClassNode classNode, MethodNode methodN
         SimpleParamsDiffSnapshot diff = EnhancedParamsDiff.create(originalDesc, modifiedDesc);
         if (!diff.isEmpty()) {
             MethodTransform patch = diff.asParameterTransformer(ParamTransformTarget.ALL, false, false);
-            patch.apply(classNode, methodNode, methodContext);
+            patch.apply(methodContext);
         }
     }
 }