Transfer API: bugfixes and improvements (#1762)

* Fix SingleVariantItemStorage extraction bug

* Inventory wrappers: reuse the original stack more aggressively

* Add some filtered overloads to StorageUtil methods

* Override PlayerInventoryStorage#insert to call #offer

* Add small comment regarding predicates

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

@@ -34,8 +34,8 @@ import net.fabricmc.fabric.impl.transfer.item.CursorSlotWrapper;
  * with an additional transactional wrapper for {@link PlayerInventory#offerOrDrop}.
  *
  * <p>Note that this is a wrapper around all the slots of the player inventory.
- * This may cause direct insertion to insert arbitrary items into equipment slots or other unexpected behavior.
- * To prevent this, {@link #offerOrDrop} is recommended for simple insertions.
+ * However, {@link #insert} is overriden to behave like {@link #offer}.
+ * For simple insertions, {@link #offer} or {@link #offerOrDrop} is recommended.
  * {@link #getSlots} can also be used and combined with {@link CombinedStorage} to retrieve a wrapper around a specific range of slots.
  *
  * @deprecated Experimental feature, we reserve the right to remove or change it without further notice.
@@ -69,6 +69,15 @@ public interface PlayerInventoryStorage extends InventoryStorage {
 		return CursorSlotWrapper.get(screenHandler);
 	}
 
+	/**
+	 * Insert items into this player inventory. Behaves the same as {@link #offer}.
+	 * More fine-tuned insertion, for example over a specific range of slots, is possible with {@linkplain #getSlots() the slot list}.
+	 *
+	 * @see #offer
+	 */
+	@Override
+	long insert(ItemVariant resource, long maxAmount, TransactionContext transaction);
+
 	/**
 	 * Add items to the inventory if possible, and drop any leftover items in the world, similar to {@link PlayerInventory#offerOrDrop}.
 	 *

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

@@ -32,6 +32,9 @@ import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
 /**
  * Helper functions to work with {@link Storage}s.
  *
+ * <p>Note that the functions that take a predicate iterate over the entire inventory in the worst case.
+ * If the resource is known, there will generally be a more performance efficient way.
+ *
  * @deprecated Experimental feature, we reserve the right to remove or change it without further notice.
  * The transfer API is a complex addition, and we want to be able to correct possible design mistakes.
  */
@@ -118,28 +121,40 @@ public final class StorageUtil {
 	/**
 	 * Attempt to find a resource stored in the passed storage.
 	 *
-	 * @param storage The storage to inspect, may be null.
-	 * @param transaction The current transaction, or {@code null} if a transaction should be opened for this query.
-	 * @param <T> The type of the stored resources.
+	 * @see #findStoredResource(Storage, Predicate, TransactionContext)
 	 * @return A non-blank resource stored in the storage, or {@code null} if none could be found.
 	 */
 	@Nullable
 	public static <T> T findStoredResource(@Nullable Storage<T> storage, @Nullable TransactionContext transaction) {
+		return findStoredResource(storage, r -> true, transaction);
+	}
+
+	/**
+	 * Attempt to find a resource stored in the passed storage that matches the passed filter.
+	 *
+	 * @param storage The storage to inspect, may be null.
+	 * @param filter The filter. Only a resource for which this filter returns {@code true} will be returned.
+	 * @param transaction The current transaction, or {@code null} if a transaction should be opened for this query.
+	 * @param <T> The type of the stored resources.
+	 * @return A non-blank resource stored in the storage that matches the filter, or {@code null} if none could be found.
+	 */
+	@Nullable
+	public static <T> T findStoredResource(@Nullable Storage<T> storage, Predicate<T> filter, @Nullable TransactionContext transaction) {
 		if (storage == null) return null;
 
 		if (transaction == null) {
 			try (Transaction outer = Transaction.openOuter()) {
-				return findStoredResourceInner(storage, outer);
+				return findStoredResourceInner(storage, filter, outer);
 			}
 		} else {
-			return findStoredResourceInner(storage, transaction);
+			return findStoredResourceInner(storage, filter, transaction);
 		}
 	}
 
 	@Nullable
-	private static <T> T findStoredResourceInner(Storage<T> storage, TransactionContext transaction) {
+	private static <T> T findStoredResourceInner(Storage<T> storage, Predicate<T> filter, TransactionContext transaction) {
 		for (StorageView<T> view : storage.iterable(transaction)) {
-			if (!view.isResourceBlank()) {
+			if (!view.isResourceBlank() && filter.test(view.getResource())) {
 				return view.getResource();
 			}
 		}
@@ -150,13 +165,25 @@ public final class StorageUtil {
 	/**
 	 * Attempt to find a resource stored in the passed storage that can be extracted.
 	 *
-	 * @param storage The storage to inspect, may be null.
-	 * @param transaction The current transaction, or {@code null} if a transaction should be opened for this query.
-	 * @param <T> The type of the stored resources.
+	 * @see #findExtractableResource(Storage, Predicate, TransactionContext)
 	 * @return A non-blank resource stored in the storage that can be extracted, or {@code null} if none could be found.
 	 */
 	@Nullable
 	public static <T> T findExtractableResource(@Nullable Storage<T> storage, @Nullable TransactionContext transaction) {
+		return findExtractableResource(storage, r -> true, transaction);
+	}
+
+	/**
+	 * Attempt to find a resource stored in the passed storage that matches the passed filter and can be extracted.
+	 *
+	 * @param storage The storage to inspect, may be null.
+	 * @param filter The filter. Only a resource for which this filter returns {@code true} will be returned.
+	 * @param transaction The current transaction, or {@code null} if a transaction should be opened for this query.
+	 * @param <T> The type of the stored resources.
+	 * @return A non-blank resource stored in the storage that matches the filter and can be extracted, or {@code null} if none could be found.
+	 */
+	@Nullable
+	public static <T> T findExtractableResource(@Nullable Storage<T> storage, Predicate<T> filter, @Nullable TransactionContext transaction) {
 		if (storage == null) return null;
 
 		try (Transaction nested = Transaction.openNested(transaction)) {
@@ -164,7 +191,7 @@ public final class StorageUtil {
 				// Extract below could change the resource, so we have to query it before extracting.
 				T resource = view.getResource();
 
-				if (!view.isResourceBlank() && view.extract(resource, Long.MAX_VALUE, nested) > 0) {
+				if (!view.isResourceBlank() && filter.test(resource) && view.extract(resource, Long.MAX_VALUE, nested) > 0) {
 					// Will abort the extraction.
 					return resource;
 				}
@@ -177,15 +204,28 @@ public final class StorageUtil {
 	/**
 	 * Attempt to find a resource stored in the passed storage that can be extracted, and how much of it can be extracted.
 	 *
-	 * @param storage The storage to inspect, may be null.
-	 * @param transaction The current transaction, or {@code null} if a transaction should be opened for this query.
-	 * @param <T> The type of the stored resources.
-	 * @return A non-blank resource stored in the storage that can be extracted and the strictly positive amount of it that can be extracted,
+	 * @see #findExtractableContent(Storage, Predicate, TransactionContext)
+	 * @return A non-blank resource stored in the storage that can be extracted, and the strictly positive amount of it that can be extracted,
 	 * or {@code null} if none could be found.
 	 */
 	@Nullable
 	public static <T> ResourceAmount<T> findExtractableContent(@Nullable Storage<T> storage, @Nullable TransactionContext transaction) {
-		T extractableResource = findExtractableResource(storage, transaction);
+		return findExtractableContent(storage, r -> true, transaction);
+	}
+
+	/**
+	 * Attempt to find a resource stored in the passed storage that can be extracted and matches the filter, and how much of it can be extracted.
+	 *
+	 * @param storage The storage to inspect, may be null.
+	 * @param filter The filter. Only a resource for which this filter returns {@code true} will be returned.
+	 * @param transaction The current transaction, or {@code null} if a transaction should be opened for this query.
+	 * @param <T> The type of the stored resources.
+	 * @return A non-blank resource stored in the storage that can be extracted and matches the filter, and the strictly positive amount of it that can be extracted,
+	 * or {@code null} if none could be found.
+	 */
+	@Nullable
+	public static <T> ResourceAmount<T> findExtractableContent(@Nullable Storage<T> storage, Predicate<T> filter, @Nullable TransactionContext transaction) {
+		T extractableResource = findExtractableResource(storage, filter, transaction);
 
 		if (extractableResource != null) {
 			long extractableAmount = storage.simulateExtract(extractableResource, Long.MAX_VALUE, transaction);

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

@@ -94,7 +94,7 @@ public abstract class SingleVariantItemStorage<T extends TransferVariant<?>> imp
 	 * then edit the NBT of the stack so it contains the correct resource and amount.
 	 *
 	 * <p>When the new amount is 0, it is recommended that the subtags corresponding to the resource and amount
-	 * be removed, for example using {@link ItemStack#removeSubTag}, so that newly-crafted containers can stack with
+	 * be removed, for example using {@link ItemStack#removeSubNbt}, so that newly-crafted containers can stack with
 	 * emptied containers.
 	 *
 	 * @param currentVariant Variant to which the modification should be applied.
@@ -184,7 +184,7 @@ public abstract class SingleVariantItemStorage<T extends TransferVariant<?>> imp
 		}
 
 		if (extracted > 0) {
-			if (tryUpdateStorage(resource, maxAmount - extracted, transaction)) {
+			if (tryUpdateStorage(resource, amount - extracted, transaction)) {
 				return extracted;
 			}
 		}

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

@@ -79,9 +79,10 @@ class InventorySlotWrapper extends SingleStackStorage {
 		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.
+		if (!original.isEmpty() && original.getItem() == currentStack.getItem()) {
+			// None is empty and the items match: just update the amount and NBT, and reuse the original stack.
 			original.setCount(currentStack.getCount());
+			original.setNbt(currentStack.hasNbt() ? currentStack.getNbt().copy() : null);
 			setStack(original);
 		} else {
 			// Otherwise assume everything was taken from original so empty it.

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

@@ -40,6 +40,11 @@ class PlayerInventoryStorageImpl extends InventoryStorageImpl implements PlayerI
 		this.playerInventory = playerInventory;
 	}
 
+	@Override
+	public long insert(ItemVariant resource, long maxAmount, TransactionContext transaction) {
+		return offer(resource, maxAmount, transaction);
+	}
+
 	@Override
 	public long offer(ItemVariant resource, long amount, TransactionContext tx) {
 		StoragePreconditions.notBlankNotNegative(resource, amount);

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

@@ -71,6 +71,8 @@ public class BaseStorageTests {
 		assertEquals(0L, noWater.simulateExtract(water, BUCKET, null));
 		// The fluid should be visible.
 		assertEquals(water, StorageUtil.findStoredResource(noWater, null));
+		// Test the filter.
+		assertEquals(null, StorageUtil.findStoredResource(noWater, fv -> fv.isOf(Fluids.LAVA), null));
 		// But it can't be extracted, even through a storage view.
 		assertEquals(null, StorageUtil.findExtractableResource(noWater, null));
 		assertEquals(null, StorageUtil.findExtractableContent(noWater, null));

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

@@ -175,5 +175,14 @@ class FluidItemTests {
 						null
 				)
 		);
+		// Test the filtering.
+		assertEquals(
+				null,
+				StorageUtil.findExtractableContent(
+						ContainerItemContext.withInitial(new ItemStack(Items.WATER_BUCKET)).find(FluidStorage.ITEM),
+						FluidVariant::hasNbt, // Only allow NBT -> won't match anything.
+						null
+				)
+		);
 	}
 }

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

@@ -26,6 +26,7 @@ import net.minecraft.inventory.SimpleInventory;
 import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
+import net.minecraft.nbt.NbtCompound;
 import net.minecraft.screen.ScreenHandler;
 import net.minecraft.util.math.Direction;
 
@@ -68,6 +69,21 @@ class ItemTests {
 		}
 
 		if (stack != inv.getStack(0)) throw new AssertionError("Stack should have stayed the same.");
+
+		// Also edit the stack when the item matches, even when the NBT and the count change.
+		ItemVariant oldVariant = ItemVariant.of(Items.DIAMOND);
+		NbtCompound testTag = new NbtCompound();
+		testTag.putInt("energy", 42);
+		ItemVariant newVariant = ItemVariant.of(Items.DIAMOND, testTag);
+
+		try (Transaction tx = Transaction.openOuter()) {
+			invWrapper.extract(oldVariant, 2, tx);
+			invWrapper.insert(newVariant, 5, tx);
+			tx.commit();
+		}
+
+		if (stack != inv.getStack(0)) throw new AssertionError("Stack should have stayed the same.");
+		if (!stackEquals(stack, newVariant, 5)) throw new AssertionError("Failed to update stack NBT or count.");
 	}
 
 	private static void testInventoryWrappers() {
@@ -113,7 +129,11 @@ class ItemTests {
 	}
 
 	private static boolean stackEquals(ItemStack stack, Item item, int count) {
-		return stack.getItem() == item && stack.getCount() == count;
+		return stackEquals(stack, ItemVariant.of(item), count);
+	}
+
+	private static boolean stackEquals(ItemStack stack, ItemVariant variant, int count) {
+		return variant.matches(stack) && stack.getCount() == count;
 	}
 
 	private static class TestSidedInventory extends SimpleInventory implements SidedInventory {

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

@@ -18,6 +18,8 @@ package net.fabricmc.fabric.test.transfer.unittests;
 
 import static net.fabricmc.fabric.test.transfer.unittests.TestUtil.assertEquals;
 
+import java.util.function.Function;
+
 import net.minecraft.entity.player.PlayerInventory;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
@@ -25,16 +27,20 @@ import net.minecraft.item.Items;
 import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
 import net.fabricmc.fabric.api.transfer.v1.item.PlayerInventoryStorage;
 import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
+import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
 
 public class PlayerInventoryStorageTests {
 	public static void run() {
-		testStacking();
+		// Ensure that offer stacks as expected.
+		testStacking(playerInv -> playerInv::offer);
+		// Also test that the behavior of insert matches that of offer.
+		testStacking(playerInv -> playerInv::insert);
 	}
 
-	private static void testStacking() {
-		// A bit hacky... but nothing should try using the player inventory as long as we don't call drop.
+	private static void testStacking(Function<PlayerInventoryStorage, InsertionFunction> inserterBuilder) {
+		// A bit hacky... but nothing should try using the null player entity as long as we don't call drop.
 		PlayerInventory inv = new PlayerInventory(null);
-		PlayerInventoryStorage wrapper = PlayerInventoryStorage.of(inv);
+		InsertionFunction inserter = inserterBuilder.apply(PlayerInventoryStorage.of(inv));
 
 		// Fill everything with stone besides the first two inventory slots.
 		inv.selectedSlot = 3;
@@ -48,14 +54,14 @@ public class PlayerInventoryStorageTests {
 		ItemVariant stone = ItemVariant.of(Items.STONE);
 
 		try (Transaction tx = Transaction.openOuter()) {
-			assertEquals(1L, wrapper.offer(stone, 1, tx));
+			assertEquals(1L, inserter.insert(stone, 1, tx));
 
 			// Should have went into the main stack
 			assertEquals(64, inv.main.get(3).getCount());
 		}
 
 		try (Transaction tx = Transaction.openOuter()) {
-			assertEquals(2L, wrapper.offer(stone, 2, tx));
+			assertEquals(2L, inserter.insert(stone, 2, tx));
 
 			// Should have went into the main and offhand stacks.
 			assertEquals(64, inv.main.get(3).getCount());
@@ -66,7 +72,7 @@ public class PlayerInventoryStorageTests {
 
 		// Should be just enough to fill existing stacks, but not touch slots 0, 1 and 2.
 		try (Transaction tx = Transaction.openOuter()) {
-			assertEquals(toInsertStacking, wrapper.offer(stone, toInsertStacking, tx));
+			assertEquals(toInsertStacking, inserter.insert(stone, toInsertStacking, tx));
 
 			assertEquals(64, inv.main.get(3).getCount());
 			assertEquals(64, inv.offHand.get(0).getCount());
@@ -80,13 +86,17 @@ public class PlayerInventoryStorageTests {
 			}
 
 			// Now insertion should fill the remaining stacks
-			assertEquals(150L, wrapper.offer(stone, 150, tx));
+			assertEquals(150L, inserter.insert(stone, 150, tx));
 			assertEquals(64, inv.main.get(0).getCount());
 			assertEquals(64, inv.main.get(1).getCount());
 			assertEquals(22, inv.main.get(2).getCount());
 
 			// Only 64 - 22 = 42 room left!
-			assertEquals(42L, wrapper.offer(stone, Long.MAX_VALUE, tx));
+			assertEquals(42L, inserter.insert(stone, Long.MAX_VALUE, tx));
 		}
 	}
+
+	private interface InsertionFunction {
+		long insert(ItemVariant variant, long maxAmount, TransactionContext transaction);
+	}
 }

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

@@ -80,7 +80,7 @@ public class SingleVariantItemStorageTests {
 
 		try (Transaction tx = Transaction.openOuter()) {
 			// Test extract along the way.
-			assertEquals(BUCKET, storage.extract(LAVA, BUCKET, tx));
+			assertEquals(BUCKET, storage.extract(LAVA, 10 * BUCKET, tx));
 
 			tx.commit();
 		}