From 33a20007059c95991fae7fb15fbed749fd590505 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Thu, 19 Aug 2021 00:42:35 +0700
Subject: [PATCH] Make setChanged be more robust

---
 .../dbotthepony/mc/otm/menu/MatteryMenu.java  | 23 +++++++++++++++++++
 .../mc/otm/menu/slot/MatterySlot.java         | 16 ++++++++++---
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/main/java/ru/dbotthepony/mc/otm/menu/MatteryMenu.java b/src/main/java/ru/dbotthepony/mc/otm/menu/MatteryMenu.java
index f502386b8..f04ab8bc4 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/menu/MatteryMenu.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/menu/MatteryMenu.java
@@ -9,6 +9,7 @@ import net.minecraft.world.inventory.MenuType;
 import net.minecraft.world.inventory.Slot;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.level.block.entity.BlockEntity;
+import ru.dbotthepony.mc.otm.menu.slot.MatterySlot;
 import ru.dbotthepony.mc.otm.menu.widget.AbstractWidget;
 
 import javax.annotation.Nullable;
@@ -106,6 +107,22 @@ public abstract class MatteryMenu extends AbstractContainerMenu {
 	abstract protected int getWorkingSlotStart();
 	abstract protected int getWorkingSlotEnd();
 
+	protected void notifySlotsWatchCurrentStack() {
+		for (var slot : slots) {
+			if (slot instanceof MatterySlot slot1) {
+				slot1.watchCurrentStack();
+			}
+		}
+	}
+
+	protected void notifySlotsUnwatchCurrentStack() {
+		for (var slot : slots) {
+			if (slot instanceof MatterySlot slot1) {
+				slot1.unwatchCurrentStack();
+			}
+		}
+	}
+
 	// This method receive Player interactor and slot_index where Shift + Right click occurred
 	// It shall return item stack that got moved
 	@Override
@@ -126,6 +143,7 @@ public abstract class MatteryMenu extends AbstractContainerMenu {
 
 		ItemStack moved = ItemStack.EMPTY;
 		Slot get_slot = this.slots.get(slot_index);
+		notifySlotsWatchCurrentStack();
 
 		if (get_slot.hasItem()) {
 			ItemStack slot_item = get_slot.getItem();
@@ -135,10 +153,12 @@ public abstract class MatteryMenu extends AbstractContainerMenu {
 				// Moving FROM machine TO inventory
 
 				if (!moveItemStackTo(slot_item, inventory_slot_index_start, inventory_slot_index_end + 1, false)) {
+					notifySlotsUnwatchCurrentStack();
 					return ItemStack.EMPTY;
 				}
 			} else if (!moveItemStackTo(slot_item, start, end, false)) {
 				// Moving FROM inventory TO machine
+				notifySlotsUnwatchCurrentStack();
 				return ItemStack.EMPTY;
 			}
 
@@ -147,8 +167,11 @@ public abstract class MatteryMenu extends AbstractContainerMenu {
 			} else {
 				get_slot.setChanged();
 			}
+
+			notifySlotsWatchCurrentStack();
 		}
 
+		notifySlotsUnwatchCurrentStack();
 		return moved;
 	}
 }
diff --git a/src/main/java/ru/dbotthepony/mc/otm/menu/slot/MatterySlot.java b/src/main/java/ru/dbotthepony/mc/otm/menu/slot/MatterySlot.java
index 01bd6f397..1dd536b6c 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/menu/slot/MatterySlot.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/menu/slot/MatterySlot.java
@@ -4,6 +4,7 @@ import net.minecraft.world.Container;
 import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.inventory.Slot;
 import net.minecraft.world.item.ItemStack;
+import ru.dbotthepony.mc.otm.OverdriveThatMatters;
 import ru.dbotthepony.mc.otm.container.MatteryContainer;
 
 // Just a sign that this slot needs auto draw
@@ -20,6 +21,15 @@ public class MatterySlot extends Slot {
 	}
 
 	private ItemStack last_stack;
+	private ItemStack watching_stack;
+
+	public void watchCurrentStack() {
+		watching_stack = super.getItem().copy();
+	}
+
+	public void unwatchCurrentStack() {
+		watching_stack = null;
+	}
 
 	@Override
 	public ItemStack getItem() {
@@ -31,7 +41,7 @@ public class MatterySlot extends Slot {
 	@Override
 	public void setChanged() {
 		if (container instanceof MatteryContainer container1) {
-			ItemStack old = last_stack.copy();
+			ItemStack old = watching_stack != null ? watching_stack.copy() : last_stack.copy();
 			container1.setChanged(getSlotIndex(), super.getItem(), old);
 		} else {
 			super.setChanged();
@@ -41,7 +51,7 @@ public class MatterySlot extends Slot {
 	@Override
 	public void set(ItemStack p_40240_) {
 		if (container instanceof MatteryContainer container1) {
-			ItemStack old = container1.getItem(getSlotIndex()).copy();
+			ItemStack old = watching_stack != null ? watching_stack.copy() : container1.getItem(getSlotIndex()).copy();
 			container1.startIgnore();
 			container1.setItem(getSlotIndex(), p_40240_);
 			container1.stopIgnore();
@@ -54,7 +64,7 @@ public class MatterySlot extends Slot {
 	@Override
 	public void onTake(Player p_150645_, ItemStack p_150646_) {
 		if (container instanceof MatteryContainer container1) {
-			ItemStack old = last_stack.copy();
+			ItemStack old = watching_stack != null ? watching_stack.copy() : last_stack.copy();
 			container1.setChanged(getSlotIndex(), p_150646_.copy(), old);
 			last_stack = ItemStack.EMPTY;
 		} else {