Skip to content

Commit

Permalink
Transform WrapOperation to target instanceof call
Browse files Browse the repository at this point in the history
  • Loading branch information
Su5eD committed Jul 18, 2024
1 parent 303e382 commit 1f0604d
Show file tree
Hide file tree
Showing 17 changed files with 454 additions and 180 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -96,8 +96,8 @@ public LocalVariableLookup dirtyLocalsTable() {
@Override
public MethodQualifier getTargetMethodQualifier() {
// Get method targets
List<String> methodRefs = methodAnnotation().<List<String>>getValue("method").orElseThrow().get();
if (methodRefs.size() > 1) {
List<String> methodRefs = methodAnnotation().<List<String>>getValue("method").map(AnnotationValueHandle::get).orElseGet(Collections::emptyList);
if (methodRefs.size() != 1) {
// We only support single method targets for now
return null;
}
Expand Down Expand Up @@ -246,7 +246,7 @@ private InsnList computeSlicedInsns(ISliceContext context, AnnotationNode annota

@Nullable
private TargetPair findInjectionTarget(ClassLookup lookup) {
Pair<ClassNode, List<MethodNode>> pair = findInjectionTargetCandidates(lookup);
Pair<ClassNode, List<MethodNode>> pair = findInjectionTargetCandidates(lookup, false);
if (pair == null) {
return null;
}
Expand All @@ -263,7 +263,7 @@ private TargetPair findInjectionTarget(ClassLookup lookup) {
}

@Nullable
public Pair<ClassNode, List<MethodNode>> findInjectionTargetCandidates(ClassLookup lookup) {
public Pair<ClassNode, List<MethodNode>> findInjectionTargetCandidates(ClassLookup lookup, boolean ignoreDesc) {
// Find target method qualifier
MethodQualifier qualifier = getTargetMethodQualifier();
if (qualifier == null || qualifier.name() == null) {
Expand All @@ -288,7 +288,7 @@ public Pair<ClassNode, List<MethodNode>> findInjectionTargetCandidates(ClassLook
// Find target method in class
String desc = qualifier.desc();
List<MethodNode> candidates = targetClass.methods.stream()
.filter(mtd -> mtd.name.equals(qualifier.name()) && (desc == null || mtd.desc.equals(desc)))
.filter(mtd -> mtd.name.equals(qualifier.name()) && (ignoreDesc || desc == null || mtd.desc.equals(desc)))
.toList();
// If there's multiple candidates, try removing bouncer methods
if (candidates.size() > 1 && desc == null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,22 @@ public static List<AbstractInsnNode> findMethodCallParamInsns(MethodNode methodN
return interpreter.getTargetArgs();
}

@Nullable
public static List<AbstractInsnNode> findFullMethodCallParamInsns(MethodNode methodNode, MethodInsnNode minsn) {
List<AbstractInsnNode> insns = findMethodCallParamInsns(methodNode, minsn);
if (minsn != null) {
List<AbstractInsnNode> fullInsns = new ArrayList<>();
for (AbstractInsnNode i = insns.getFirst(); i != null ; i = i.getNext()) {
if (i == minsn) {
break;
}
fullInsns.add(i);
}
return fullInsns;
}
return null;
}

public static <T> List<T> analyzeMethod(MethodNode methodNode, NaryOperationHandler<T> handler) {
return analyzeMethod(methodNode, (insn, values) -> true, handler);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
package org.sinytra.adapter.patch.analysis;

import com.mojang.datafixers.util.Pair;
import org.jetbrains.annotations.Nullable;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.FrameNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.MethodNode;
import org.sinytra.adapter.patch.api.MethodContext;

import java.util.*;
import java.util.stream.Stream;

public class MethodLabelComparator {
public record ComparisonResult(List<List<AbstractInsnNode>> patchedLabels, List<AbstractInsnNode> cleanLabel) {}

@Nullable
public static ComparisonResult findPatchedLabels(AbstractInsnNode cleanInjectionInsn, MethodContext methodContext) {
List<List<AbstractInsnNode>> cleanLabels = getLabelsInMethod(methodContext.findCleanInjectionTarget().methodNode());
List<List<AbstractInsnNode>> cleanLabelsOriginal = List.copyOf(cleanLabels);

List<List<AbstractInsnNode>> cleanMatchedLabels = cleanLabels.stream()
.filter(insns -> insns.contains(cleanInjectionInsn))
.toList();
if (cleanMatchedLabels.size() != 1) {
return null;
}
List<AbstractInsnNode> cleanLabel = cleanMatchedLabels.getFirst();

List<List<AbstractInsnNode>> dirtyLabels = getLabelsInMethod(methodContext.findDirtyInjectionTarget().methodNode());
List<List<AbstractInsnNode>> dirtyLabelsOriginal = List.copyOf(dirtyLabels);

Map<List<AbstractInsnNode>, List<AbstractInsnNode>> matchedLabels = new LinkedHashMap<>();
for (List<AbstractInsnNode> cleanInsns : cleanLabelsOriginal) {
List<List<AbstractInsnNode>> candidates = new ArrayList<>();

for (List<AbstractInsnNode> dirtyInsns : dirtyLabels) {
if (InstructionMatcher.test(cleanInsns, dirtyInsns, InsnComparator.IGNORE_VAR_INDEX | InsnComparator.IGNORE_LINE_NUMBERS)) {
candidates.add(dirtyInsns);
}
}
// TODO Try and come up with something better
// This prevents messing up the order of labels
// Without this countermeasure, it might happen that a label that was deleted will match a seemingly identical label somewhere else in the method, which is wrong
// We disable any duplicated until we can properly handle such cases
if (candidates.size() == 1) {
List<AbstractInsnNode> dirtyInsns = candidates.getFirst();
matchedLabels.put(cleanInsns, dirtyInsns);
cleanLabels.remove(cleanInsns);
dirtyLabels.remove(dirtyInsns);
}
}

Pair<List<AbstractInsnNode>, List<AbstractInsnNode>> patchRange = findPatchHunkRange(cleanLabel, cleanLabelsOriginal, matchedLabels);
if (patchRange == null) {
return null;
}

List<List<AbstractInsnNode>> patchedLabels = dirtyLabelsOriginal.subList(dirtyLabelsOriginal.indexOf(patchRange.getFirst()) + 1, dirtyLabelsOriginal.indexOf(patchRange.getSecond()));
return new ComparisonResult(patchedLabels, cleanLabel);
}

@Nullable
private static Pair<List<AbstractInsnNode>, List<AbstractInsnNode>> findPatchHunkRange(List<AbstractInsnNode> cleanLabel, List<List<AbstractInsnNode>> cleanLabels, Map<List<AbstractInsnNode>, List<AbstractInsnNode>> matchedLabels) {
// Find last matched dirty label BEFORE the injection point
List<AbstractInsnNode> dirtyLabelBefore = Stream.iterate(cleanLabels.indexOf(cleanLabel), i -> i >= 0, i -> i - 1)
.map(i -> matchedLabels.get(cleanLabels.get(i)))
.filter(Objects::nonNull)
.findFirst()
.orElse(null);
if (dirtyLabelBefore == null) {
return null;
}

// Find first matched dirty label AFTER the injection point
List<AbstractInsnNode> dirtyLabelAfter = Stream.iterate(cleanLabels.indexOf(cleanLabel), i -> i < cleanLabels.size(), i -> i + 1)
.map(i -> matchedLabels.get(cleanLabels.get(i)))
.filter(Objects::nonNull)
.findFirst()
.orElse(null);
if (dirtyLabelAfter == null) {
return null;
}

return Pair.of(dirtyLabelBefore, dirtyLabelAfter);
}

private static List<List<AbstractInsnNode>> getLabelsInMethod(MethodNode methodNode) {
List<List<AbstractInsnNode>> list = new ArrayList<>();
List<AbstractInsnNode> workingList = null;
for (AbstractInsnNode insn : methodNode.instructions) {
if (insn instanceof FrameNode) {
continue;
}
if (insn instanceof LabelNode) {
if (workingList != null) {
list.add(workingList);
}
workingList = new ArrayList<>();
}
workingList.add(insn);
}
return list;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ public interface MethodContext {
List<AbstractInsnNode> computeInjectionTargetInsns(@Nullable TargetPair target);

@Nullable
Pair<ClassNode, List<MethodNode>> findInjectionTargetCandidates(ClassLookup lookup);
Pair<ClassNode, List<MethodNode>> findInjectionTargetCandidates(ClassLookup lookup, boolean ignoreDesc);

void updateDescription(List<Type> parameters);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,5 +32,6 @@ public class MixinConstants {
public static final String UNIQUE = "Lorg/spongepowered/asm/mixin/Unique;";
public static final String SHADOW = "Lorg/spongepowered/asm/mixin/Shadow;";
public static final String COERCE = "Lorg/spongepowered/asm/mixin/injection/Coerce;";
public static final String CONSTANT = "Lorg/spongepowered/asm/mixin/injection/Constant;";
public static final List<Integer> LVT_COMPATIBILITY_LEVELS = List.of(FabricUtil.COMPATIBILITY_0_10_0, FabricUtil.COMPATIBILITY_0_9_2);
}
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Predicate;
import java.util.function.Supplier;

public interface Patch {
static ClassPatchBuilder builder() {
Expand Down Expand Up @@ -44,6 +45,10 @@ public Result or(Result other) {
}
return this;
}

public Result orElseGet(Supplier<Result> other) {
return this == PASS ? other.get() : this;
}
}

interface Builder<T extends Builder<T>> extends MethodTransformBuilder<T> {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
package org.sinytra.adapter.patch.transformer;

import com.google.common.collect.ImmutableList;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.GeneratorAdapter;
import org.objectweb.asm.commons.Method;
import org.objectweb.asm.tree.*;
import org.sinytra.adapter.patch.analysis.MethodCallAnalyzer;
import org.sinytra.adapter.patch.api.*;
import org.sinytra.adapter.patch.util.AdapterUtil;
import org.sinytra.adapter.patch.util.OpcodeUtil;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.stream.Stream;

public record MirrorableExtractMixin(String destinationClass, MethodInsnNode destinationMethodInvocation) implements MethodTransform {
@Override
public Patch.Result apply(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, PatchContext context) {
Type selfType = Type.getObjectType(methodContext.findDirtyInjectionTarget().classNode().name);
Type[] params = Type.getArgumentTypes(this.destinationMethodInvocation.desc);
int selfIndex = Stream.of(Stream.iterate(0, i -> i < params.length, i -> i + 1)
.filter(i -> params[i].equals(selfType))
.toList())
.filter(list -> list.size() == 1)
.map(List::getFirst)
.findFirst()
.orElse(-1);
if (selfIndex == -1) {
return Patch.Result.PASS;
}

List<AbstractInsnNode> callInsns = MethodCallAnalyzer.findMethodCallParamInsns(methodContext.findDirtyInjectionTarget().methodNode(), this.destinationMethodInvocation);
if (callInsns == null || callInsns.size() <= selfIndex) {
return Patch.Result.PASS;
}
AbstractInsnNode selfParamInsn = callInsns.get(selfIndex);
if (!(selfParamInsn instanceof VarInsnNode varInsn) || varInsn.getOpcode() != Opcodes.ALOAD || varInsn.var != 0) {
return Patch.Result.PASS;
}
// Cool, out instance is passed into the method. Now let's inject there and call the old mixin method
ClassNode generatedTarget = methodContext.patchContext().environment().classGenerator().getOrGenerateMixinClass(methodContext.getMixinClass(), this.destinationClass, null);
methodContext.patchContext().environment().refmapHolder().copyEntries(methodContext.getMixinClass().name, generatedTarget.name);
// Generate a method with the same injector annotation
MethodNode originalMixinMethod = methodContext.getMixinMethod();
String name = originalMixinMethod.name + "$adapter$mirror$" + AdapterUtil.randomString(5);
List<Type> originalParams = List.of(Type.getArgumentTypes(originalMixinMethod.desc));
List<Type> newParams = ImmutableList.<Type>builder().add(Type.getArgumentTypes(this.destinationMethodInvocation.desc)).add(AdapterUtil.CI_TYPE).build();
// Make sure we have all required params
if (!new HashSet<>(newParams).containsAll(originalParams)) {
return Patch.Result.PASS;
}

String desc = Type.getMethodDescriptor(Type.VOID_TYPE, newParams.toArray(Type[]::new));
// Change target
BundledMethodTransform.builder().modifyTarget(this.destinationMethodInvocation.name + this.destinationMethodInvocation.desc).apply(methodContext);
MethodNode invokerMixinMethod = (MethodNode) generatedTarget.visitMethod(Opcodes.ACC_PRIVATE | Opcodes.ACC_STATIC, name, desc, null, null);
invokerMixinMethod.visibleAnnotations = new ArrayList<>(originalMixinMethod.visibleAnnotations);
// Make original mixin a unique public method
originalMixinMethod.access = OpcodeUtil.setAccessVisibility(originalMixinMethod.access, Opcodes.ACC_PUBLIC);
originalMixinMethod.visibleAnnotations.remove(methodContext.methodAnnotation().unwrap());
originalMixinMethod.visitAnnotation(MixinConstants.UNIQUE, true);
// Now call the original mixin
GeneratorAdapter gen = new GeneratorAdapter(invokerMixinMethod, invokerMixinMethod.access, invokerMixinMethod.name, invokerMixinMethod.desc);
gen.newLabel();
gen.loadArg(selfIndex);
for (Type type : originalParams) {
gen.loadArg(newParams.indexOf(type));
}
gen.invokeVirtual(selfType, new Method(originalMixinMethod.name, originalMixinMethod.desc));
gen.newLabel();
gen.returnValue();
gen.newLabel();
gen.endMethod();
return Patch.Result.APPLY;
}
}
Loading

0 comments on commit 1f0604d

Please sign in to comment.