Another wrap operation case

Sinytra · Aug 13, 2024 · aade659 · aade659
1 parent 02ccc2d
commit aade659
Show file tree

Hide file tree

Showing 15 changed files with 331 additions and 54 deletions.
diff --git a/...on/src/main/java/org/sinytra/adapter/patch/analysis/params/LayeredParamsDiffSnapshot.java b/...on/src/main/java/org/sinytra/adapter/patch/analysis/params/LayeredParamsDiffSnapshot.java
@@ -11,6 +11,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.function.Consumer;
 
 public record LayeredParamsDiffSnapshot(List<ParamModification> modifications) implements ParamsDiffSnapshot {
@@ -21,7 +22,7 @@ public interface ParamModification {
 
         boolean satisfiesIndexLimit(int index);
 
-        ParameterTransformer asParameterTransformer();
+        ParameterTransformer asParameterTransformer(Set<Flags> flags);
     }
 
     public record InsertParam(int index, Type type) implements ParamModification {
@@ -36,7 +37,7 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
             return new InjectParameterTransform(this.index, this.type);
         }
     }
@@ -53,8 +54,8 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
-            return new ReplaceParametersTransformer(this.index, this.type);
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
+            return new ReplaceParametersTransformer(this.index, this.type, flags.contains(Flags.UPGRADE_WRAP_OP));
         }
     }
 
@@ -70,7 +71,7 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
             return new SwapParametersTransformer(this.from, this.to);
         }
     }
@@ -87,7 +88,7 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
             return new MoveParametersTransformer(this.from, this.to);
         }
     }
@@ -104,7 +105,7 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
             return new RemoveParameterTransformer(this.index);
         }
     }
@@ -121,7 +122,7 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
             return new InlineParameterTransformer(this.target, this.adapter);
         }
     }
@@ -138,7 +139,7 @@ public boolean satisfiesIndexLimit(int index) {
         }
 
         @Override
-        public ParameterTransformer asParameterTransformer() {
+        public ParameterTransformer asParameterTransformer(Set<Flags> flags) {
             return new SubstituteParameterTransformer(this.target, this.substitute);
         }
     }
@@ -196,8 +197,8 @@ public LayeredParamsDiffSnapshot offset(int offset, int limit) {
     }
 
     @Override
-    public MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset, boolean upgradeWrapOperation) {
-        List<ParameterTransformer> transformers = this.modifications.stream().map(ParamModification::asParameterTransformer).toList();
+    public MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset, Set<Flags> flags) {
+        List<ParameterTransformer> transformers = this.modifications.stream().map(paramModification -> paramModification.asParameterTransformer(flags)).toList();
         return TransformParameters.builder().transform(transformers).withOffset(withOffset).targetType(type).build();
     }
 

diff --git a/definition/src/main/java/org/sinytra/adapter/patch/analysis/params/ParamsDiffSnapshot.java b/definition/src/main/java/org/sinytra/adapter/patch/analysis/params/ParamsDiffSnapshot.java
@@ -5,9 +5,15 @@
 import org.sinytra.adapter.patch.api.MethodTransform;
 import org.sinytra.adapter.patch.transformer.operation.param.ParamTransformTarget;
 
+import java.util.EnumSet;
 import java.util.List;
+import java.util.Set;
 
 public interface ParamsDiffSnapshot {
+    enum Flags {
+        UPGRADE_WRAP_OP;
+    } 
+
     boolean isEmpty();
 
     List<Pair<Integer, Type>> insertions();
@@ -19,8 +25,8 @@ public interface ParamsDiffSnapshot {
     ParamsDiffSnapshot offset(int offset, int limit);
 
     default MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset) {
-        return asParameterTransformer(type, withOffset, true);
+        return asParameterTransformer(type, withOffset, EnumSet.of(Flags.UPGRADE_WRAP_OP));
     }
 
-    MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset, boolean upgradeWrapOperation);
+    MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset, Set<Flags> flags);
 }
diff --git a/...ion/src/main/java/org/sinytra/adapter/patch/analysis/params/SimpleParamsDiffSnapshot.java b/...ion/src/main/java/org/sinytra/adapter/patch/analysis/params/SimpleParamsDiffSnapshot.java
@@ -21,6 +21,7 @@
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.function.UnaryOperator;
@@ -88,7 +89,7 @@ public SimpleParamsDiffSnapshot offset(int offset, int limit) {
     }
 
     @Override
-    public MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset, boolean upgradeWrapOperation) {
+    public MethodTransform asParameterTransformer(ParamTransformTarget type, boolean withOffset, Set<Flags> flags) {
         List<MethodTransform> list = new ArrayList<>();
         SimpleParamsDiffSnapshot light = new SimpleParamsDiffSnapshot(List.of(), this.replacements, this.swaps, this.substitutes, this.removals, this.moves, this.inlines);
         if (!light.isEmpty()) {
@@ -97,7 +98,7 @@ public MethodTransform asParameterTransformer(ParamTransformTarget type, boolean
         if (!this.insertions.isEmpty()) {
             list.add(TransformParameters.builder()
                 .transform(this.insertions.stream()
-                    .<ParameterTransformer>map(p -> new InjectParameterTransform(p.getFirst(), p.getSecond(), upgradeWrapOperation))
+                    .<ParameterTransformer>map(p -> new InjectParameterTransform(p.getFirst(), p.getSecond(), flags.contains(Flags.UPGRADE_WRAP_OP)))
                     .toList())
                 .withOffset(withOffset)
                 .targetType(type)

diff --git a/definition/src/main/java/org/sinytra/adapter/patch/fixes/MethodUpgrader.java b/definition/src/main/java/org/sinytra/adapter/patch/fixes/MethodUpgrader.java
@@ -24,6 +24,7 @@
 
 import java.util.Comparator;
 import java.util.List;
+import java.util.Set;
 
 public final class MethodUpgrader {
 
@@ -75,7 +76,7 @@ private static void upgradeModifyExpValue(MethodNode methodNode, MethodContext m
         // Create diff
         SimpleParamsDiffSnapshot diff = EnhancedParamsDiff.create(originalDesc, modifiedDesc);
         if (!diff.isEmpty()) {
-            MethodTransform patch = diff.asParameterTransformer(ParamTransformTarget.ALL, false, false);
+            MethodTransform patch = diff.asParameterTransformer(ParamTransformTarget.ALL, false, Set.of());
             patch.apply(methodContext);
         }
     }
@@ -130,7 +131,7 @@ public static void upgradeCapturedLocals(MethodNode methodNode, MethodContext me
         LayeredParamsDiffSnapshot diff = EnhancedParamsDiff.createLayered(expected, required);
         if (!diff.isEmpty()) {
             List<ParameterTransformer> transformers = diff.modifications().stream()
-                .map(LayeredParamsDiffSnapshot.ParamModification::asParameterTransformer)
+                .map(paramModification -> paramModification.asParameterTransformer(Set.of()))
                 .toList();
             MethodTransform patch = TransformParameters.builder().transform(transformers).withOffset().targetType(ParamTransformTarget.METHOD).build();
             patch.apply(methodContext);
@@ -155,7 +156,31 @@ private static void upgradeWrapOperation(MethodNode methodNode, MethodContext me
         // Create diff
         SimpleParamsDiffSnapshot diff = EnhancedParamsDiff.create(originalDesc, modifiedDesc);
         if (!diff.isEmpty()) {
-            MethodTransform patch = diff.asParameterTransformer(ParamTransformTarget.ALL, false, false);
+            MethodTransform patch = diff.asParameterTransformer(ParamTransformTarget.ALL, false, Set.of());
+            patch.apply(methodContext);
+        }
+    }
+
+    public static void upgradeWrapOperationLayered(MethodContext methodContext, MethodQualifier cleanQualifier, MethodQualifier dirtyQualifier) {
+        if (dirtyQualifier.owner() == null || cleanQualifier.desc() == null) {
+            return;
+        }
+        List<Type> originalTargetDesc = List.of(Type.getArgumentTypes(cleanQualifier.desc()));
+        List<Type> modifiedTargetDesc = List.of(Type.getArgumentTypes(dirtyQualifier.desc()));
+        MethodNode methodNode = methodContext.getMixinMethod();
+        List<Type> originalDesc = List.of(Type.getArgumentTypes(methodNode.desc));
+        List<Type> modifiedDesc = ImmutableList.<Type>builder()
+            // Add instance parameter
+            .add(Type.getType(dirtyQualifier.owner()))
+            // Add target parameters
+            .addAll(modifiedTargetDesc)
+            // Add everything after the original owner and target args (such as captured locals)
+            .addAll(originalDesc.subList(1 + originalTargetDesc.size(), originalDesc.size()))
+            .build();
+        // Create diff
+        LayeredParamsDiffSnapshot diff = EnhancedParamsDiff.createLayered(originalDesc, modifiedDesc);
+        if (!diff.isEmpty()) {
+            MethodTransform patch = diff.asParameterTransformer(ParamTransformTarget.ALL, false, Set.of());
             patch.apply(methodContext);
         }
     }

diff --git a/definition/src/main/java/org/sinytra/adapter/patch/fixes/SimpleTypeAdapter.java b/definition/src/main/java/org/sinytra/adapter/patch/fixes/SimpleTypeAdapter.java
@@ -4,6 +4,8 @@
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.InsnList;
 
+import java.util.function.BiConsumer;
+
 public record SimpleTypeAdapter(Type from, Type to, TypePatch adapter) implements TypeAdapter {
     public interface TypePatch {
         void apply(InsnList list, AbstractInsnNode target);
@@ -13,4 +15,12 @@ public interface TypePatch {
     public void apply(InsnList list, AbstractInsnNode target) {
         this.adapter.apply(list, target);
     }
+
+    @Override
+    public TypeAdapter andThen(BiConsumer<InsnList, AbstractInsnNode> consumer) {
+        return new SimpleTypeAdapter(this.from, this.to, (l, t) -> {
+            consumer.accept(l, t);
+            this.adapter.apply(l, t);
+        });
+    }
 }
diff --git a/definition/src/main/java/org/sinytra/adapter/patch/fixes/SupplierTypeAdapter.java b/definition/src/main/java/org/sinytra/adapter/patch/fixes/SupplierTypeAdapter.java
@@ -7,6 +7,7 @@
 import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.TypeInsnNode;
 
+import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 
 public class SupplierTypeAdapter implements TypeAdapterProvider {
@@ -27,5 +28,10 @@ public void apply(InsnList list, AbstractInsnNode target) {
             patch.add(new TypeInsnNode(Opcodes.CHECKCAST, this.to.getInternalName()));
             list.insert(target, patch);
         }
+
+        @Override
+        public TypeAdapter andThen(BiConsumer<InsnList, AbstractInsnNode> consumer) {
+            throw new UnsupportedOperationException();
+        }
     }
 }
diff --git a/definition/src/main/java/org/sinytra/adapter/patch/fixes/TypeAdapter.java b/definition/src/main/java/org/sinytra/adapter/patch/fixes/TypeAdapter.java
@@ -4,10 +4,14 @@
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.InsnList;
 
+import java.util.function.BiConsumer;
+
 public interface TypeAdapter {
     Type from();
 
     Type to();
 
     void apply(InsnList list, AbstractInsnNode target);
+
+    TypeAdapter andThen(BiConsumer<InsnList, AbstractInsnNode> consumer);
 }
diff --git a/...ain/java/org/sinytra/adapter/patch/transformer/dynamic/DynamicModifyVarAtReturnPatch.java b/...ain/java/org/sinytra/adapter/patch/transformer/dynamic/DynamicModifyVarAtReturnPatch.java
@@ -1,11 +1,8 @@
 package org.sinytra.adapter.patch.transformer.dynamic;
 
 import com.mojang.datafixers.util.Pair;
-import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.*;
-import org.objectweb.asm.tree.analysis.SourceInterpreter;
-import org.objectweb.asm.tree.analysis.SourceValue;
 import org.sinytra.adapter.patch.analysis.MethodCallAnalyzer;
 import org.sinytra.adapter.patch.analysis.selector.AnnotationHandle;
 import org.sinytra.adapter.patch.analysis.selector.AnnotationValueHandle;
@@ -16,7 +13,6 @@
 import org.spongepowered.asm.mixin.injection.struct.InjectionInfo;
 import org.spongepowered.asm.mixin.injection.struct.Target;
 
-import javax.annotation.Nullable;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
@@ -153,32 +149,4 @@ private static Pair<AbstractInsnNode, Integer> getTargetPair(ClassNode classNode
         int local = discriminator.findLocal(ctx);
         return Pair.of(targetInsn, local);
     }
-
-    private static class MethodCallInterpreter extends SourceInterpreter {
-        private final MethodInsnNode targetInsn;
-        private List<AbstractInsnNode> targetArgs;
-
-        public MethodCallInterpreter(MethodInsnNode targetInsn) {
-            super(Opcodes.ASM9);
-            this.targetInsn = targetInsn;
-        }
-
-        @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);
-        }
-    }
 }
diff --git a/...on/src/main/java/org/sinytra/adapter/patch/transformer/dynfix/DynFixMethodComparison.java b/...on/src/main/java/org/sinytra/adapter/patch/transformer/dynfix/DynFixMethodComparison.java
@@ -64,7 +64,8 @@ public FixResult apply(ClassNode classNode, MethodNode methodNode, MethodContext
 
         if (methodContext.methodAnnotation().matchesDesc(MixinConstants.WRAP_OPERATION)) {
             return handleWrapOperationToInstanceOf(cleanInjectionInsn, comparisonResult.cleanLabel(), hunkLabels, methodContext)
-                .or(() -> handleWrapOpertationNewInjectionPoint(cleanInjectionInsn, comparisonResult.cleanLabel(), hunkLabels, methodContext))
+                .or(() -> handleWrapOperationAdaptedTarget(cleanInjectionInsn, hunkLabels, methodContext))
+                .or(() -> handleWrapOperationNewInjectionPoint(cleanInjectionInsn, comparisonResult.cleanLabel(), hunkLabels, methodContext))
                 .or(() -> handleTargetModification(hunkLabels, methodContext))
                 .orElse(null);
         }
@@ -181,7 +182,25 @@ private static FixResult handleModifyArgInjectionPoint(AbstractInsnNode cleanInj
         return null;
     }
 
-    private static Optional<FixResult> handleWrapOpertationNewInjectionPoint(AbstractInsnNode cleanInjectionInsn, List<AbstractInsnNode> cleanLabel, List<List<AbstractInsnNode>> hunkLabels, MethodContext methodContext) {
+    private static Optional<FixResult> handleWrapOperationAdaptedTarget(AbstractInsnNode cleanInjectionInsn, List<List<AbstractInsnNode>> hunkLabels, MethodContext methodContext) {
+        if (!(cleanInjectionInsn instanceof MethodInsnNode minsn) || hunkLabels.size() != 1) {
+            return Optional.empty();
+        }
+
+        List<MethodInsnNode> methodCalls = hunkLabels.getFirst().stream()
+            .filter(i -> i instanceof MethodInsnNode)
+            .map(i -> (MethodInsnNode) i)
+            .toList();
+        if (methodCalls.size() != 1) {
+            return Optional.empty();
+        }
+
+        Patch.Result result = WrapOperationSurgeon.tryUpgrade(methodContext, minsn, methodCalls.getLast());
+
+        return Optional.ofNullable(FixResult.of(result, PatchAuditTrail.Match.FULL));
+    } 
+
+    private static Optional<FixResult> handleWrapOperationNewInjectionPoint(AbstractInsnNode cleanInjectionInsn, List<AbstractInsnNode> cleanLabel, List<List<AbstractInsnNode>> hunkLabels, MethodContext methodContext) {
         if (!(cleanInjectionInsn instanceof MethodInsnNode minsn) || hunkLabels.size() != 1) {
             return Optional.empty();
         }