From 86bae2c0e9492e8827eb5be335c0a3615f6f8e59 Mon Sep 17 00:00:00 2001
From: Technici4n <13494793+Technici4n@users.noreply.github.com>
Date: Thu, 9 Sep 2021 19:47:06 +0200
Subject: [PATCH] Work around vanilla capturing ItemStack references (#1700)

---
 .../v1/item/base/SingleStackStorage.java      |  6 ++++-
 .../transfer/item/InventorySlotWrapper.java   | 22 ++++++++++++++++
 .../test/transfer/unittests/ItemTests.java    | 25 +++++++++++++++++++
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/item/base/SingleStackStorage.java b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/item/base/SingleStackStorage.java
index e29936a4d..713adde58 100644
--- a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/item/base/SingleStackStorage.java
+++ b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/item/base/SingleStackStorage.java
@@ -115,6 +115,7 @@ public abstract class SingleStackStorage extends SnapshotParticipant<ItemStack>
 
 			if (insertedAmount > 0) {
 				updateSnapshots(transaction);
+				currentStack = getStack();
 
 				if (currentStack.isEmpty()) {
 					currentStack = insertedVariant.toStack(insertedAmount);
@@ -142,6 +143,7 @@ public abstract class SingleStackStorage extends SnapshotParticipant<ItemStack>
 
 			if (extracted > 0) {
 				this.updateSnapshots(transaction);
+				currentStack = getStack();
 				currentStack.decrement(extracted);
 				setStack(currentStack);
 			}
@@ -154,7 +156,9 @@ public abstract class SingleStackStorage extends SnapshotParticipant<ItemStack>
 
 	@Override
 	protected final ItemStack createSnapshot() {
-		return getStack().copy();
+		ItemStack original = getStack();
+		setStack(original.copy());
+		return original;
 	}
 
 	@Override
diff --git a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/InventorySlotWrapper.java b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/InventorySlotWrapper.java
index 2d46f0b98..64564684e 100644
--- a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/InventorySlotWrapper.java
+++ b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/InventorySlotWrapper.java
@@ -34,6 +34,7 @@ class InventorySlotWrapper extends SingleStackStorage {
 	 */
 	private final InventoryStorageImpl storage;
 	final int slot;
+	private ItemStack lastReleasedSnapshot = null;
 
 	InventorySlotWrapper(InventoryStorageImpl storage, int slot) {
 		this.storage = storage;
@@ -66,4 +67,25 @@ class InventorySlotWrapper extends SingleStackStorage {
 		storage.markDirtyParticipant.updateSnapshots(transaction);
 		super.updateSnapshots(transaction);
 	}
+
+	@Override
+	protected void releaseSnapshot(ItemStack snapshot) {
+		lastReleasedSnapshot = snapshot;
+	}
+
+	@Override
+	protected void onFinalCommit() {
+		// Try to apply the change to the original stack
+		ItemStack original = lastReleasedSnapshot;
+		ItemStack currentStack = getStack();
+
+		if (!original.isEmpty() && !currentStack.isEmpty() && ItemStack.canCombine(original, currentStack)) {
+			// None is empty and the contents match: just update the amount and reuse the original stack.
+			original.setCount(currentStack.getCount());
+			setStack(original);
+		} else {
+			// Otherwise assume everything was taken from original so empty it.
+			original.setCount(0);
+		}
+	}
 }
diff --git a/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/ItemTests.java b/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/ItemTests.java
index 191ccea66..fb4af6c67 100644
--- a/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/ItemTests.java
+++ b/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/ItemTests.java
@@ -40,11 +40,36 @@ import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
  */
 class ItemTests {
 	public static void run() {
+		testStackReference();
 		testInventoryWrappers();
 		testLimitedStackCountInventory();
 		testLimitedStackCountItem();
 	}
 
+	private static void testStackReference() {
+		// Ensure that Inventory wrappers will try to mutate the backing stack as much as possible.
+		// In many cases, MC code captures a reference to the ItemStack so we want to edit that stack directly
+		// and not a copy whenever we can. Obviously this can't be perfect, but we try to cover as many cases as possible.
+		SimpleInventory inv = new SimpleInventory(new ItemStack(Items.DIAMOND, 2));
+		InventoryStorage invWrapper = InventoryStorage.of(inv, null);
+		ItemStack stack = inv.getStack(0);
+
+		// Simulate should correctly reset the stack.
+		try (Transaction tx = Transaction.openOuter()) {
+			invWrapper.extract(ItemVariant.of(Items.DIAMOND), 2, tx);
+		}
+
+		if (stack != inv.getStack(0)) throw new AssertionError("Stack should have stayed the same.");
+
+		// Commit should try to edit the original stack when it is feasible to do so.
+		try (Transaction tx = Transaction.openOuter()) {
+			invWrapper.extract(ItemVariant.of(Items.DIAMOND), 1, tx);
+			tx.commit();
+		}
+
+		if (stack != inv.getStack(0)) throw new AssertionError("Stack should have stayed the same.");
+	}
+
 	private static void testInventoryWrappers() {
 		ItemVariant emptyBucket = ItemVariant.of(Items.BUCKET);
 		TestSidedInventory testInventory = new TestSidedInventory();