Merge pull request #84 from LlamaLad7/feature/cancellable-sugar

Feature/cancellable sugar

src/main/java/com/llamalad7/mixinextras/sugar/Cancellable.java

@@ -0,0 +1,31 @@
+package com.llamalad7.mixinextras.sugar;
+
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.ModifyConstant;
+import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Allows you to receive a cancellable {@link CallbackInfo} or {@link CallbackInfoReturnable} as appropriate
+ * from any kind of injector. This allows you to optionally cancel the target method without being forced to use
+ * {@link Inject @Inject}.
+ * <p>
+ * The same {@link CallbackInfo}s will be passed to every handler method in a chain of
+ * {@link WrapOperation @WrapOperation}s (i.e. any number of {@link WrapOperation @WrapOperation}s and at most one inner
+ * {@link Redirect @Redirect} / {@link ModifyConstant @ModifyConstant}). This means you can choose to use the
+ * {@link CallbackInfo#isCancelled()} and {@link CallbackInfoReturnable#getReturnValue()} methods to see if the wrapped
+ * handler cancelled, so you can respond accordingly.
+ * <p>
+ * See <a href="https://github.com/LlamaLad7/MixinExtras/wiki/Cancellable">the wiki article</a> for more info.
+ */
+@Target(ElementType.PARAMETER)
+@Retention(RetentionPolicy.CLASS)
+public @interface Cancellable {
+}

src/main/java/com/llamalad7/mixinextras/sugar/impl/CancellableSugarApplicator.java

@@ -0,0 +1,115 @@
+package com.llamalad7.mixinextras.sugar.impl;
+
+import com.llamalad7.mixinextras.injector.StackExtension;
+import com.llamalad7.mixinextras.utils.ASMUtils;
+import com.llamalad7.mixinextras.utils.Decorations;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.*;
+import org.spongepowered.asm.mixin.injection.struct.InjectionInfo;
+import org.spongepowered.asm.mixin.injection.struct.InjectionNodes.InjectionNode;
+import org.spongepowered.asm.mixin.injection.struct.Target;
+
+class CancellableSugarApplicator extends SugarApplicator {
+    CancellableSugarApplicator(InjectionInfo info, SugarParameter parameter) {
+        super(info, parameter);
+    }
+
+    @Override
+    void validate(Target target, InjectionNode node) {
+    }
+
+    @Override
+    void prepare(Target target, InjectionNode node) {
+    }
+
+    @Override
+    void inject(Target target, InjectionNode node, StackExtension stack) {
+        Type ciType = Type.getObjectType(target.getCallbackInfoClass());
+        if (!ciType.equals(paramType)) {
+            throw new IllegalStateException(
+                    String.format(
+                            "@Cancellable sugar has wrong type! Expected %s but got %s!",
+                            ciType.getClassName(),
+                            paramType.getClassName()
+                    )
+            );
+        }
+        int ciIndex = getOrCreateCi(target, node, stack, ciType);
+        stack.extra(1);
+        target.insns.insertBefore(node.getCurrentTarget(), new VarInsnNode(Opcodes.ALOAD, ciIndex));
+    }
+
+    @Override
+    int postProcessingPriority() {
+        // Early, we don't care about being particularly tight compared to e.g. `@Local`s.
+        return -1000;
+    }
+
+    private int getOrCreateCi(Target target, InjectionNode node, StackExtension stack, Type ciType) {
+        if (node.hasDecoration(Decorations.CANCELLABLE_CI_INDEX)) {
+            return node.getDecoration(Decorations.CANCELLABLE_CI_INDEX);
+        }
+        int ciIndex = target.allocateLocal();
+        target.addLocalVariable(ciIndex, "callbackInfo" + ciIndex, ciType.getDescriptor());
+        node.decorate(Decorations.CANCELLABLE_CI_INDEX, ciIndex);
+
+        InsnList init = new InsnList();
+        init.add(new TypeInsnNode(Opcodes.NEW, ciType.getInternalName()));
+        init.add(new InsnNode(Opcodes.DUP));
+        init.add(new LdcInsnNode(target.method.name));
+        init.add(new InsnNode(Opcodes.ICONST_1));
+        init.add(new MethodInsnNode(
+                Opcodes.INVOKESPECIAL,
+                ciType.getInternalName(),
+                "<init>",
+                "(Ljava/lang/String;Z)V",
+                false
+        ));
+        init.add(new VarInsnNode(Opcodes.ASTORE, ciIndex));
+        target.insertBefore(node, init);
+        stack.extra(4);
+
+        SugarPostProcessingExtension.enqueuePostProcessing(this, () -> {
+            InsnList cancellation = new InsnList();
+            LabelNode notCancelled = new LabelNode();
+            cancellation.add(new VarInsnNode(Opcodes.ALOAD, ciIndex));
+            cancellation.add(new MethodInsnNode(
+                    Opcodes.INVOKEVIRTUAL,
+                    ciType.getInternalName(),
+                    "isCancelled",
+                    "()Z",
+                    false
+            ));
+            cancellation.add(new JumpInsnNode(Opcodes.IFEQ, notCancelled));
+            cancellation.add(new VarInsnNode(Opcodes.ALOAD, ciIndex));
+            if (target.returnType.equals(Type.VOID_TYPE)) {
+                cancellation.add(new InsnNode(Opcodes.RETURN));
+            } else if (ASMUtils.isPrimitive(target.returnType)) {
+                cancellation.add(new MethodInsnNode(
+                        Opcodes.INVOKEVIRTUAL,
+                        ciType.getInternalName(),
+                        "getReturnValue" + target.returnType.getDescriptor(),
+                        "()" + target.returnType.getDescriptor(),
+                        false
+                ));
+                cancellation.add(new InsnNode(target.returnType.getOpcode(Opcodes.IRETURN)));
+            } else {
+                cancellation.add(new MethodInsnNode(
+                        Opcodes.INVOKEVIRTUAL,
+                        ciType.getInternalName(),
+                        "getReturnValue",
+                        "()Ljava/lang/Object;",
+                        false
+                ));
+                cancellation.add(new TypeInsnNode(Opcodes.CHECKCAST, target.returnType.getInternalName()));
+                cancellation.add(new InsnNode(Opcodes.ARETURN));
+            }
+            cancellation.add(notCancelled);
+            target.insns.insert(node.getCurrentTarget(), cancellation);
+            // No need to adjust the stack because we only increase the height by at most 2, which is covered by
+            // our bump of 4 earlier.
+        });
+        return ciIndex;
+    }
+}

src/main/java/com/llamalad7/mixinextras/sugar/impl/LocalSugarApplicator.java

@@ -73,6 +73,12 @@ void inject(Target target, InjectionNode node, StackExtension stack) {
         }
     }
 
+    @Override
+    int postProcessingPriority() {
+        // Late, we need to be tight around the handler calls to ensure proper initialization and disposal.
+        return 1000;
+    }
+
     private void initAndLoadLocalRef(Target target, InjectionNode node, int index, StackExtension stack) {
         String refName = LocalRefClassGenerator.getForType(targetLocalType);
         int refIndex = getOrCreateRef(target, node, index, refName, stack);

src/main/java/com/llamalad7/mixinextras/sugar/impl/SugarApplicator.java

@@ -2,8 +2,10 @@
 
 import com.llamalad7.mixinextras.injector.StackExtension;
 import com.llamalad7.mixinextras.service.MixinExtrasService;
+import com.llamalad7.mixinextras.sugar.Cancellable;
 import com.llamalad7.mixinextras.sugar.Local;
 import com.llamalad7.mixinextras.sugar.Share;
+import com.llamalad7.mixinextras.utils.ASMUtils;
 import com.llamalad7.mixinextras.utils.CompatibilityHelper;
 import org.apache.commons.lang3.tuple.Pair;
 import org.objectweb.asm.Type;
@@ -26,6 +28,7 @@ abstract class SugarApplicator {
 
     static {
         List<Pair<Class<? extends Annotation>, Class<? extends SugarApplicator>>> sugars = Arrays.asList(
+                Pair.of(Cancellable.class, CancellableSugarApplicator.class),
                 Pair.of(Local.class, LocalSugarApplicator.class),
                 Pair.of(Share.class, ShareSugarApplicator.class)
         );
@@ -60,6 +63,15 @@ abstract class SugarApplicator {
 
     abstract void inject(Target target, InjectionNode node, StackExtension stack);
 
+    int postProcessingPriority() {
+        throw new UnsupportedOperationException(
+                String.format(
+                        "Sugar type %s does not support post-processing! Please inform LlamaLad7!",
+                        ASMUtils.annotationToString(sugar)
+                )
+        );
+    }
+
     static SugarApplicator create(InjectionInfo info, SugarParameter parameter) {
         try {
             Class<? extends SugarApplicator> clazz = MAP.get(parameter.sugar.desc);

src/main/java/com/llamalad7/mixinextras/sugar/impl/SugarPostProcessingExtension.java

@@ -5,16 +5,14 @@
 import org.spongepowered.asm.mixin.transformer.ext.IExtension;
 import org.spongepowered.asm.mixin.transformer.ext.ITargetClassContext;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 public class SugarPostProcessingExtension implements IExtension {
-    private static final Map<String, List<Runnable>> POST_PROCESSING_TASKS = new HashMap<>();
+    private static final Map<String, List<Task>> POST_PROCESSING_TASKS = new HashMap<>();
 
     static void enqueuePostProcessing(SugarApplicator applicator, Runnable task) {
-        POST_PROCESSING_TASKS.computeIfAbsent(applicator.info.getClassNode().name, k -> new ArrayList<>()).add(task);
+        POST_PROCESSING_TASKS.computeIfAbsent(applicator.info.getClassNode().name, k -> new ArrayList<>())
+                .add(new Task(applicator.postProcessingPriority(), task));
     }
 
     @Override
@@ -29,14 +27,33 @@ public void preApply(ITargetClassContext context) {
     @Override
     public void postApply(ITargetClassContext context) {
         String targetName = context.getClassNode().name;
-        List<Runnable> tasks = POST_PROCESSING_TASKS.get(targetName);
+        List<Task> tasks = POST_PROCESSING_TASKS.remove(targetName);
         if (tasks != null) {
-            tasks.forEach(Runnable::run);
-            POST_PROCESSING_TASKS.remove(targetName);
+            Collections.sort(tasks);
+            tasks.forEach(Task::run);
         }
     }
 
     @Override
     public void export(MixinEnvironment env, String name, boolean force, ClassNode classNode) {
     }
+
+    private static class Task implements Comparable<Task> {
+        private final int priority;
+        private final Runnable body;
+
+        public Task(int priority, Runnable body) {
+            this.priority = priority;
+            this.body = body;
+        }
+
+        public void run() {
+            body.run();
+        }
+
+        @Override
+        public int compareTo(Task o) {
+            return Integer.compare(priority, o.priority);
+        }
+    }
 }

src/main/java/com/llamalad7/mixinextras/utils/Decorations.java

@@ -33,4 +33,9 @@ public class Decorations {
      * Stores that this node has been wrapped by a {@link WrapOperation}.
      */
     public static final String WRAPPED = "mixinextras_wrappedOperation";
+
+    /**
+     * Stores the shared CallbackInfo local index for this target instruction.
+     */
+    public static final String CANCELLABLE_CI_INDEX = "mixinextras_cancellableCiIndex";
 }