Always generate buttons and use as reference for slot collection

Author: NotRyken

common/src/main/java/dev/terminalmc/clientsort/client/gui/TriggerButtonManager.java

@@ -55,7 +55,7 @@ public abstract class TriggerButton extends Button {
 
     private final AbstractContainerScreen<?> screen;
     public final Container container;
-    private final Slot referenceSlot;
+    public final Slot referenceSlot;
     public final boolean isPlayerInv;
 
     private final WidgetSprites sprites;

common/src/main/java/dev/terminalmc/clientsort/client/gui/widget/TriggerButton.java

@@ -20,6 +20,7 @@
 import dev.terminalmc.clientsort.client.ClientSort;
 import dev.terminalmc.clientsort.client.compat.itemlocks.ItemLocksWrapper;
 import dev.terminalmc.clientsort.client.config.ClassPolicy;
+import dev.terminalmc.clientsort.client.gui.TriggerButtonManager;
 import dev.terminalmc.clientsort.client.inventory.Scope;
 import dev.terminalmc.clientsort.client.inventory.helper.ContainerScreenHelper;
 import dev.terminalmc.clientsort.client.inventory.operator.client.ClientCreativeOperator;
@@ -42,6 +43,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 
 import static dev.terminalmc.clientsort.ClientSort.debug;
@@ -62,6 +64,7 @@ public abstract class SingleUseOperator {
     private boolean hasOperated = false;
     protected final AbstractContainerScreen<?> screen;
     protected final ContainerScreenHelper<? extends AbstractContainerScreen<?>> screenHelper;
+    protected final Operation op;
     /**
      * The slot that was hovered when sorting was triggered.
      */
@@ -94,13 +97,15 @@ public abstract class SingleUseOperator {
      */
     protected final ItemStack[] otherScopeStacks;
 
-    protected SingleUseOperator(AbstractContainerScreen<?> screen, Slot originSlot) {
+    protected SingleUseOperator(AbstractContainerScreen<?> screen, Slot originSlot, Operation op) {
         this.screen = screen;
         this.screenHelper = ContainerScreenHelper.of(screen);
         this.originSlot = originSlot;
+        this.op = op;
+
         // Collect slots in origin scope
         Scope originScope = screenHelper.getScope(originSlot);
-        originScopeSlots = collectSlots(originScope);
+        originScopeSlots = collectSlots(originSlot);
         if (debug()) {
             ClientSort.LOG.info(
                     "Discovered {} slots in origin scope ({} - {}):",
@@ -116,13 +121,20 @@ protected SingleUseOperator(AbstractContainerScreen<?> screen, Slot originSlot)
             originScopeStacks[i] = originScopeSlots[i].getItem().copy();
         }
 
+        // Get the other container, if any
+        LocalPlayer player = Objects.requireNonNull(Minecraft.getInstance().player);
+        @Nullable Slot otherSlot = originSlot.container == player.getInventory()
+                ? TriggerButtonManager.getContainerRefSlot(op)
+                : TriggerButtonManager.getPlayerRefSlot(op);
+        if (otherSlot == null || op.equals(Operation.SORT)) {
+            otherScopeSlots = new Slot[]{};
+            otherScopeStacks = new ItemStack[]{};
+            return;
+        }
+
         // Collect slots in other container scope, if any
-        Scope otherScope = switch (originScope) {
-            case PLAYER_INV -> Scope.CONTAINER_INV;
-            case CONTAINER_INV -> Scope.PLAYER_INV;
-            default -> Scope.INVALID;
-        };
-        otherScopeSlots = collectSlots(otherScope);
+        Scope otherScope = screenHelper.getScope(otherSlot);
+        otherScopeSlots = collectSlots(otherSlot);
         if (debug()) {
             ClientSort.LOG.info(
                     "Discovered {} slots in other scope ({} - {}):",
@@ -142,8 +154,9 @@ protected SingleUseOperator(AbstractContainerScreen<?> screen, Slot originSlot)
     /**
      * Finds all the valid inventory menu slots that are in {@code scope}.
      */
-    private Slot[] collectSlots(Scope scope) {
+    private Slot[] collectSlots(Slot refSlot) {
         LocalPlayer player = Minecraft.getInstance().player;
+        Scope scope = screenHelper.getScope(refSlot);
         if (scope == Scope.INVALID)
             return new Slot[0];
 
@@ -153,6 +166,12 @@ private Slot[] collectSlots(Scope scope) {
             int slotId = ((ISlot) slot).clientsort$getIndexInMenu();
             int slotIdx = ((ISlot) slot).clientsort$getIndexInContainer();
 
+            //noinspection ConstantValue
+            if (slot.container == null)
+                continue;
+            if (slot.container != refSlot.container)
+                continue;
+
             // Ignore slots in different scope
             if (screenHelper.getScope(slot) != scope)
                 continue;
@@ -295,7 +314,7 @@ public static boolean transfer(
                 && ClientServices.PLATFORM.canSendToServer(operation.type)) {
             if (debug())
                 ClientSort.LOG.info("Preparing server operator for {}", operation.name());
-            return new ServerOperator(screen, originSlot);
+            return new ServerOperator(screen, originSlot, operation);
         }
 
         // Check that there is not already an op running
@@ -313,11 +332,11 @@ public static boolean transfer(
                 && screen instanceof CreativeModeInventoryScreen) {
             if (debug())
                 ClientSort.LOG.info("Preparing client-creative operator for {}", operation.name());
-            return new ClientCreativeOperator(screen, originSlot);
+            return new ClientCreativeOperator(screen, originSlot, operation);
         } else {
             if (debug())
                 ClientSort.LOG.info("Preparing client-survival operator for {}", operation.name());
-            return new ClientSurvivalOperator(screen, originSlot);
+            return new ClientSurvivalOperator(screen, originSlot, operation);
         }
     }
 

common/src/main/java/dev/terminalmc/clientsort/client/inventory/operator/SingleUseOperator.java

@@ -18,6 +18,7 @@
 
 import dev.terminalmc.clientsort.ClientSort;
 import dev.terminalmc.clientsort.client.interaction.InteractionManager;
+import dev.terminalmc.clientsort.client.inventory.operator.Operation;
 import dev.terminalmc.clientsort.client.util.SoundManager;
 import dev.terminalmc.clientsort.util.inject.ISlot;
 import net.minecraft.client.Minecraft;
@@ -36,8 +37,12 @@
  */
 public class ClientCreativeOperator extends ClientOperator {
 
-    public ClientCreativeOperator(AbstractContainerScreen<?> screen, Slot originSlot) {
-        super(screen, originSlot);
+    public ClientCreativeOperator(
+            AbstractContainerScreen<?> screen,
+            Slot originSlot,
+            Operation op
+    ) {
+        super(screen, originSlot, op);
     }
 
     /**

common/src/main/java/dev/terminalmc/clientsort/client/inventory/operator/client/ClientCreativeOperator.java

@@ -18,6 +18,7 @@
 
 import dev.terminalmc.clientsort.client.ClientSort;
 import dev.terminalmc.clientsort.client.interaction.InteractionManager;
+import dev.terminalmc.clientsort.client.inventory.operator.Operation;
 import dev.terminalmc.clientsort.client.inventory.operator.SingleUseOperator;
 import dev.terminalmc.clientsort.client.order.SortContext;
 import dev.terminalmc.clientsort.client.order.SortOrder;
@@ -32,8 +33,8 @@
  */
 public abstract class ClientOperator extends SingleUseOperator {
 
-    public ClientOperator(AbstractContainerScreen<?> screen, Slot originSlot) {
-        super(screen, originSlot);
+    public ClientOperator(AbstractContainerScreen<?> screen, Slot originSlot, Operation op) {
+        super(screen, originSlot, op);
     }
 
     /**

common/src/main/java/dev/terminalmc/clientsort/client/inventory/operator/client/ClientOperator.java

@@ -20,6 +20,7 @@
 import dev.terminalmc.clientsort.ClientSort;
 import dev.terminalmc.clientsort.client.interaction.InteractionManager;
 import dev.terminalmc.clientsort.client.interaction.InteractionManager.InteractionEvent;
+import dev.terminalmc.clientsort.client.inventory.operator.Operation;
 import dev.terminalmc.clientsort.client.util.SoundManager;
 import dev.terminalmc.clientsort.mixin.client.accessor.AbstractContainerScreenAccessor;
 import dev.terminalmc.clientsort.util.inject.ISlot;
@@ -44,8 +45,12 @@
  */
 public class ClientSurvivalOperator extends ClientOperator {
 
-    public ClientSurvivalOperator(AbstractContainerScreen<?> screen, Slot originSlot) {
-        super(screen, originSlot);
+    public ClientSurvivalOperator(
+            AbstractContainerScreen<?> screen,
+            Slot originSlot,
+            Operation op
+    ) {
+        super(screen, originSlot, op);
     }
 
     /**

common/src/main/java/dev/terminalmc/clientsort/client/inventory/operator/client/ClientSurvivalOperator.java

@@ -17,6 +17,7 @@
 package dev.terminalmc.clientsort.client.inventory.operator.server;
 
 import dev.terminalmc.clientsort.client.ClientSort;
+import dev.terminalmc.clientsort.client.inventory.operator.Operation;
 import dev.terminalmc.clientsort.client.inventory.operator.SingleUseOperator;
 import dev.terminalmc.clientsort.client.network.handler.CollectResultHandler;
 import dev.terminalmc.clientsort.client.network.handler.SortResultHandler;
@@ -49,8 +50,8 @@
  */
 public class ServerOperator extends SingleUseOperator {
 
-    public ServerOperator(AbstractContainerScreen<?> screen, Slot originSlot) {
-        super(screen, originSlot);
+    public ServerOperator(AbstractContainerScreen<?> screen, Slot originSlot, Operation op) {
+        super(screen, originSlot, op);
     }
 
     @Override
@@ -76,7 +77,7 @@ protected void sort(SortOrder sortOrder) {
                     }
                 };
 
-                ServerOperator sorter = new ServerOperator(screen, originSlot);
+                ServerOperator sorter = new ServerOperator(screen, originSlot, op);
                 int[] slotMapping = sorter.createSlotMapping(sortOrder);
                 if (debug())
                     ClientSort.LOG.info("Sending payload for operation SORT");