Optimize Brass Funnel Performance in EXACTLY Mode

src/main/java/com/simibubi/create/foundation/item/ItemHelper.java

@@ -5,6 +5,10 @@
 import java.util.function.Function;
 import java.util.function.Predicate;
 
+import com.simibubi.create.foundation.utility.VisitedItemStackTracker;
+
+import com.simibubi.create.foundation.utility.VisitedItemStackTracker.SlotAmountRecord;
+
 import org.jetbrains.annotations.Nullable;
 
 import org.apache.commons.lang3.mutable.MutableInt;
@@ -182,70 +186,64 @@ public static ItemStack extract(IItemHandler inv, Predicate<ItemStack> test, int
 
 	public static ItemStack extract(IItemHandler inv, Predicate<ItemStack> test, ExtractionCountMode mode, int amount,
 									boolean simulate) {
-		ItemStack extracting = ItemStack.EMPTY;
-		boolean amountRequired = mode == ExtractionCountMode.EXACTLY;
-		boolean checkHasEnoughItems = amountRequired;
-		boolean hasEnoughItems = !checkHasEnoughItems;
-		boolean potentialOtherMatch = false;
-		int maxExtractionCount = amount;
-
-		Extraction:
-		do {
-			extracting = ItemStack.EMPTY;
-
-			for (int slot = 0; slot < inv.getSlots(); slot++) {
-				ItemStack slotStack = inv.getStackInSlot(slot);
-				if (slotStack.isEmpty())
+		if (mode == ItemHelper.ExtractionCountMode.EXACTLY) {
+			VisitedItemStackTracker tracker = new VisitedItemStackTracker();
+			for (int i = 0; i < inv.getSlots(); i++) {
+				ItemStack stackIn = inv.getStackInSlot(i);
+				if (stackIn.isEmpty() || stackIn.getMaxStackSize() < amount)
 					continue;
-				int amountToExtractFromThisSlot =
-					Math.min(maxExtractionCount - extracting.getCount(), slotStack.getMaxStackSize());
-				ItemStack stack = inv.extractItem(slot, amountToExtractFromThisSlot, true);
 
-				if (stack.isEmpty())
+				ItemStack extracted = inv.extractItem(i, Math.min(stackIn.getCount(), amount), true);
+				if (extracted.isEmpty() || !test.test(extracted))
 					continue;
-				if (!test.test(stack))
-					continue;
-				if (!extracting.isEmpty() && !canItemStackAmountsStack(stack, extracting)) {
-					potentialOtherMatch = true;
-					continue;
-				}
 
-				if (extracting.isEmpty())
-					extracting = stack.copy();
-				else
-					extracting.grow(stack.getCount());
-
-				if (!simulate && hasEnoughItems)
-					inv.extractItem(slot, stack.getCount(), false);
-
-				if (extracting.getCount() >= maxExtractionCount) {
-					if (checkHasEnoughItems) {
-						hasEnoughItems = true;
-						checkHasEnoughItems = false;
-						continue Extraction;
-					} else {
-						break Extraction;
+				VisitedItemStackTracker.SlotAmountRecord slotRecord = tracker.update(extracted.copy(), i);
+				if (slotRecord.totalAmount >= amount) {
+					ItemStack result = extracted.copyWithCount(amount);
+					if (!simulate) {
+						for (int slot: slotRecord.slots) {
+							int extractAmount = Math.min(inv.getStackInSlot(slot).getCount(), amount);
+							amount -= inv.extractItem(slot, extractAmount, false).getCount();
+						}
 					}
+					return result;
 				}
 			}
+		} else {
+			VisitedItemStackTracker.SlotAmountRecord slotRecord = new SlotAmountRecord();
+			ItemStack result = ItemStack.EMPTY;
+			int maxExtractAmount = amount;
+			for (int i = 0; i < inv.getSlots(); i++) {
+				ItemStack stackIn = inv.getStackInSlot(i);
+				if (stackIn.isEmpty() || (!result.isEmpty() && !ItemStack.isSameItemSameComponents(result, stackIn)))
+					continue;
 
-			if (!extracting.isEmpty() && !hasEnoughItems && potentialOtherMatch) {
-				ItemStack blackListed = extracting.copy();
-				test = test.and(i -> !ItemStack.isSameItemSameComponents(i, blackListed));
-				continue;
-			}
-
-			if (checkHasEnoughItems)
-				checkHasEnoughItems = false;
-			else
-				break Extraction;
-
-		} while (true);
+				ItemStack extracted = inv.extractItem(i, Math.min(stackIn.getCount(), maxExtractAmount), true);
+				if (extracted.isEmpty() || !test.test(extracted))
+					continue;
 
-		if (amountRequired && extracting.getCount() < amount)
-			return ItemStack.EMPTY;
+				if (result.isEmpty()) {
+					result = extracted.copy();
+					maxExtractAmount = Math.min(amount, extracted.getMaxStackSize());
+				}
 
-		return extracting;
+				slotRecord.add(i, extracted.getCount());
+				if (slotRecord.totalAmount >= maxExtractAmount)
+					break;
+			}
+			if (!result.isEmpty()) {
+				amount = Math.min(slotRecord.totalAmount, maxExtractAmount);
+				result.setCount(amount);
+				if (!simulate) {
+					for (int slot: slotRecord.slots) {
+						int extractAmount = Math.min(inv.getStackInSlot(slot).getCount(), amount);
+						amount -= inv.extractItem(slot, extractAmount, false).getCount();
+					}
+				}
+				return result;
+			}
+		}
+		return ItemStack.EMPTY;
 	}
 
 	public static ItemStack extract(IItemHandler inv, Predicate<ItemStack> test,

src/main/java/com/simibubi/create/foundation/utility/VisitedItemStackTracker.java

@@ -0,0 +1,54 @@
+package com.simibubi.create.foundation.utility;
+
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import net.minecraft.world.item.ItemStack;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class VisitedItemStackTracker {
+	Map<AbstractItemStack, SlotAmountRecord> itemStackMap = new HashMap<>();
+
+	public SlotAmountRecord update(ItemStack itemStack, int slot) {
+		SlotAmountRecord record = itemStackMap.computeIfAbsent(new AbstractItemStack(itemStack), k -> new SlotAmountRecord());
+		record.add(slot, itemStack.getCount());
+		return record;
+	}
+
+	public static class AbstractItemStack {
+		private final ItemStack stack;
+		private boolean initialized = false;
+		private int hashCode;
+
+		AbstractItemStack(ItemStack stack) {
+			this.stack = stack;
+		}
+
+		@Override
+		public int hashCode() {
+			if (!initialized) {
+				initialized = true;
+				hashCode = stack.getItem().hashCode();
+				hashCode = hashCode * 31 + stack.getComponents().hashCode();
+			}
+			return hashCode;
+		}
+
+		@Override
+		public boolean equals(Object other) {
+			if (!(other instanceof AbstractItemStack otherStack))
+				return false;
+			return ItemStack.isSameItemSameComponents(stack, otherStack.stack);
+		}
+	}
+
+	public static class SlotAmountRecord {
+		public final IntArrayList slots = new IntArrayList();
+		public int totalAmount = 0;
+
+		public void add(int slot, int amount) {
+			slots.add(slot);
+			totalAmount += amount;
+		}
+	}
+}