From 9ab11d529d27f9847c218058066d69d77cd44380 Mon Sep 17 00:00:00 2001
From: Technici4n <13494793+technici4n@users.noreply.github.com>
Date: Sun, 26 Nov 2023 13:00:57 +0000
Subject: [PATCH] Fix empty storage iterator returning views that become empty
 during iteration (#3423)

(cherry picked from commit aaf9c9690c0f6511924af65d83ebaff56ec540b2)
---
 .../api/transfer/v1/storage/Storage.java      |  3 +-
 .../fabric/impl/transfer/TransferApiImpl.java | 38 -------------------
 .../transfer/unittests/BaseStorageTests.java  | 25 ++++++++++++
 3 files changed, 27 insertions(+), 39 deletions(-)

diff --git a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/storage/Storage.java b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/storage/Storage.java
index 7398140fb..5ab49372a 100644
--- a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/storage/Storage.java
+++ b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/storage/Storage.java
@@ -18,6 +18,7 @@ package net.fabricmc.fabric.api.transfer.v1.storage;
 
 import java.util.Iterator;
 
+import com.google.common.collect.Iterators;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
 
@@ -152,7 +153,7 @@ public interface Storage<T> extends Iterable<StorageView<T>> {
 	 * @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());
+		return Iterators.filter(iterator(), view -> view.getAmount() > 0 && !view.isResourceBlank());
 	}
 
 	/**
diff --git a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/TransferApiImpl.java b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/TransferApiImpl.java
index 274915a55..cb8a1bd87 100644
--- a/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/TransferApiImpl.java
+++ b/fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/TransferApiImpl.java
@@ -94,44 +94,6 @@ 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
diff --git a/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/BaseStorageTests.java b/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/BaseStorageTests.java
index d8f9b22ac..1f39a034f 100644
--- a/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/BaseStorageTests.java
+++ b/fabric-transfer-api-v1/src/testmod/java/net/fabricmc/fabric/test/transfer/unittests/BaseStorageTests.java
@@ -19,11 +19,15 @@ package net.fabricmc.fabric.test.transfer.unittests;
 import static net.fabricmc.fabric.api.transfer.v1.fluid.FluidConstants.BUCKET;
 import static net.fabricmc.fabric.test.transfer.unittests.TestUtil.assertEquals;
 
+import java.util.Iterator;
+
 import net.minecraft.fluid.Fluids;
 
 import net.fabricmc.fabric.api.transfer.v1.fluid.FluidVariant;
+import net.fabricmc.fabric.api.transfer.v1.fluid.base.SingleFluidStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
 import net.fabricmc.fabric.api.transfer.v1.storage.StorageUtil;
+import net.fabricmc.fabric.api.transfer.v1.storage.StorageView;
 import net.fabricmc.fabric.api.transfer.v1.storage.base.FilteringStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleVariantStorage;
 import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
@@ -31,6 +35,7 @@ import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
 public class BaseStorageTests {
 	public static void run() {
 		testFilteringStorage();
+		testNonEmptyIteratorWithModifiedView();
 	}
 
 	private static void testFilteringStorage() {
@@ -92,4 +97,24 @@ public class BaseStorageTests {
 
 		assertEquals(BUCKET, StorageUtil.simulateExtract(storage, lava, BUCKET, null));
 	}
+
+	/**
+	 * Regression test for <a href="https://github.com/FabricMC/fabric/issues/3414">
+	 * {@code nonEmptyIterator} not handling views that become empty during iteration correctly</a>.
+	 */
+	private static void testNonEmptyIteratorWithModifiedView() {
+		SingleVariantStorage<FluidVariant> storage = SingleFluidStorage.withFixedCapacity(BUCKET, () -> { });
+		storage.variant = FluidVariant.of(Fluids.WATER);
+
+		Iterator<StorageView<FluidVariant>> iterator = storage.nonEmptyIterator();
+		storage.amount = BUCKET;
+		// Iterator should have a next element now
+		assertEquals(true, iterator.hasNext());
+		assertEquals(storage, iterator.next());
+
+		iterator = storage.nonEmptyIterator();
+		storage.amount = 0;
+		// Iterator should not have a next element...
+		assertEquals(false, iterator.hasNext());
+	}
 }