Fix 2x2 crafting

7 files changed, 193 insertions, 24 deletions

diff --git a/api/src/main/java/me/shedaniel/rei/api/common/display/Display.java b/api/src/main/java/me/shedaniel/rei/api/common/display/Display.java
index 82505f971..bf6dc0cef 100644
--- a/api/src/main/java/me/shedaniel/rei/api/common/display/Display.java
+++ b/api/src/main/java/me/shedaniel/rei/api/common/display/Display.java
@@ -25,8 +25,11 @@ package me.shedaniel.rei.api.common.display;
 
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.entry.EntryIngredient;
+import me.shedaniel.rei.api.common.entry.EntryStack;
+import me.shedaniel.rei.api.common.entry.InputIngredient;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfo;
 import me.shedaniel.rei.api.common.transfer.info.MenuSerializationContext;
+import me.shedaniel.rei.api.common.util.CollectionUtils;
 import me.shedaniel.rei.impl.display.DisplaySpec;
 import net.minecraft.resources.ResourceLocation;
 import org.jetbrains.annotations.ApiStatus;
@@ -49,10 +52,15 @@ public interface Display extends DisplaySpec {
      */
     List<EntryIngredient> getInputEntries();
     
+    @Deprecated
     default List<EntryIngredient> getInputEntries(MenuSerializationContext<?, ?, ?> context, MenuInfo<?, ?> info, boolean fill) {
         return getInputEntries();
     }
     
+    default List<InputIngredient<EntryStack<?>>> getInputIngredients(MenuSerializationContext<?, ?, ?> context, MenuInfo<?, ?> info, boolean fill) {
+        return CollectionUtils.mapIndexed(getInputEntries(context, info, fill), InputIngredient::of);
+    }
+    
     /**
      * @return a list of outputs
      */
diff --git a/api/src/main/java/me/shedaniel/rei/api/common/entry/InputIngredient.java b/api/src/main/java/me/shedaniel/rei/api/common/entry/InputIngredient.java
new file mode 100644
index 000000000..99dd3dcaa
--- /dev/null
+++ b/api/src/main/java/me/shedaniel/rei/api/common/entry/InputIngredient.java
@@ -0,0 +1,72 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021, 2022 shedaniel
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package me.shedaniel.rei.api.common.entry;
+
+import me.shedaniel.rei.api.common.entry.type.EntryType;
+import me.shedaniel.rei.api.common.util.CollectionUtils;
+
+import java.util.Collections;
+import java.util.List;
+
+public interface InputIngredient<T> {
+    static <T> InputIngredient<T> empty(int index) {
+        return of(index, Collections.emptyList());
+    }
+    
+    static <T> InputIngredient<T> of(int index, List<T> ingredient) {
+        return new InputIngredient<T>() {
+            @Override
+            public List<T> get() {
+                return ingredient;
+            }
+            
+            @Override
+            public int getIndex() {
+                return index;
+            }
+        };
+    }
+    
+    static <T> InputIngredient<T> withType(InputIngredient<EntryStack<?>> ingredient, EntryType<T> type) {
+        return new InputIngredient<T>() {
+            @SuppressWarnings("RedundantTypeArguments")
+            List<T> list = CollectionUtils.<EntryStack<?>, T>filterAndMap(ingredient.get(),
+                    stack -> stack.getType() == type, EntryStack::castValue);
+            
+            @Override
+            public List<T> get() {
+                return list;
+            }
+            
+            @Override
+            public int getIndex() {
+                return ingredient.getIndex();
+            }
+        };
+    }
+    
+    List<T> get();
+    
+    int getIndex();
+}
diff --git a/api/src/main/java/me/shedaniel/rei/api/common/transfer/info/MenuInfo.java b/api/src/main/java/me/shedaniel/rei/api/common/transfer/info/MenuInfo.java
index ad9a0e3cc..ca3b8ba74 100644
--- a/api/src/main/java/me/shedaniel/rei/api/common/transfer/info/MenuInfo.java
+++ b/api/src/main/java/me/shedaniel/rei/api/common/transfer/info/MenuInfo.java
@@ -25,6 +25,7 @@ package me.shedaniel.rei.api.common.transfer.info;
 
 import com.mojang.blaze3d.vertex.PoseStack;
 import it.unimi.dsi.fastutil.ints.IntList;
+import it.unimi.dsi.fastutil.ints.IntSet;
 import me.shedaniel.math.Rectangle;
 import me.shedaniel.rei.api.client.gui.widgets.Slot;
 import me.shedaniel.rei.api.client.gui.widgets.Widget;
@@ -32,6 +33,7 @@ import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.Display;
 import me.shedaniel.rei.api.common.display.DisplaySerializerRegistry;
 import me.shedaniel.rei.api.common.entry.EntryStack;
+import me.shedaniel.rei.api.common.entry.InputIngredient;
 import me.shedaniel.rei.api.common.entry.type.VanillaEntryTypes;
 import me.shedaniel.rei.api.common.transfer.RecipeFinder;
 import me.shedaniel.rei.api.common.transfer.RecipeFinderPopulator;
@@ -133,10 +135,12 @@ public interface MenuInfo<T extends AbstractContainerMenu, D extends Display> ex
      *                otherwise it should be aligned for the display category
      * @return the list of lists of items
      */
+    @Deprecated
+    @ApiStatus.ScheduledForRemoval
     default List<List<ItemStack>> getInputs(MenuInfoContext<T, ?, D> context, boolean fill) {
         if (context.getDisplay() == null) return Collections.emptyList();
-        return CollectionUtils.map(context.getDisplay().getInputEntries(context, this, fill), inputEntry ->
-                CollectionUtils.<EntryStack<?>, ItemStack>filterAndMap(inputEntry,
+        return CollectionUtils.map(context.getDisplay().getInputIngredients(context, this, fill), inputEntry ->
+                CollectionUtils.<EntryStack<?>, ItemStack>filterAndMap(inputEntry.get(),
                         stack -> stack.getType() == VanillaEntryTypes.ITEM, EntryStack::castValue));
     }
     
@@ -153,6 +157,20 @@ public interface MenuInfo<T extends AbstractContainerMenu, D extends Display> ex
     }
     
     /**
+     * Returns the inputs of the {@link Display}. The nested lists are possible stacks for that specific slot.
+     *
+     * @param context the context of the transfer
+     * @param fill    whether this call is for a fill or not, if it is for a fill, the returned list should be aligned for the menu,
+     *                otherwise it should be aligned for the display category
+     * @return the list of lists of items
+     */
+    default List<InputIngredient<ItemStack>> getInputsIndexed(MenuInfoContext<T, ?, D> context, boolean fill) {
+        if (context.getDisplay() == null) return Collections.emptyList();
+        return CollectionUtils.map(context.getDisplay().getInputIngredients(context, this, fill), entry ->
+                InputIngredient.withType(entry, VanillaEntryTypes.ITEM));
+    }
+    
+    /**
      * Serializes the {@link Display} as {@link CompoundTag}, sent to the server for further info for the transfer.
      *
      * @param context the context of the transfer
@@ -181,6 +199,7 @@ public interface MenuInfo<T extends AbstractContainerMenu, D extends Display> ex
      *
      * @param context        the context of the transfer
      * @param inputs         the list of inputs
+     * @param missing        the list of missing stacks
      * @param missingIndices the indices of the missing stacks
      * @param matrices       the rendering transforming matrices
      * @param mouseX         the mouse x position
@@ -190,7 +209,7 @@ public interface MenuInfo<T extends AbstractContainerMenu, D extends Display> ex
      * @param bounds         the bounds of the display
      */
     @Environment(EnvType.CLIENT)
-    default void renderMissingInput(MenuInfoContext<T, ?, D> context, List<List<ItemStack>> inputs, IntList missingIndices, PoseStack matrices, int mouseX, int mouseY,
+    default void renderMissingInput(MenuInfoContext<T, ?, D> context, List<InputIngredient<ItemStack>> inputs, List<InputIngredient<ItemStack>> missing, IntSet missingIndices, PoseStack matrices, int mouseX, int mouseY,
             float delta, List<Widget> widgets, Rectangle bounds) {
         int i = 0;
         for (Widget widget : widgets) {
@@ -205,4 +224,25 @@ public interface MenuInfo<T extends AbstractContainerMenu, D extends Display> ex
             }
         }
     }
+    
+    /**
+     * Renders the missing ingredients of the transfer.
+     * The indices of the missing stacks are provided, this aligns with the list returned by {@link #getInputs(MenuInfoContext, boolean)}.
+     *
+     * @param context        the context of the transfer
+     * @param inputs         the list of inputs
+     * @param missingIndices the indices of the missing stacks
+     * @param matrices       the rendering transforming matrices
+     * @param mouseX         the mouse x position
+     * @param mouseY         the mouse y position
+     * @param delta          the delta frame time
+     * @param widgets        the widgets set-up by the category
+     * @param bounds         the bounds of the display
+     */
+    @Environment(EnvType.CLIENT)
+    @Deprecated
+    @ApiStatus.ScheduledForRemoval
+    default void renderMissingInput(MenuInfoContext<T, ?, D> context, List<List<ItemStack>> inputs, IntList missingIndices, PoseStack matrices, int mouseX, int mouseY,
+            float delta, List<Widget> widgets, Rectangle bounds) {
+    }
 }
diff --git a/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java b/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java
index 6efbc1901..7d00c76ec 100644
--- a/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java
+++ b/api/src/main/java/me/shedaniel/rei/api/common/util/CollectionUtils.java
@@ -130,6 +130,20 @@ public class CollectionUtils {
         return l;
     }
     
+    public static <T, R> List<R> mapIndexed(Iterable<T> list, IndexedFunction<T, R> function) {
+        List<R> l = list instanceof Collection ? new ArrayList<>(((Collection<T>) list).size() + 1) : new ArrayList<>();
+        int i = 0;
+        for (T t : list) {
+            l.add(function.apply(i++, t));
+        }
+        return l;
+    }
+    
+    @FunctionalInterface
+    public interface IndexedFunction<T, R> {
+        R apply(int index, T object);
+    }
+    
     public static <T, R> List<R> flatMap(Iterable<T> list, Function<T, Collection<R>> function) {
         List<R> l = new ArrayList<>();
         for (T t : list) {
diff --git a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCraftingDisplay.java b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCraftingDisplay.java
index adc046c0e..cf5ce260f 100644
--- a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCraftingDisplay.java
+++ b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCraftingDisplay.java
@@ -31,6 +31,7 @@ import me.shedaniel.rei.api.common.display.SimpleGridMenuDisplay;
 import me.shedaniel.rei.api.common.display.basic.BasicDisplay;
 import me.shedaniel.rei.api.common.entry.EntryIngredient;
 import me.shedaniel.rei.api.common.entry.EntryStack;
+import me.shedaniel.rei.api.common.entry.InputIngredient;
 import me.shedaniel.rei.api.common.entry.type.VanillaEntryTypes;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfo;
 import me.shedaniel.rei.api.common.transfer.info.MenuSerializationContext;
@@ -140,31 +141,33 @@ public abstract class DefaultCraftingDisplay<C extends Recipe<?>> extends BasicD
     }
     
     @Override
-    public List<EntryIngredient> getInputEntries(MenuSerializationContext<?, ?, ?> context, MenuInfo<?, ?> info, boolean fill) {
-        int craftingWidth = Math.max(3, getInputWidth());
-        int craftingHeight = Math.max(3, getInputHeight());
+    public List<InputIngredient<EntryStack<?>>> getInputIngredients(MenuSerializationContext<?, ?, ?> context, MenuInfo<?, ?> info, boolean fill) {
+        int inputWidth = Math.max(3, getInputWidth());
+        int inputHeight = Math.max(3, getInputHeight());
+        int craftingWidth = 3, craftingHeight = 3;
         
-        if (info instanceof SimpleGridMenuInfo) {
+        if (info instanceof SimpleGridMenuInfo && fill) {
             craftingWidth = ((SimpleGridMenuInfo<AbstractContainerMenu, ?>) info).getCraftingWidth(context.getMenu());
             craftingHeight = ((SimpleGridMenuInfo<AbstractContainerMenu, ?>) info).getCraftingHeight(context.getMenu());
         }
         
-        EntryIngredient[][] grid = new EntryIngredient[craftingWidth][craftingHeight];
+        InputIngredient<EntryStack<?>>[][] grid = new InputIngredient[Math.max(inputWidth, craftingWidth)][Math.max(inputHeight, craftingHeight)];
         
         List<EntryIngredient> inputEntries = getInputEntries();
         for (int i = 0; i < inputEntries.size(); i++) {
-            grid[i % getInputWidth()][i / getInputWidth()] = inputEntries.get(i);
+            grid[i % getInputWidth()][i / getInputWidth()] = InputIngredient.of(getSlotWithSize(getInputWidth(), i, craftingWidth), inputEntries.get(i));
         }
         
-        List<EntryIngredient> list = new ArrayList<>(craftingWidth * craftingHeight);
+        List<InputIngredient<EntryStack<?>>> list = new ArrayList<>(craftingWidth * craftingHeight);
         for (int i = 0, n = craftingWidth * craftingHeight; i < n; i++) {
-            list.add(EntryIngredient.empty());
+            list.add(InputIngredient.empty(i));
         }
         
         for (int x = 0; x < craftingWidth; x++) {
             for (int y = 0; y < craftingHeight; y++) {
                 if (grid[x][y] != null) {
-                    list.set(craftingWidth * y + x, grid[x][y]);
+                    int index = craftingWidth * y + x;
+                    list.set(index, grid[x][y]);
                 }
             }
         }
diff --git a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultShapelessDisplay.java b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultShapelessDisplay.java
index 7634edfdf..502c5c61b 100644
--- a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultShapelessDisplay.java
+++ b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultShapelessDisplay.java
@@ -49,4 +49,16 @@ public class DefaultShapelessDisplay extends DefaultCraftingDisplay<ShapelessRec
     public int getHeight() {
         return recipe.get().getIngredients().size() > 4 ? 3 : 2;
     }
+    
+    @SuppressWarnings("OptionalGetWithoutIsPresent")
+    @Override
+    public int getInputWidth() {
+        return Math.min(recipe.get().getIngredients().size(), 3);
+    }
+    
+    @SuppressWarnings("OptionalGetWithoutIsPresent")
+    @Override
+    public int getInputHeight() {
+        return (int) Math.ceil(recipe.get().getIngredients().size() / 3.0);
+    }
 }
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java b/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java
index d37872ff1..e428812e3 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java
@@ -25,22 +25,25 @@ package me.shedaniel.rei.plugin.autocrafting;
 
 import io.netty.buffer.Unpooled;
 import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.ints.IntLinkedOpenHashSet;
 import it.unimi.dsi.fastutil.ints.IntList;
+import it.unimi.dsi.fastutil.ints.IntSet;
 import me.shedaniel.architectury.networking.NetworkManager;
 import me.shedaniel.rei.RoughlyEnoughItemsNetwork;
 import me.shedaniel.rei.api.client.ClientHelper;
 import me.shedaniel.rei.api.client.registry.transfer.TransferHandler;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.Display;
+import me.shedaniel.rei.api.common.entry.InputIngredient;
 import me.shedaniel.rei.api.common.transfer.RecipeFinder;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfo;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfoContext;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfoRegistry;
 import me.shedaniel.rei.api.common.transfer.info.MenuTransferException;
+import me.shedaniel.rei.api.common.util.CollectionUtils;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.minecraft.client.Minecraft;
-import net.minecraft.client.gui.screens.Screen;
 import net.minecraft.client.gui.screens.inventory.AbstractContainerScreen;
 import net.minecraft.client.gui.screens.recipebook.RecipeUpdateListener;
 import net.minecraft.network.FriendlyByteBuf;
@@ -49,6 +52,7 @@ import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.inventory.AbstractContainerMenu;
 import net.minecraft.world.item.ItemStack;
 
+import java.util.ArrayList;
 import java.util.List;
 
 @Environment(EnvType.CLIENT)
@@ -75,12 +79,19 @@ public class DefaultCategoryHandler implements TransferHandler {
                 return Result.createNotApplicable();
             }
         }
-        List<List<ItemStack>> input = menuInfo.getInputs(menuInfoContext, false);
-        IntList intList = hasItems(menu, menuInfo, display, input);
-        if (!intList.isEmpty()) {
+        List<InputIngredient<ItemStack>> input = menuInfo.getInputsIndexed(menuInfoContext, false);
+        List<InputIngredient<ItemStack>> missing = hasItemsIndexed(menu, menuInfo, display, input);
+        if (!missing.isEmpty()) {
+            IntList missingIndices = new IntArrayList(missing.size());
+            for (InputIngredient<ItemStack> ingredient : missing) {
+                missingIndices.add(ingredient.getIndex());
+            }
+            IntSet missingIndicesSet = new IntLinkedOpenHashSet(missingIndices);
+            List<List<ItemStack>> oldInputs = CollectionUtils.map(input, InputIngredient::get);
             return Result.createFailed(new TranslatableComponent("error.rei.not.enough.materials"))
                     .renderer((matrices, mouseX, mouseY, delta, widgets, bounds, d) -> {
-                        menuInfo.renderMissingInput(menuInfoContext, input, intList, matrices, mouseX, mouseY, delta, widgets, bounds);
+                        menuInfo.renderMissingInput(menuInfoContext, oldInputs, missingIndices, matrices, mouseX, mouseY, delta, widgets, bounds);
+                        menuInfo.renderMissingInput(menuInfoContext, input, missing, missingIndicesSet, matrices, mouseX, mouseY, delta, widgets, bounds);
                     });
         }
         if (!ClientHelper.getInstance().canUseMovePackets()) {
@@ -138,14 +149,23 @@ public class DefaultCategoryHandler implements TransferHandler {
     }
     
     public IntList hasItems(AbstractContainerMenu menu, MenuInfo<AbstractContainerMenu, Display> info, Display display, List<List<ItemStack>> inputs) {
+        List<InputIngredient<ItemStack>> missing = hasItemsIndexed(menu, info, display,
+                CollectionUtils.mapIndexed(inputs, InputIngredient::of));
+        IntList ids = new IntArrayList(missing.size());
+        for (InputIngredient<ItemStack> ingredient : missing) {
+            ids.add(ingredient.getIndex());
+        }
+        return ids;
+    }
+    
+    public List<InputIngredient<ItemStack>> hasItemsIndexed(AbstractContainerMenu menu, MenuInfo<AbstractContainerMenu, Display> info, Display display, List<InputIngredient<ItemStack>> inputs) {
         // Create a clone of player's inventory, and count
         RecipeFinder recipeFinder = new RecipeFinder();
         info.getRecipeFinderPopulator().populate(ofContext(menu, info, display), recipeFinder);
-        IntList intList = new IntArrayList();
-        for (int i = 0; i < inputs.size(); i++) {
-            List<ItemStack> possibleStacks = inputs.get(i);
-            boolean done = possibleStacks.isEmpty();
-            for (ItemStack possibleStack : possibleStacks) {
+        List<InputIngredient<ItemStack>> missing = new ArrayList<>();
+        for (InputIngredient<ItemStack> possibleStacks : inputs) {
+            boolean done = possibleStacks.get().isEmpty();
+            for (ItemStack possibleStack : possibleStacks.get()) {
                 if (!done) {
                     int invRequiredCount = possibleStack.getCount();
                     int key = RecipeFinder.getItemId(possibleStack);
@@ -160,9 +180,9 @@ public class DefaultCategoryHandler implements TransferHandler {
                 }
             }
             if (!done) {
-                intList.add(i);
+                missing.add(possibleStacks);
             }
         }
-        return intList;
+        return missing;
     }
 }