Fix #9

AlmostReliable · Feb 4, 2024 · 740675e · 740675e
1 parent ecc7413
commit 740675e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 4 deletions.
diff --git a/Common/src/main/java/com/almostreliable/morejs/Debug.java b/Common/src/main/java/com/almostreliable/morejs/Debug.java
@@ -1,5 +1,5 @@
 package com.almostreliable.morejs;
 
 public class Debug {
-    public static final boolean ENCHANTMENT = true;
+    public static final boolean ENCHANTMENT = false;
 }
diff --git a/.../src/main/java/com/almostreliable/morejs/features/enchantment/EnchantmentMenuProcess.java b/.../src/main/java/com/almostreliable/morejs/features/enchantment/EnchantmentMenuProcess.java
@@ -49,7 +49,10 @@ public void setCurrentItem(ItemStack currentItem) {
     }
 
     public void clearEnchantments() {
-        if (Debug.ENCHANTMENT) MoreJS.LOG.warn("<{}> Clearing enchantments", player);
+        if (Debug.ENCHANTMENT) {
+            MoreJS.LOG.warn("");
+            MoreJS.LOG.warn("<{}> Clearing enchantments", player);
+        }
         enchantments.clear();
     }
 
@@ -67,6 +70,7 @@ public void setEnchantments(int index, List<EnchantmentInstance> enchantments) {
             var s = formatEnchantments(enchantments);
             MoreJS.LOG.info("<{}> Setting enchantments for index {} [{}]", player, index, s);
         }
+
         this.enchantments.put(index, new ArrayList<>(enchantments));
     }
 
@@ -114,4 +118,8 @@ public void abortEvent(ItemStack item) {
         clearEnchantments();
         setState(EnchantmentState.IDLE);
     }
+
+    public ItemStack getCurrentItem() {
+        return currentItem;
+    }
 }
diff --git a/Common/src/main/java/com/almostreliable/morejs/mixin/enchanting/EnchantmentMenuMixin.java b/Common/src/main/java/com/almostreliable/morejs/mixin/enchanting/EnchantmentMenuMixin.java
@@ -27,7 +27,8 @@
 import javax.annotation.Nullable;
 import java.util.List;
 
-@Mixin(value = EnchantmentMenu.class, priority = 42) // lower priority to ensure we run first. 42 is the answer to everything.
+// lower priority to ensure we run first. 42 is the answer to everything.
+@Mixin(value = EnchantmentMenu.class, priority = 42)
 public abstract class EnchantmentMenuMixin extends AbstractContainerMenu implements EnchantmentMenuExtension {
     @Unique
     private EnchantmentMenuProcess morejs$process;
@@ -141,7 +142,7 @@ private void handleEnchantmentGetter(ItemStack itemStack, int index, int powerLe
                 cir.setReturnValue(false);
             }
 
-            if(e.itemWasChanged()) {
+            if (e.itemWasChanged()) {
                 cir.setReturnValue(false);
                 ItemStack newItem = e.getItem().copy();
                 this.morejs$process.abortEvent(newItem);
@@ -150,6 +151,17 @@ private void handleEnchantmentGetter(ItemStack itemStack, int index, int powerLe
         });
     }
 
+    @Inject(method = "clickMenuButton", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/inventory/ContainerLevelAccess;execute(Ljava/util/function/BiConsumer;)V", shift = At.Shift.AFTER))
+    private void clickMenuButton$postClear(Player player, int i, CallbackInfoReturnable<Boolean> cir) {
+        this.access.execute((level, pos) -> {
+            ItemStack currentItem = this.morejs$process.getCurrentItem();
+            ItemStack secondItem = this.enchantSlots.getItem(0);
+            if (!ItemStack.matches(currentItem, secondItem)) {
+                this.morejs$process.abortEvent(ItemStack.EMPTY);
+            }
+        });
+    }
+
     @Override
     public void broadcastChanges() {
         if (this.morejs$process.isFreezeBroadcast()) {