2 files changed, 50 insertions, 20 deletions
diff --git a/api/src/main/java/me/shedaniel/rei/api/common/transfer/ItemRecipeFinder.java b/api/src/main/java/me/shedaniel/rei/api/common/transfer/ItemRecipeFinder.java
index 00646b72a..59480aaba 100644
--- a/api/src/main/java/me/shedaniel/rei/api/common/transfer/ItemRecipeFinder.java
+++ b/api/src/main/java/me/shedaniel/rei/api/common/transfer/ItemRecipeFinder.java
@@ -33,12 +33,14 @@ import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemStack;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
 public class ItemRecipeFinder {
     private final Interner<ItemKey> keys = Interners.newWeakInterner();
-    private final RecipeFinder<ItemKey> finder = new RecipeFinder<>();
+    private final RecipeFinder<ItemKey, Ingredient> finder = new RecipeFinder<>();
     
     public boolean contains(ItemStack item) {
         return finder.contains(ofKey(item));
@@ -74,27 +76,54 @@ public class ItemRecipeFinder {
     }
     
     public boolean findRecipe(List<List<ItemStack>> list, int maxCrafts, @Nullable Consumer<ItemStack> output) {
-        return finder.findRecipe(CollectionUtils.map(list, this::ofKeys), maxCrafts, itemKey -> {
+        return finder.findRecipe(toIngredients(list), maxCrafts, flatten(itemStack -> {
             if (output != null) {
-                output.accept(new ItemStack(itemKey.item(), 1, itemKey.patch()));
+                output.accept(itemStack);
             }
-        });
+        }));
     }
     
     public int countRecipeCrafts(List<List<ItemStack>> list, int maxCrafts, @Nullable Consumer<ItemStack> output) {
-        return finder.countRecipeCrafts(CollectionUtils.map(list, this::ofKeys), maxCrafts, itemKey -> {
+        return finder.countRecipeCrafts(toIngredients(list), maxCrafts, flatten(itemStack -> {
             if (output != null) {
-                output.accept(new ItemStack(itemKey.item(), 1, itemKey.patch()));
+                output.accept(itemStack);
             }
-        });
+        }));
     }
     
     private ItemKey ofKey(ItemStack itemStack) {
         return keys.intern(new ItemKey(itemStack.getItemHolder(), itemStack.getComponentsPatch()));
     }
     
-    private RecipeFinder.Ingredient<ItemKey> ofKeys(List<ItemStack> itemStack) {
-        return new RecipeFinder.Ingredient<>(CollectionUtils.map(itemStack, this::ofKey));
+    private Ingredient ofKeys(int index, List<ItemStack> itemStack) {
+        return new Ingredient(index, CollectionUtils.map(itemStack, this::ofKey));
+    }
+    
+    private List<Ingredient> toIngredients(List<List<ItemStack>> list) {
+        List<Ingredient> ingredients = new ArrayList<>();
+        
+        for (int i = 0; i < list.size(); i++) {
+            List<ItemStack> stacks = list.get(i);
+            if (!stacks.isEmpty()) {
+                ingredients.add(ofKeys(i, stacks));
+            }
+        }
+        
+        return ingredients;
+    }
+    
+    private static BiConsumer<ItemKey, Ingredient> flatten(Consumer<ItemStack> consumer) {
+        int[] lastIndex = {-1};
+        return (itemKey, ingredient) -> {
+            for (int i = lastIndex[0] + 1; i < ingredient.index(); i++) {
+                consumer.accept(ItemStack.EMPTY);
+            }
+            consumer.accept(new ItemStack(itemKey.item(), 1, itemKey.patch()));
+            lastIndex[0] = ingredient.index();
+        };
+    }
+    
+    private record Ingredient(int index, List<ItemKey> elements) implements RecipeFinder.Ingredient<ItemKey> {
     }
     
     private record ItemKey(Holder<Item> item, DataComponentPatch patch) {
diff --git a/api/src/main/java/me/shedaniel/rei/api/common/transfer/RecipeFinder.java b/api/src/main/java/me/shedaniel/rei/api/common/transfer/RecipeFinder.java
index 9b260594c..c19d12ae1 100644
--- a/api/src/main/java/me/shedaniel/rei/api/common/transfer/RecipeFinder.java
+++ b/api/src/main/java/me/shedaniel/rei/api/common/transfer/RecipeFinder.java
@@ -32,9 +32,9 @@ import org.jetbrains.annotations.Nullable;
 import java.util.BitSet;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Consumer;
+import java.util.function.BiConsumer;
 
-public class RecipeFinder<T> {
+public class RecipeFinder<T, I extends RecipeFinder.Ingredient<T>> {
     public final Reference2IntOpenHashMap<T> amounts = new Reference2IntOpenHashMap<>();
     
     public boolean contains(T item) {
@@ -56,11 +56,11 @@ public class RecipeFinder<T> {
         this.amounts.addTo(item, amount);
     }
     
-    public boolean findRecipe(List<Ingredient<T>> list, int maxCrafts, @Nullable Consumer<T> output) {
+    public boolean findRecipe(List<I> list, int maxCrafts, @Nullable BiConsumer<T, I> output) {
         return new Filter(list).tryPick(maxCrafts, output);
     }
     
-    public int countRecipeCrafts(List<Ingredient<T>> list, int maxCrafts, @Nullable Consumer<T> output) {
+    public int countRecipeCrafts(List<I> list, int maxCrafts, @Nullable BiConsumer<T, I> output) {
         return new Filter(list).tryPickAll(maxCrafts, output);
     }
     
@@ -69,14 +69,14 @@ public class RecipeFinder<T> {
     }
     
     class Filter {
-        private final List<Ingredient<T>> ingredients;
+        private final List<I> ingredients;
         private final int ingredientCount;
         private final List<T> items;
         private final int itemCount;
         private final BitSet data;
         private final IntList path = new IntArrayList();
         
-        public Filter(final List<Ingredient<T>> list) {
+        public Filter(final List<I> list) {
             this.ingredients = list;
             this.ingredientCount = this.ingredients.size();
             this.items = this.getUniqueAvailableIngredientItems();
@@ -87,7 +87,7 @@ public class RecipeFinder<T> {
         
         private void setInitialConnections() {
             for (int i = 0; i < this.ingredientCount; i++) {
-                List<T> list = ((Ingredient<T>) this.ingredients.get(i)).elements();
+                List<T> list = this.ingredients.get(i).elements();
                 
                 for (int j = 0; j < this.itemCount; j++) {
                     if (list.contains(this.items.get(j))) {
@@ -97,7 +97,7 @@ public class RecipeFinder<T> {
             }
         }
         
-        public boolean tryPick(int maxCrafts, @Nullable Consumer<T> output) {
+        public boolean tryPick(int maxCrafts, @Nullable BiConsumer<T, I> output) {
             if (maxCrafts <= 0) {
                 return true;
             } else {
@@ -117,7 +117,7 @@ public class RecipeFinder<T> {
                                     this.unassign(m, l);
                                     put(this.items.get(m), maxCrafts);
                                     if (bl2) {
-                                        output.accept(this.items.get(m));
+                                        output.accept(this.items.get(m), this.ingredients.get(l));
                                     }
                                     break;
                                 }
@@ -362,7 +362,7 @@ public class RecipeFinder<T> {
             this.data.clear(i, i + j);
         }
         
-        public int tryPickAll(int i, @Nullable Consumer<T> output) {
+        public int tryPickAll(int i, @Nullable BiConsumer<T, I> output) {
             int j = 0;
             int k = Math.min(i, this.getMinIngredientCount()) + 1;
             
@@ -403,6 +403,7 @@ public class RecipeFinder<T> {
         }
     }
     
-    public record Ingredient<T>(List<T> elements) {
+    public interface Ingredient<T> {
+        List<T> elements();
     }
 }