Add slot binding

2 files changed, 208 insertions, 134 deletions

diff --git a/src/main/java/moe/nea/firmament/init/HandledScreenRiser.java b/src/main/java/moe/nea/firmament/init/HandledScreenRiser.java
index bcb754c..1fbbe45 100644
--- a/src/main/java/moe/nea/firmament/init/HandledScreenRiser.java
+++ b/src/main/java/moe/nea/firmament/init/HandledScreenRiser.java
@@ -2,8 +2,8 @@
 package moe.nea.firmament.init;
 
 import me.shedaniel.mm.api.ClassTinkerers;
-import net.minecraft.client.gui.screen.Screen;
-import net.minecraft.client.gui.screen.ingame.HandledScreen;
+import net.minecraft.client.gui.Element;
+import net.minecraft.client.gui.ParentElement;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.ClassNode;
@@ -17,69 +17,139 @@ import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.VarInsnNode;
 
 import java.lang.reflect.Modifier;
+import java.util.function.Consumer;
 
 public class HandledScreenRiser extends RiserUtils {
-    @IntermediaryName(net.minecraft.client.gui.screen.Screen.class)
-    String Screen;
-    @IntermediaryName(net.minecraft.client.gui.screen.ingame.HandledScreen.class)
-    String HandledScreen;
-    Type mouseScrolledDesc = Type.getMethodType(Type.BOOLEAN_TYPE, Type.DOUBLE_TYPE, Type.DOUBLE_TYPE, Type.DOUBLE_TYPE, Type.DOUBLE_TYPE);
-    String mouseScrolled = remapper.mapMethodName("intermediary", "net.minecraft.class_364", "method_25401",
-                                                  mouseScrolledDesc.getDescriptor());
+	@IntermediaryName(net.minecraft.client.gui.screen.Screen.class)
+	String Screen;
+	@IntermediaryName(net.minecraft.client.gui.screen.ingame.HandledScreen.class)
+	String HandledScreen;
+	Type mouseScrolledDesc = Type.getMethodType(Type.BOOLEAN_TYPE, Type.DOUBLE_TYPE, Type.DOUBLE_TYPE, Type.DOUBLE_TYPE, Type.DOUBLE_TYPE);
+	String mouseScrolled = remapper.mapMethodName("intermediary", "net.minecraft.class_364", "method_25401",
+	                                              mouseScrolledDesc.getDescriptor());
+	// boolean keyReleased(int keyCode, int scanCode, int modifiers)
+	Type keyReleasedDesc = Type.getMethodType(Type.BOOLEAN_TYPE, Type.INT_TYPE, Type.INT_TYPE, Type.INT_TYPE);
+	String keyReleased = remapper.mapMethodName("intermediary", Intermediary.<Element>className(),
+	                                            Intermediary.methodName(Element::keyReleased),
+	                                            keyReleasedDesc.getDescriptor());
 
-    @Override
-    public void addTinkerers() {
-        ClassTinkerers.addTransformation(HandledScreen, this::handle, true);
-    }
 
-    void handle(ClassNode classNode) {
-        MethodNode mouseScrolledNode = findMethod(classNode, mouseScrolled, mouseScrolledDesc);
-        if (mouseScrolledNode == null) {
-            mouseScrolledNode = new MethodNode(
-                Modifier.PUBLIC,
-                mouseScrolled,
-                mouseScrolledDesc.getDescriptor(),
-                null,
-                new String[0]
-            );
-            var insns = mouseScrolledNode.instructions;
-            // ALOAD 0, load this
-            insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
-            // DLOAD 1-4, load the 4 argument doubles. Note that since doubles are two entries wide we skip 2 each time.
-            insns.add(new VarInsnNode(Opcodes.DLOAD, 1));
-            insns.add(new VarInsnNode(Opcodes.DLOAD, 3));
-            insns.add(new VarInsnNode(Opcodes.DLOAD, 5));
-            insns.add(new VarInsnNode(Opcodes.DLOAD, 7));
-            // INVOKESPECIAL call super method
-            insns.add(new MethodInsnNode(Opcodes.INVOKESPECIAL, getTypeForClassName(Screen).getInternalName(), mouseScrolled, mouseScrolledDesc.getDescriptor()));
-            // IRETURN return int on stack (booleans are int at runtime)
-            insns.add(new InsnNode(Opcodes.IRETURN));
-            classNode.methods.add(mouseScrolledNode);
-        }
+	@Override
+	public void addTinkerers() {
+		ClassTinkerers.addTransformation(HandledScreen, this::addMouseScroll, true);
+		ClassTinkerers.addTransformation(HandledScreen, this::addKeyReleased, true);
+	}
 
-        var insns = new InsnList();
-        // ALOAD 0, load this
-        insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
-        // DLOAD 1-4, load the 4 argument doubles. Note that since doubles are two entries wide we skip 2 each time.
-        insns.add(new VarInsnNode(Opcodes.DLOAD, 1));
-        insns.add(new VarInsnNode(Opcodes.DLOAD, 3));
-        insns.add(new VarInsnNode(Opcodes.DLOAD, 5));
-        insns.add(new VarInsnNode(Opcodes.DLOAD, 7));
-        // INVOKEVIRTUAL call custom handler
-        insns.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL,
-                                     getTypeForClassName(HandledScreen).getInternalName(),
-                                     "mouseScrolled_firmament",
-                                     mouseScrolledDesc.getDescriptor()));
-        // Create jump target (but not insert it yet)
-        var jumpIfFalse = new LabelNode();
-        // IFEQ (if returned boolean == 0), jump to jumpIfFalse
-        insns.add(new JumpInsnNode(Opcodes.IFEQ, jumpIfFalse));
-        // LDC 1 (as int, which is what booleans are at runtime)
-        insns.add(new LdcInsnNode(1));
-        // IRETURN return int on stack (booleans are int at runtime)
-        insns.add(new InsnNode(Opcodes.IRETURN));
-        insns.add(jumpIfFalse);
-        mouseScrolledNode.instructions.insert(insns);
-    }
+	/**
+	 * Insert a handler that roughly inserts the following code at the beginning of the instruction list:
+	 * <code><pre>
+	 * if (insertInvoke(insertLoads)) return true
+	 * </pre></code>
+	 *
+	 * @param node         The method node to prepend the instructions to
+	 * @param insertLoads  insert all the loads, including the {@code this} parameter
+	 * @param insertInvoke insert the invokevirtual/invokestatic call
+	 */
+	void insertTrueHandler(MethodNode node,
+	                       Consumer<InsnList> insertLoads,
+	                       Consumer<InsnList> insertInvoke) {
+
+		var insns = new InsnList();
+		insertLoads.accept(insns);
+		insertInvoke.accept(insns);
+		// Create jump target (but not insert it yet)
+		var jumpIfFalse = new LabelNode();
+		// IFEQ (if returned boolean == 0), jump to jumpIfFalse
+		insns.add(new JumpInsnNode(Opcodes.IFEQ, jumpIfFalse));
+		// LDC 1 (as int, which is what booleans are at runtime)
+		insns.add(new LdcInsnNode(1));
+		// IRETURN return int on stack (booleans are int at runtime)
+		insns.add(new InsnNode(Opcodes.IRETURN));
+		insns.add(jumpIfFalse);
+		node.instructions.insert(insns);
+	}
+
+	void addKeyReleased(ClassNode classNode) {
+		var keyReleasedNode = findMethod(classNode, keyReleased, keyReleasedDesc);
+		if (keyReleasedNode == null) {
+			keyReleasedNode = new MethodNode(
+				Modifier.PUBLIC,
+				keyReleased,
+				keyReleasedDesc.getDescriptor(),
+				null,
+				new String[0]
+			);
+			var insns = keyReleasedNode.instructions;
+			// ALOAD 0, load this
+			insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
+			// ILOAD 1-3, load args
+			insns.add(new VarInsnNode(Opcodes.ILOAD, 1));
+			insns.add(new VarInsnNode(Opcodes.ILOAD, 2));
+			insns.add(new VarInsnNode(Opcodes.ILOAD, 3));
+			// INVOKESPECIAL call super method
+			insns.add(new MethodInsnNode(Opcodes.INVOKESPECIAL, getTypeForClassName(Screen).getInternalName(),
+			                             keyReleased, keyReleasedDesc.getDescriptor()));
+			// IRETURN return int on stack (booleans are int at runtime)
+			insns.add(new InsnNode(Opcodes.IRETURN));
+			classNode.methods.add(keyReleasedNode);
+		}
+		insertTrueHandler(keyReleasedNode, insns -> {
+			// ALOAD 0, load this
+			insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
+			// ILOAD 1-3, load args
+			insns.add(new VarInsnNode(Opcodes.ILOAD, 1));
+			insns.add(new VarInsnNode(Opcodes.ILOAD, 2));
+			insns.add(new VarInsnNode(Opcodes.ILOAD, 3));
+		}, insns -> {
+			// INVOKEVIRTUAL call custom handler
+			insns.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL,
+			                             getTypeForClassName(HandledScreen).getInternalName(),
+			                             "keyReleased_firmament",
+			                             keyReleasedDesc.getDescriptor()));
+		});
+	}
+
+	void addMouseScroll(ClassNode classNode) {
+		MethodNode mouseScrolledNode = findMethod(classNode, mouseScrolled, mouseScrolledDesc);
+		if (mouseScrolledNode == null) {
+			mouseScrolledNode = new MethodNode(
+				Modifier.PUBLIC,
+				mouseScrolled,
+				mouseScrolledDesc.getDescriptor(),
+				null,
+				new String[0]
+			);
+			var insns = mouseScrolledNode.instructions;
+			// ALOAD 0, load this
+			insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
+			// DLOAD 1-4, load the 4 argument doubles. Note that since doubles are two entries wide we skip 2 each time.
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 1));
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 3));
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 5));
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 7));
+			// INVOKESPECIAL call super method
+			insns.add(new MethodInsnNode(Opcodes.INVOKESPECIAL, getTypeForClassName(Screen).getInternalName(), mouseScrolled, mouseScrolledDesc.getDescriptor()));
+			// IRETURN return int on stack (booleans are int at runtime)
+			insns.add(new InsnNode(Opcodes.IRETURN));
+			classNode.methods.add(mouseScrolledNode);
+		}
+
+		insertTrueHandler(mouseScrolledNode, insns -> {
+			// ALOAD 0, load this
+			insns.add(new VarInsnNode(Opcodes.ALOAD, 0));
+			// DLOAD 1-4, load the 4 argument doubles. Note that since doubles are two entries wide we skip 2 each time.
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 1));
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 3));
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 5));
+			insns.add(new VarInsnNode(Opcodes.DLOAD, 7));
+		}, insns -> {
+			// INVOKEVIRTUAL call custom handler
+			insns.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL,
+			                             getTypeForClassName(HandledScreen).getInternalName(),
+			                             "mouseScrolled_firmament",
+			                             mouseScrolledDesc.getDescriptor()));
+
+		});
+	}
 
 }
diff --git a/src/main/java/moe/nea/firmament/mixins/MixinHandledScreen.java b/src/main/java/moe/nea/firmament/mixins/MixinHandledScreen.java
index 741ba4b..fd50c72 100644
--- a/src/main/java/moe/nea/firmament/mixins/MixinHandledScreen.java
+++ b/src/main/java/moe/nea/firmament/mixins/MixinHandledScreen.java
@@ -24,77 +24,81 @@ import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
 @Mixin(HandledScreen.class)
 public abstract class MixinHandledScreen<T extends ScreenHandler> {
 
-    @Shadow
-    @Final
-    protected T handler;
-
-    @Shadow
-    public abstract T getScreenHandler();
-
-    @Shadow
-    protected int y;
-    @Shadow
-    protected int x;
-    @Unique
-    PlayerInventory playerInventory;
-
-    @Inject(method = "<init>", at = @At("TAIL"))
-    public void savePlayerInventory(ScreenHandler handler, PlayerInventory inventory, Text title, CallbackInfo ci) {
-        this.playerInventory = inventory;
-    }
-
-    @Inject(method = "keyPressed", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;handleHotbarKeyPressed(II)Z", shift = At.Shift.BEFORE), cancellable = true)
-    public void onKeyPressed(int keyCode, int scanCode, int modifiers, CallbackInfoReturnable<Boolean> cir) {
-        if (HandledScreenKeyPressedEvent.Companion.publish(new HandledScreenKeyPressedEvent((HandledScreen<?>) (Object) this, keyCode, scanCode, modifiers)).getCancelled()) {
-            cir.setReturnValue(true);
-        }
-    }
-
-    @Inject(method = "mouseClicked", at = @At("HEAD"), cancellable = true)
-    public void onMouseClicked(double mouseX, double mouseY, int button, CallbackInfoReturnable<Boolean> cir) {
-        if (HandledScreenClickEvent.Companion.publish(new HandledScreenClickEvent((HandledScreen<?>) (Object) this, mouseX, mouseY, button)).getCancelled()) {
-            cir.setReturnValue(true);
-        }
-    }
-
-    @Inject(method = "render", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;drawForeground(Lnet/minecraft/client/gui/DrawContext;II)V", shift = At.Shift.AFTER))
-    public void onAfterRenderForeground(DrawContext context, int mouseX, int mouseY, float delta, CallbackInfo ci) {
-        context.getMatrices().push();
-        context.getMatrices().translate(-x, -y, 0);
-        HandledScreenForegroundEvent.Companion.publish(new HandledScreenForegroundEvent((HandledScreen<?>) (Object) this, context, mouseX, mouseY, delta));
-        context.getMatrices().pop();
-    }
-
-    @Inject(method = "onMouseClick(Lnet/minecraft/screen/slot/Slot;IILnet/minecraft/screen/slot/SlotActionType;)V", at = @At("HEAD"), cancellable = true)
-    public void onMouseClickedSlot(Slot slot, int slotId, int button, SlotActionType actionType, CallbackInfo ci) {
-        if (slotId == -999 && getScreenHandler() != null && actionType == SlotActionType.PICKUP) { // -999 is code for "clicked outside the main window"
-            ItemStack cursorStack = getScreenHandler().getCursorStack();
-            if (cursorStack != null && IsSlotProtectedEvent.shouldBlockInteraction(slot, SlotActionType.THROW, cursorStack)) {
-                ci.cancel();
-                return;
-            }
-        }
-        if (IsSlotProtectedEvent.shouldBlockInteraction(slot, actionType)) {
-            ci.cancel();
-            return;
-        }
-        if (actionType == SlotActionType.SWAP && 0 <= button && button < 9) {
-            if (IsSlotProtectedEvent.shouldBlockInteraction(new Slot(playerInventory, button, 0, 0), actionType)) {
-                ci.cancel();
-            }
-        }
-    }
-
-
-    @Inject(method = "render", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;drawSlot(Lnet/minecraft/client/gui/DrawContext;Lnet/minecraft/screen/slot/Slot;)V", shift = At.Shift.AFTER), locals = LocalCapture.CAPTURE_FAILHARD)
-    public void onAfterDrawSlot(DrawContext context, int mouseX, int mouseY, float delta, CallbackInfo ci, int i, int j, int k, Slot slot) {
-        SlotRenderEvents.After event = new SlotRenderEvents.After(context, slot, mouseX, mouseY, delta);
-        SlotRenderEvents.After.Companion.publish(event);
-    }
-
-    @Inject(method = "render", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;drawSlot(Lnet/minecraft/client/gui/DrawContext;Lnet/minecraft/screen/slot/Slot;)V", shift = At.Shift.BEFORE), locals = LocalCapture.CAPTURE_FAILHARD)
-    public void onBeforeDrawSlot(DrawContext context, int mouseX, int mouseY, float delta, CallbackInfo ci, int i, int j, int k, Slot slot) {
-        SlotRenderEvents.Before event = new SlotRenderEvents.Before(context, slot, mouseX, mouseY, delta);
-        SlotRenderEvents.Before.Companion.publish(event);
-    }
+	@Shadow
+	@Final
+	protected T handler;
+
+	@Shadow
+	public abstract T getScreenHandler();
+
+	@Shadow
+	protected int y;
+	@Shadow
+	protected int x;
+	@Unique
+	PlayerInventory playerInventory;
+
+	@Inject(method = "<init>", at = @At("TAIL"))
+	public void savePlayerInventory(ScreenHandler handler, PlayerInventory inventory, Text title, CallbackInfo ci) {
+		this.playerInventory = inventory;
+	}
+
+	@Inject(method = "keyPressed", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;handleHotbarKeyPressed(II)Z", shift = At.Shift.BEFORE), cancellable = true)
+	public void onKeyPressed(int keyCode, int scanCode, int modifiers, CallbackInfoReturnable<Boolean> cir) {
+		if (HandledScreenKeyPressedEvent.Companion.publish(new HandledScreenKeyPressedEvent((HandledScreen<?>) (Object) this, keyCode, scanCode, modifiers)).getCancelled()) {
+			cir.setReturnValue(true);
+		}
+	}
+
+	@Inject(method = "mouseClicked", at = @At("HEAD"), cancellable = true)
+	public void onMouseClicked(double mouseX, double mouseY, int button, CallbackInfoReturnable<Boolean> cir) {
+		if (HandledScreenClickEvent.Companion.publish(new HandledScreenClickEvent((HandledScreen<?>) (Object) this, mouseX, mouseY, button)).getCancelled()) {
+			cir.setReturnValue(true);
+		}
+	}
+
+	boolean keyReleased_firmament(int keyCode, int scanCode, int modifiers) {
+		return HandledScreenKeyReleasedEvent.Companion.publish(new HandledScreenKeyReleasedEvent((HandledScreen<?>) (Object) this, keyCode, scanCode, modifiers)).getCancelled();
+	}
+
+	@Inject(method = "render", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;drawForeground(Lnet/minecraft/client/gui/DrawContext;II)V", shift = At.Shift.AFTER))
+	public void onAfterRenderForeground(DrawContext context, int mouseX, int mouseY, float delta, CallbackInfo ci) {
+		context.getMatrices().push();
+		context.getMatrices().translate(-x, -y, 0);
+		HandledScreenForegroundEvent.Companion.publish(new HandledScreenForegroundEvent((HandledScreen<?>) (Object) this, context, mouseX, mouseY, delta));
+		context.getMatrices().pop();
+	}
+
+	@Inject(method = "onMouseClick(Lnet/minecraft/screen/slot/Slot;IILnet/minecraft/screen/slot/SlotActionType;)V", at = @At("HEAD"), cancellable = true)
+	public void onMouseClickedSlot(Slot slot, int slotId, int button, SlotActionType actionType, CallbackInfo ci) {
+		if (slotId == -999 && getScreenHandler() != null && actionType == SlotActionType.PICKUP) { // -999 is code for "clicked outside the main window"
+			ItemStack cursorStack = getScreenHandler().getCursorStack();
+			if (cursorStack != null && IsSlotProtectedEvent.shouldBlockInteraction(slot, SlotActionType.THROW, cursorStack)) {
+				ci.cancel();
+				return;
+			}
+		}
+		if (IsSlotProtectedEvent.shouldBlockInteraction(slot, actionType)) {
+			ci.cancel();
+			return;
+		}
+		if (actionType == SlotActionType.SWAP && 0 <= button && button < 9) {
+			if (IsSlotProtectedEvent.shouldBlockInteraction(new Slot(playerInventory, button, 0, 0), actionType)) {
+				ci.cancel();
+			}
+		}
+	}
+
+
+	@Inject(method = "render", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;drawSlot(Lnet/minecraft/client/gui/DrawContext;Lnet/minecraft/screen/slot/Slot;)V", shift = At.Shift.AFTER), locals = LocalCapture.CAPTURE_FAILHARD)
+	public void onAfterDrawSlot(DrawContext context, int mouseX, int mouseY, float delta, CallbackInfo ci, int i, int j, int k, Slot slot) {
+		SlotRenderEvents.After event = new SlotRenderEvents.After(context, slot, mouseX, mouseY, delta);
+		SlotRenderEvents.After.Companion.publish(event);
+	}
+
+	@Inject(method = "render", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/gui/screen/ingame/HandledScreen;drawSlot(Lnet/minecraft/client/gui/DrawContext;Lnet/minecraft/screen/slot/Slot;)V", shift = At.Shift.BEFORE), locals = LocalCapture.CAPTURE_FAILHARD)
+	public void onBeforeDrawSlot(DrawContext context, int mouseX, int mouseY, float delta, CallbackInfo ci, int i, int j, int k, Slot slot) {
+		SlotRenderEvents.Before event = new SlotRenderEvents.Before(context, slot, mouseX, mouseY, delta);
+		SlotRenderEvents.Before.Companion.publish(event);
+	}
 }