Transfer API: Add slotted storage and non-empty iterator (#2908)

* Transfer API: Add non-empty iterator

* Add SlottedStorage

* Add StorageUtil.extractAny

* Undeprecate ContainerItemContext.withInitial

* Add licenses

* Revert "Undeprecate ContainerItemContext.withInitial"

This reverts commit dcf123eb332ff642cdbd5fda0d8d2237794d93fc.

* Tweaks

* Make SlottedStorage#getSlots return a view, remove useless field, add UnmodifiableView annotations

* Remove useless @inheritDoc

* Fix infinite loop in the tests

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/context/ContainerItemContext.java

@@ -20,6 +20,7 @@ import java.util.List;
 
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.UnmodifiableView;
 
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.entity.player.PlayerInventory;
@@ -293,5 +294,6 @@ public interface ContainerItemContext {
 	 *
 	 * @return An unmodifiable list containing additional slots of this context. If no additional slot is available, the list is empty.
 	 */
+	@UnmodifiableView
 	List<SingleSlotStorage<ItemVariant>> getAdditionalSlots();
 }

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

@@ -21,6 +21,7 @@ import java.util.Objects;
 
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.UnmodifiableView;
 
 import net.minecraft.entity.player.PlayerInventory;
 import net.minecraft.inventory.Inventory;
@@ -28,7 +29,7 @@ import net.minecraft.inventory.SidedInventory;
 import net.minecraft.inventory.SimpleInventory;
 import net.minecraft.util.math.Direction;
 
-import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
+import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.base.CombinedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleSlotStorage;
 import net.fabricmc.fabric.impl.transfer.item.InventoryStorageImpl;
@@ -50,7 +51,7 @@ import net.fabricmc.fabric.impl.transfer.item.InventoryStorageImpl;
  */
 @ApiStatus.Experimental
 @ApiStatus.NonExtendable
-public interface InventoryStorage extends Storage<ItemVariant> {
+public interface InventoryStorage extends SlottedStorage<ItemVariant> {
 	/**
 	 * Return a wrapper around an {@link Inventory}.
 	 *
@@ -69,11 +70,16 @@ public interface InventoryStorage extends Storage<ItemVariant> {
 	 * Retrieve an unmodifiable list of the wrappers for the slots in this inventory.
 	 * Each wrapper corresponds to a single slot in the inventory.
 	 */
+	@Override
+	@UnmodifiableView
 	List<SingleSlotStorage<ItemVariant>> getSlots();
 
-	/**
+	@Override
-	 * Retrieve a wrapper around a specific slot of the inventory.
+	default int getSlotCount() {
-	 */
+		return getSlots().size();
+	}
+
+	@Override
 	default SingleSlotStorage<ItemVariant> getSlot(int slot) {
 		return getSlots().get(slot);
 	}

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

@@ -32,10 +32,12 @@ import net.minecraft.util.Identifier;
 import net.minecraft.util.math.Direction;
 
 import net.fabricmc.fabric.api.lookup.v1.block.BlockApiLookup;
-import net.fabricmc.fabric.api.transfer.v1.storage.base.SidedStorageBlockEntity;
 import net.fabricmc.fabric.api.transfer.v1.item.base.SingleStackStorage;
+import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.CombinedSlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.base.CombinedStorage;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.SidedStorageBlockEntity;
 import net.fabricmc.fabric.impl.transfer.item.ComposterWrapper;
 import net.fabricmc.fabric.mixin.transfer.DoubleInventoryAccessor;
 
@@ -119,10 +121,10 @@ public final class ItemStorage {
 
 					// For double chests, we need to retrieve a wrapper for each part separately.
 					if (inventoryToWrap instanceof DoubleInventoryAccessor accessor) {
-						Storage<ItemVariant> first = InventoryStorage.of(accessor.fabric_getFirst(), direction);
+						SlottedStorage<ItemVariant> first = InventoryStorage.of(accessor.fabric_getFirst(), direction);
-						Storage<ItemVariant> second = InventoryStorage.of(accessor.fabric_getSecond(), direction);
+						SlottedStorage<ItemVariant> second = InventoryStorage.of(accessor.fabric_getSecond(), direction);
 
-						return new CombinedStorage<>(List.of(first, second));
+						return new CombinedSlottedStorage<>(List.of(first, second));
 					}
 				} else {
 					inventoryToWrap = inventory;

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

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.api.transfer.v1.storage;
+
+import java.util.List;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.UnmodifiableView;
+
+import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleSlotStorage;
+import net.fabricmc.fabric.api.transfer.v1.context.ContainerItemContext;
+import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
+
+/**
+ * A {@link Storage} implementation made of indexed slots.
+ *
+ * <p>Please note that some storages may not implement this interface.
+ * It is up to the storage implementation to decide whether to implement this interface or not.
+ * Checking whether a storage is slotted can be done using {@code instanceof}.
+ *
+ * @param <T> The type of the stored resources.
+ *
+ * <b>Experimental feature</b>, 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.
+ */
+@ApiStatus.Experimental
+public interface SlottedStorage<T> extends Storage<T> {
+	/**
+	 * Retrieve the number of slots in this storage.
+	 */
+	int getSlotCount();
+
+	/**
+	 * Retrieve a specific slot of this storage.
+	 *
+	 * @throws IndexOutOfBoundsException If the slot index is out of bounds.
+	 */
+	SingleSlotStorage<T> getSlot(int slot);
+
+	/**
+	 * Retrieve a list containing all the slots of this storage. <b>The list must not be modified.</b>
+	 *
+	 * <p>This function can be used to interface with code that requires a slot list,
+	 * for example {@link StorageUtil#insertStacking} or {@link ContainerItemContext#getAdditionalSlots()}.
+	 *
+	 * <p>It is guaranteed that calling this function is fast.
+	 * The default implementation returns a view over the storage that delegates to {@link #getSlotCount} and {@link #getSlot}.
+	 */
+	@UnmodifiableView
+	default List<SingleSlotStorage<T>> getSlots() {
+		return TransferApiImpl.makeListView(this);
+	}
+}

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

@@ -151,6 +151,39 @@ public interface Storage<T> extends Iterable<StorageView<T>> {
 	@Override
 	Iterator<StorageView<T>> iterator();
 
+	/**
+	 * Same as {@link #iterator()}, but the iterator is guaranteed to skip over empty views,
+	 * i.e. views that {@linkplain StorageView#isResourceBlank() contain blank resources} or have a zero {@linkplain StorageView#getAmount() amount}.
+	 *
+	 * <p>This can provide a large performance benefit over {@link #iterator()} if the caller is only interested in non-empty views,
+	 * for example because it is trying to extract resources from the storage.
+	 *
+	 * <p>This function should only be overridden if the storage is able to provide an optimized iterator over non-empty views,
+	 * for example because it is keeping an index of non-empty views.
+	 * Otherwise, the default implementation simply calls {@link #iterator()} and filters out empty views.
+	 *
+	 * <p>When implementing this function, note that the guarantees of {@link #iterator()} still apply.
+	 * In particular, {@link #insert} and {@link #extract} may be called safely during iteration.
+	 *
+	 * @return An iterator over the non-empty views of this storage. Calling remove on the iterator is not allowed.
+	 */
+	default Iterator<StorageView<T>> nonEmptyIterator() {
+		return TransferApiImpl.filterEmptyViews(iterator());
+	}
+
+	/**
+	 * Convenient helper to get an {@link Iterable} over the {@linkplain #nonEmptyIterator() non-empty views} of this storage, for use in for-each loops.
+	 *
+	 * <p><pre>{@code
+	 * for (StorageView<T> view : storage.nonEmptyViews()) {
+	 *     // Do something with the view
+	 * }
+	 * }</pre>
+	 */
+	default Iterable<StorageView<T>> nonEmptyViews() {
+		return this::nonEmptyIterator;
+	}
+
 	/**
 	 * Return a view over this storage, for a specific resource, or {@code null} if none is quickly available.
 	 *

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

@@ -88,8 +88,7 @@ public final class StorageUtil {
 		long totalMoved = 0;
 
 		try (Transaction iterationTransaction = Transaction.openNested(transaction)) {
-			for (StorageView<T> view : from) {
+			for (StorageView<T> view : from.nonEmptyViews()) {
-				if (view.isResourceBlank()) continue;
 				T resource = view.getResource();
 				if (!filter.test(resource)) continue;
 				long maxExtracted;
@@ -124,6 +123,32 @@ public final class StorageUtil {
 		return totalMoved;
 	}
 
+	/**
+	 * Try to extract any resource from a storage, up to a maximum amount.
+	 *
+	 * <p>This function will only ever pull from one storage view of the storage, even if multiple storage views contain the same resource.
+	 *
+	 * @param storage The storage, may be null.
+	 * @param maxAmount The maximum to extract.
+	 * @param transaction The transaction this operation is part of.
+	 * @return A non-blank resource and the strictly positive amount of it that was extracted from the storage,
+	 * or {@code null} if none could be found.
+	 */
+	@Nullable
+	public static <T> ResourceAmount<T> extractAny(@Nullable Storage<T> storage, long maxAmount, TransactionContext transaction) {
+		StoragePreconditions.notNegative(maxAmount);
+
+		if (storage == null) return null;
+
+		for (StorageView<T> view : storage.nonEmptyViews()) {
+			T resource = view.getResource();
+			long amount = view.extract(resource, maxAmount, transaction);
+			if (amount > 0) return new ResourceAmount<>(resource, amount);
+		}
+
+		return null;
+	}
+
 	/**
 	 * Try to insert up to some amount of a resource into a list of storage slots, trying to "stack" first,
 	 * i.e. prioritizing slots that already contain the resource.
@@ -131,7 +156,7 @@ public final class StorageUtil {
 	 * @return How much was inserted.
 	 * @see Storage#insert
 	 */
-	public static <T> long insertStacking(List<SingleSlotStorage<T>> slots, T resource, long maxAmount, TransactionContext transaction) {
+	public static <T> long insertStacking(List<? extends SingleSlotStorage<T>> slots, T resource, long maxAmount, TransactionContext transaction) {
 		StoragePreconditions.notNegative(maxAmount);
 		long amount = 0;
 
@@ -150,6 +175,27 @@ public final class StorageUtil {
 		return amount;
 	}
 
+	/**
+	 * Insert resources in a storage, attempting to stack them with existing resources first if possible.
+	 *
+	 * @param storage The storage, may be null.
+	 * @param resource The resource to insert. May not be blank.
+	 * @param maxAmount The maximum amount of resource to insert. May not be negative.
+	 * @param transaction The transaction this operation is part of.
+	 * @return A nonnegative integer not greater than maxAmount: the amount that was inserted.
+	 */
+	public static <T> long tryInsertStacking(@Nullable Storage<T> storage, T resource, long maxAmount, TransactionContext transaction) {
+		StoragePreconditions.notNegative(maxAmount);
+
+		if (storage instanceof SlottedStorage<T> slottedStorage) {
+			return insertStacking(slottedStorage.getSlots(), resource, maxAmount, transaction);
+		} else if (storage != null) {
+			return storage.insert(resource, maxAmount, transaction);
+		} else {
+			return 0;
+		}
+	}
+
 	/**
 	 * Attempt to find a resource stored in the passed storage.
 	 *
@@ -174,8 +220,8 @@ public final class StorageUtil {
 		Objects.requireNonNull(filter, "Filter may not be null");
 		if (storage == null) return null;
 
-		for (StorageView<T> view : storage) {
+		for (StorageView<T> view : storage.nonEmptyViews()) {
-			if (!view.isResourceBlank() && filter.test(view.getResource())) {
+			if (filter.test(view.getResource())) {
 				return view.getResource();
 			}
 		}
@@ -209,11 +255,11 @@ public final class StorageUtil {
 		if (storage == null) return null;
 
 		try (Transaction nested = Transaction.openNested(transaction)) {
-			for (StorageView<T> view : storage) {
+			for (StorageView<T> view : storage.nonEmptyViews()) {
 				// Extract below could change the resource, so we have to query it before extracting.
 				T resource = view.getResource();
 
-				if (!view.isResourceBlank() && filter.test(resource) && view.extract(resource, Long.MAX_VALUE, nested) > 0) {
+				if (filter.test(resource) && view.extract(resource, Long.MAX_VALUE, nested) > 0) {
 					// Will abort the extraction.
 					return resource;
 				}

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

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.api.transfer.v1.storage.base;
+
+import java.util.List;
+
+import org.jetbrains.annotations.ApiStatus;
+
+import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
+import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
+
+/**
+ * A {@link Storage} wrapping multiple slotted storages.
+ * Same as {@link CombinedStorage}, but for {@link SlottedStorage}s.
+ *
+ * @param <T> The type of the stored resources.
+ * @param <S> The class of every part. {@code ? extends Storage<T>} can be used if the parts are of different types.
+ *
+ * <b>Experimental feature</b>, 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.
+ */
+@ApiStatus.Experimental
+public class CombinedSlottedStorage<T, S extends SlottedStorage<T>> extends CombinedStorage<T, S> implements SlottedStorage<T> {
+	public CombinedSlottedStorage(List<S> parts) {
+		super(parts);
+	}
+
+	@Override
+	public int getSlotCount() {
+		int count = 0;
+
+		for (S part : parts) {
+			count += part.getSlotCount();
+		}
+
+		return count;
+	}
+
+	@Override
+	public SingleSlotStorage<T> getSlot(int slot) {
+		int updatedSlot = slot;
+
+		for (SlottedStorage<T> part : parts) {
+			if (updatedSlot < part.getSlotCount()) {
+				return part.getSlot(updatedSlot);
+			}
+
+			updatedSlot -= part.getSlotCount();
+		}
+
+		throw new IndexOutOfBoundsException("Slot " + slot + " is out of bounds. This storage has size " + getSlotCount());
+	}
+}

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

@@ -20,7 +20,7 @@ import java.util.Iterator;
 
 import org.jetbrains.annotations.ApiStatus;
 
-import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
+import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.StorageView;
 import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
 
@@ -34,9 +34,23 @@ import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
  * The transfer API is a complex addition, and we want to be able to correct possible design mistakes.
  */
 @ApiStatus.Experimental
-public interface SingleSlotStorage<T> extends Storage<T>, StorageView<T> {
+public interface SingleSlotStorage<T> extends SlottedStorage<T>, StorageView<T> {
 	@Override
 	default Iterator<StorageView<T>> iterator() {
 		return TransferApiImpl.singletonIterator(this);
 	}
+
+	@Override
+	default int getSlotCount() {
+		return 1;
+	}
+
+	@Override
+	default SingleSlotStorage<T> getSlot(int slot) {
+		if (slot != 0) {
+			throw new IndexOutOfBoundsException("Slot " + slot + " does not exist in a single-slot storage.");
+		}
+
+		return this;
+	}
 }

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

@@ -47,6 +47,9 @@ public abstract class SingleVariantStorage<T extends TransferVariant<?>> extends
 
 	/**
 	 * Return the blank variant.
+	 *
+	 * <p>Note: this is called very early in the constructor.
+	 * If fields need to be accessed from this function, make sure to re-initialize {@link #variant} yourself.
 	 */
 	protected abstract T getBlankVariant();
 

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

@@ -16,16 +16,20 @@
 
 package net.fabricmc.fabric.impl.transfer;
 
+import java.util.AbstractList;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.slf4j.LoggerFactory;
 import org.slf4j.Logger;
 
+import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
 import net.fabricmc.fabric.api.transfer.v1.storage.StorageView;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleSlotStorage;
 import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
 
 public class TransferApiImpl {
@@ -64,11 +68,6 @@ public class TransferApiImpl {
 		}
 	};
 
-	/**
-	 * Not null when writing to an inventory in a transaction, null otherwise.
-	 */
-	public static final ThreadLocal<Object> SUPPRESS_SPECIAL_LOGIC = new ThreadLocal<>();
-
 	public static <T> Iterator<T> singletonIterator(T it) {
 		return new Iterator<T>() {
 			boolean hasNext = true;
@@ -89,4 +88,56 @@ public class TransferApiImpl {
 			}
 		};
 	}
+
+	public static <T> Iterator<StorageView<T>> filterEmptyViews(Iterator<StorageView<T>> iterator) {
+		return new Iterator<>() {
+			StorageView<T> next;
+
+			{
+				findNext();
+			}
+
+			private void findNext() {
+				while (iterator.hasNext()) {
+					next = iterator.next();
+
+					if (next.getAmount() > 0 && !next.isResourceBlank()) {
+						return;
+					}
+				}
+
+				next = null;
+			}
+
+			@Override
+			public boolean hasNext() {
+				return next != null;
+			}
+
+			@Override
+			public StorageView<T> next() {
+				if (!hasNext()) {
+					throw new NoSuchElementException();
+				}
+
+				StorageView<T> ret = next;
+				findNext();
+				return ret;
+			}
+		};
+	}
+
+	public static <T> List<SingleSlotStorage<T>> makeListView(SlottedStorage<T> storage) {
+		return new AbstractList<>() {
+			@Override
+			public SingleSlotStorage<T> get(int index) {
+				return storage.getSlot(index);
+			}
+
+			@Override
+			public int size() {
+				return storage.getSlotCount();
+			}
+		};
+	}
 }

fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/gametests/VanillaStorageTests.java

@@ -25,6 +25,7 @@ import net.minecraft.block.ComparatorBlock;
 import net.minecraft.block.entity.BrewingStandBlockEntity;
 import net.minecraft.block.entity.ChiseledBookshelfBlockEntity;
 import net.minecraft.block.entity.FurnaceBlockEntity;
+import net.minecraft.block.entity.HopperBlockEntity;
 import net.minecraft.block.entity.ShulkerBoxBlockEntity;
 import net.minecraft.inventory.Inventory;
 import net.minecraft.item.ItemStack;
@@ -41,6 +42,7 @@ import net.fabricmc.fabric.api.gametest.v1.FabricGameTest;
 import net.fabricmc.fabric.api.transfer.v1.item.InventoryStorage;
 import net.fabricmc.fabric.api.transfer.v1.item.ItemStorage;
 import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
 import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
 import net.fabricmc.fabric.test.transfer.mixin.AbstractFurnaceBlockEntityAccessor;
@@ -289,6 +291,23 @@ public class VanillaStorageTests {
 			tx.commit();
 		}
 
+		// Check that the inventory and slotted storages match
+		Inventory inventory = HopperBlockEntity.getInventoryAt(context.getWorld(), context.getAbsolutePos(chestPos));
+		context.assertTrue(inventory != null, "Inventory must not be null");
+
+		if (!(storage instanceof SlottedStorage<ItemVariant> slottedStorage)) {
+			throw new GameTestException("Double chest storage must be a SlottedStorage");
+		}
+
+		for (int i = 0; i < inventory.size(); ++i) {
+			ItemStack stack = inventory.getStack(i);
+			ItemVariant variant = ItemVariant.of(stack.getItem());
+			context.assertTrue(variant.matches(stack), "Item variant in slot " + i + " must match stack");
+			long expectedCount = stack.getCount();
+			long actualCount = slottedStorage.getSlot(i).getAmount();
+			context.assertTrue(expectedCount == actualCount, "Slot " + i + " should have " + expectedCount + " items, but has " + actualCount);
+		}
+
 		// Check that an update is queued for every single comparator
 		MutableInt comparatorCount = new MutableInt();