Fix #2164: Transfer API performance issues (#2189)

* Fix #2164: Transfer API performance issue

* Cache ItemStack in ItemVariantImpl

* Fix checkstyle

* Fix #2201

* Update comment in getCachedStack()

(cherry picked from commit f4563ac81e)

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/item/base/SingleStackStorage.java

@@ -85,7 +85,7 @@ public abstract class SingleStackStorage extends SnapshotParticipant<ItemStack>
 
 	@Override
 	public boolean isResourceBlank() {
-		return getResource().isBlank();
+		return getStack().isEmpty();
 	}
 
 	@Override
@@ -123,9 +123,9 @@ public abstract class SingleStackStorage extends SnapshotParticipant<ItemStack>
 				}
 
 				setStack(currentStack);
-			}
 
-			return insertedAmount;
+				return insertedAmount;
+			}
 		}
 
 		return 0;
@@ -145,9 +145,9 @@ public abstract class SingleStackStorage extends SnapshotParticipant<ItemStack>
 				currentStack = getStack();
 				currentStack.decrement(extracted);
 				setStack(currentStack);
-			}
 
-			return extracted;
+				return extracted;
+			}
 		}
 
 		return 0;

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/storage/base/SingleVariantStorage.java

@@ -81,9 +81,9 @@ public abstract class SingleVariantStorage<T extends TransferVariant<?>> extends
 				} else {
 					amount += insertedAmount;
 				}
-			}
 
-			return insertedAmount;
+				return insertedAmount;
+			}
 		}
 
 		return 0;
@@ -103,9 +103,9 @@ public abstract class SingleVariantStorage<T extends TransferVariant<?>> extends
 				if (amount == 0) {
 					variant = getBlankVariant();
 				}
-			}
 
-			return extractedAmount;
+				return extractedAmount;
+			}
 		}
 
 		return 0;

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/InventorySlotWrapper.java

@@ -21,7 +21,6 @@ import net.minecraft.item.ItemStack;
 import net.fabricmc.fabric.api.transfer.v1.item.base.SingleStackStorage;
 import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
 import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
-import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
 
 /**
  * A wrapper around a single slot of an inventory.
@@ -35,11 +34,13 @@ class InventorySlotWrapper extends SingleStackStorage {
 	 */
 	private final InventoryStorageImpl storage;
 	final int slot;
+	private final SpecialLogicInventory specialInv;
 	private ItemStack lastReleasedSnapshot = null;
 
 	InventorySlotWrapper(InventoryStorageImpl storage, int slot) {
 		this.storage = storage;
 		this.slot = slot;
+		this.specialInv = storage.inventory instanceof SpecialLogicInventory specialInv ? specialInv : null;
 	}
 
 	@Override
@@ -49,18 +50,26 @@ class InventorySlotWrapper extends SingleStackStorage {
 
 	@Override
 	protected void setStack(ItemStack stack) {
-		TransferApiImpl.SUPPRESS_SPECIAL_LOGIC.set(Boolean.TRUE);
-
-		try {
+		if (specialInv == null) {
 			storage.inventory.setStack(slot, stack);
-		} finally {
-			TransferApiImpl.SUPPRESS_SPECIAL_LOGIC.remove();
+		} else {
+			specialInv.fabric_setSuppress(true);
+
+			try {
+				storage.inventory.setStack(slot, stack);
+			} finally {
+				specialInv.fabric_setSuppress(false);
+			}
 		}
 	}
 
 	@Override
-	protected boolean canInsert(ItemVariant itemVariant) {
-		return storage.inventory.isValid(slot, itemVariant.toStack());
+	public long insert(ItemVariant insertedVariant, long maxAmount, TransactionContext transaction) {
+		if (!storage.inventory.isValid(slot, ((ItemVariantImpl) insertedVariant).getCachedStack())) {
+			return 0;
+		} else {
+			return super.insert(insertedVariant, maxAmount, transaction);
+		}
 	}
 
 	@Override

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/ItemVariantImpl.java

@@ -23,6 +23,7 @@ import org.slf4j.Logger;
 import org.jetbrains.annotations.Nullable;
 
 import net.minecraft.item.Item;
+import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.network.PacketByteBuf;
@@ -48,6 +49,10 @@ public class ItemVariantImpl implements ItemVariant {
 	private final Item item;
 	private final @Nullable NbtCompound nbt;
 	private final int hashCode;
+	/**
+	 * Lazily computed, equivalent to calling toStack(1). <b>MAKE SURE IT IS NEVER MODIFIED!</b>
+	 */
+	private volatile @Nullable ItemStack cachedStack = null;
 
 	public ItemVariantImpl(Item item, NbtCompound nbt) {
 		this.item = item;
@@ -135,4 +140,15 @@ public class ItemVariantImpl implements ItemVariant {
 	public int hashCode() {
 		return hashCode;
 	}
+
+	public ItemStack getCachedStack() {
+		ItemStack ret = cachedStack;
+
+		if (ret == null) {
+			// multiple stacks could be created at the same time by different threads, but that is not an issue
+			cachedStack = ret = toStack();
+		}
+
+		return ret;
+	}
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/SidedInventorySlotWrapper.java

@@ -40,7 +40,7 @@ class SidedInventorySlotWrapper implements SingleSlotStorage<ItemVariant> {
 
 	@Override
 	public long insert(ItemVariant resource, long maxAmount, TransactionContext transaction) {
-		if (!sidedInventory.canInsert(slotWrapper.slot, resource.toStack(), direction)) {
+		if (!sidedInventory.canInsert(slotWrapper.slot, ((ItemVariantImpl) resource).getCachedStack(), direction)) {
 			return 0;
 		} else {
 			return slotWrapper.insert(resource, maxAmount, transaction);
@@ -49,7 +49,7 @@ class SidedInventorySlotWrapper implements SingleSlotStorage<ItemVariant> {
 
 	@Override
 	public long extract(ItemVariant resource, long maxAmount, TransactionContext transaction) {
-		if (!sidedInventory.canExtract(slotWrapper.slot, resource.toStack(), direction)) {
+		if (!sidedInventory.canExtract(slotWrapper.slot, ((ItemVariantImpl) resource).getCachedStack(), direction)) {
 			return 0;
 		} else {
 			return slotWrapper.extract(resource, maxAmount, transaction);

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/SpecialLogicInventory.java

@@ -22,9 +22,13 @@ import net.minecraft.item.ItemStack;
 
 /**
  * Internal class that allows inventory instances to defer special logic until {@link InventorySlotWrapper#onFinalCommit()} is called.
- * Special logic should be suppressed when {@link net.fabricmc.fabric.impl.transfer.TransferApiImpl#SUPPRESS_SPECIAL_LOGIC} is true.
  */
 @ApiStatus.Internal
 public interface SpecialLogicInventory {
+	/**
+	 * Decide whether special logic should now be suppressed. If true, must remain suppressed until the next call.
+	 */
+	void fabric_setSuppress(boolean suppress);
+
 	void fabric_onFinalCommit(int slot, ItemStack oldStack, ItemStack newStack);
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/mixin/transfer/AbstractFurnaceBlockEntityMixin.java

@@ -18,6 +18,7 @@ package net.fabricmc.fabric.mixin.transfer;
 
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
@@ -31,7 +32,6 @@ import net.minecraft.util.collection.DefaultedList;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.world.World;
 
-import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
 import net.fabricmc.fabric.impl.transfer.item.SpecialLogicInventory;
 
 /**
@@ -45,6 +45,8 @@ public abstract class AbstractFurnaceBlockEntityMixin extends LockableContainerB
 	int cookTime;
 	@Shadow
 	int cookTimeTotal;
+	@Unique
+	private boolean fabric_suppressSpecialLogic = false;
 
 	protected AbstractFurnaceBlockEntityMixin(BlockEntityType<?> blockEntityType, BlockPos blockPos, BlockState blockState) {
 		super(blockEntityType, blockPos, blockState);
@@ -53,12 +55,17 @@ public abstract class AbstractFurnaceBlockEntityMixin extends LockableContainerB
 
 	@Inject(at = @At("HEAD"), method = "setStack", cancellable = true)
 	public void setStackSuppressUpdate(int slot, ItemStack stack, CallbackInfo ci) {
-		if (TransferApiImpl.SUPPRESS_SPECIAL_LOGIC.get() != null) {
+		if (fabric_suppressSpecialLogic) {
 			inventory.set(slot, stack);
 			ci.cancel();
 		}
 	}
 
+	@Override
+	public void fabric_setSuppress(boolean suppress) {
+		fabric_suppressSpecialLogic = suppress;
+	}
+
 	@Override
 	public void fabric_onFinalCommit(int slot, ItemStack oldStack, ItemStack newStack) {
 		if (slot == 0) {

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/mixin/transfer/LootableContainerBlockEntityMixin.java

@@ -17,25 +17,39 @@
 package net.fabricmc.fabric.mixin.transfer;
 
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Redirect;
 
 import net.minecraft.block.entity.LootableContainerBlockEntity;
+import net.minecraft.item.ItemStack;
 
-import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
+import net.fabricmc.fabric.impl.transfer.item.SpecialLogicInventory;
 
 /**
  * Defer markDirty until the outer transaction close callback when setStack is called from an inventory wrapper.
  */
 @Mixin(LootableContainerBlockEntity.class)
-public class LootableContainerBlockEntityMixin {
+public class LootableContainerBlockEntityMixin implements SpecialLogicInventory {
+	@Unique
+	private boolean fabric_suppressSpecialLogic = false;
+
 	@Redirect(
 			at = @At(value = "INVOKE", target = "Lnet/minecraft/block/entity/LootableContainerBlockEntity;markDirty()V"),
 			method = "setStack(ILnet/minecraft/item/ItemStack;)V"
 	)
 	public void fabric_redirectMarkDirty(LootableContainerBlockEntity self) {
-		if (TransferApiImpl.SUPPRESS_SPECIAL_LOGIC.get() == null) {
+		if (!fabric_suppressSpecialLogic) {
 			self.markDirty();
 		}
 	}
+
+	@Override
+	public void fabric_setSuppress(boolean suppress) {
+		fabric_suppressSpecialLogic = suppress;
+	}
+
+	@Override
+	public void fabric_onFinalCommit(int slot, ItemStack oldStack, ItemStack newStack) {
+	}
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/mixin/transfer/SimpleInventoryMixin.java

@@ -17,25 +17,39 @@
 package net.fabricmc.fabric.mixin.transfer;
 
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Redirect;
 
 import net.minecraft.inventory.SimpleInventory;
+import net.minecraft.item.ItemStack;
 
-import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
+import net.fabricmc.fabric.impl.transfer.item.SpecialLogicInventory;
 
 /**
  * Defer markDirty until the outer transaction close callback when setStack is called from an inventory wrapper.
  */
 @Mixin(SimpleInventory.class)
-public class SimpleInventoryMixin {
+public class SimpleInventoryMixin implements SpecialLogicInventory {
+	@Unique
+	private boolean fabric_suppressSpecialLogic = false;
+
 	@Redirect(
 			at = @At(value = "INVOKE", target = "Lnet/minecraft/inventory/SimpleInventory;markDirty()V"),
 			method = "setStack(ILnet/minecraft/item/ItemStack;)V"
 	)
 	public void fabric_redirectMarkDirty(SimpleInventory self) {
-		if (TransferApiImpl.SUPPRESS_SPECIAL_LOGIC.get() == null) {
+		if (!fabric_suppressSpecialLogic) {
 			self.markDirty();
 		}
 	}
+
+	@Override
+	public void fabric_setSuppress(boolean suppress) {
+		fabric_suppressSpecialLogic = suppress;
+	}
+
+	@Override
+	public void fabric_onFinalCommit(int slot, ItemStack oldStack, ItemStack newStack) {
+	}
 }

fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/ItemTests.java

@@ -16,6 +16,8 @@
 
 package net.fabricmc.fabric.test.transfer.unittests;
 
+import static net.fabricmc.fabric.test.transfer.unittests.TestUtil.assertEquals;
+
 import java.util.stream.IntStream;
 
 import org.jetbrains.annotations.Nullable;
@@ -127,6 +129,16 @@ class ItemTests {
 		if (!testInventory.getStack(1).isOf(Items.BUCKET) || testInventory.getStack(1).getCount() != 1) throw new AssertionError("Slot 1 should have been a bucket.");
 
 		checkComparatorOutput(testInventory);
+
+		// Check that we return sensible results if amount stored > capacity
+		ItemStack oversizedStack = new ItemStack(Items.DIAMOND_PICKAXE, 2);
+		SimpleInventory simpleInventory = new SimpleInventory(oversizedStack);
+		InventoryStorage wrapper = InventoryStorage.of(simpleInventory, null);
+
+		try (Transaction transaction = Transaction.openOuter()) {
+			assertEquals(0L, wrapper.insert(ItemVariant.of(oversizedStack), 10, transaction));
+			transaction.commit();
+		}
 	}
 
 	private static boolean stackEquals(ItemStack stack, Item item, int count) {