Fix #977

4 files changed, 139 insertions, 21 deletions

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 49be00f02..ba64c1529 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
@@ -24,6 +24,8 @@
 package me.shedaniel.rei.plugin.common.displays.crafting;
 
 import dev.architectury.injectables.annotations.ExpectPlatform;
+import it.unimi.dsi.fastutil.ints.IntIntImmutablePair;
+import it.unimi.dsi.fastutil.ints.IntIntPair;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.SimpleGridMenuDisplay;
 import me.shedaniel.rei.api.common.display.basic.BasicDisplay;
@@ -47,10 +49,7 @@ import net.minecraft.world.item.crafting.ShapedRecipe;
 import net.minecraft.world.item.crafting.ShapelessRecipe;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Optional;
+import java.util.*;
 
 public abstract class DefaultCraftingDisplay<C extends Recipe<?>> extends BasicDisplay implements SimpleGridMenuDisplay {
     protected Optional<C> recipe;
@@ -148,7 +147,7 @@ public abstract class DefaultCraftingDisplay<C extends Recipe<?>> extends BasicD
     }
     
     public static int getSlotWithSize(DefaultCraftingDisplay<?> display, int index, int craftingGridWidth) {
-        return getSlotWithSize(display.getInputWidth(), index, craftingGridWidth);
+        return getSlotWithSize(display.getInputWidth(craftingGridWidth, 3), index, craftingGridWidth);
     }
     
     public static int getSlotWithSize(int recipeWidth, int index, int craftingGridWidth) {
@@ -158,8 +157,24 @@ public abstract class DefaultCraftingDisplay<C extends Recipe<?>> extends BasicD
     }
     
     public static BasicDisplay.Serializer<DefaultCraftingDisplay<?>> serializer() {
-        return BasicDisplay.Serializer.<DefaultCraftingDisplay<?>>ofSimple(DefaultCustomDisplay::simple)
-                .inputProvider(display -> display.getOrganisedInputEntries(3, 3));
+        return BasicDisplay.Serializer.<DefaultCraftingDisplay<?>>of((input, output, location, tag) -> {
+            if (tag.contains("REIRecipeType")) {
+                String type = tag.getString("REIRecipeType");
+                return switch (type) {
+                    case "Shapeless" -> DefaultCustomShapelessDisplay.simple(input, output, location);
+                    case "Shaped" -> DefaultCustomShapedDisplay.simple(input, output, tag.getInt("RecipeWidth"), tag.getInt("RecipeHeight"), location);
+                    default -> throw new IllegalArgumentException("Unknown recipe type: " + type);
+                };
+            } else {
+                return DefaultCustomDisplay.simple(input, output, location);
+            }
+        }, (display, tag) -> {
+            tag.putString("REIRecipeType", display.isShapeless() ? "Shapeless" : "Shaped");
+            if (!display.isShapeless()) {
+                tag.putInt("RecipeWidth", display.getInputWidth(3, 3));
+                tag.putInt("RecipeHeight", display.getInputHeight(3, 3));
+            }
+        });
     }
     
     @Override
@@ -175,14 +190,19 @@ public abstract class DefaultCraftingDisplay<C extends Recipe<?>> extends BasicD
     }
     
     public List<InputIngredient<EntryStack<?>>> getInputIngredients(int craftingWidth, int craftingHeight) {
-        int inputWidth = Math.max(3, getInputWidth());
-        int inputHeight = Math.max(3, getInputHeight());
+        int inputWidth = getInputWidth(craftingWidth, craftingHeight);
+        int inputHeight = getInputHeight(craftingWidth, craftingHeight);
         
-        InputIngredient<EntryStack<?>>[][] grid = new InputIngredient[Math.max(inputWidth, craftingWidth)][Math.max(inputHeight, craftingHeight)];
+        Map<IntIntPair, InputIngredient<EntryStack<?>>> grid = new HashMap<>();
         
         List<EntryIngredient> inputEntries = getInputEntries();
         for (int i = 0; i < inputEntries.size(); i++) {
-            grid[i % getInputWidth()][i / getInputWidth()] = InputIngredient.of(getSlotWithSize(getInputWidth(), i, craftingWidth), inputEntries.get(i));
+            EntryIngredient stacks = inputEntries.get(i);
+            if (stacks.isEmpty()) {
+                continue;
+            }
+            int index = getSlotWithSize(inputWidth, i, craftingWidth);
+            grid.put(new IntIntImmutablePair(i % inputWidth, i / inputWidth), InputIngredient.of(index, stacks));
         }
         
         List<InputIngredient<EntryStack<?>>> list = new ArrayList<>(craftingWidth * craftingHeight);
@@ -192,9 +212,10 @@ public abstract class DefaultCraftingDisplay<C extends Recipe<?>> extends BasicD
         
         for (int x = 0; x < craftingWidth; x++) {
             for (int y = 0; y < craftingHeight; y++) {
-                if (grid[x][y] != null) {
+                InputIngredient<EntryStack<?>> ingredient = grid.get(new IntIntImmutablePair(x, y));
+                if (ingredient != null) {
                     int index = craftingWidth * y + x;
-                    list.set(index, grid[x][y]);
+                    list.set(index, ingredient);
                 }
             }
         }
diff --git a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapedDisplay.java b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapedDisplay.java
index a8bb9f9b9..547664010 100644
--- a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapedDisplay.java
+++ b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapedDisplay.java
@@ -24,6 +24,7 @@
 package me.shedaniel.rei.plugin.common.displays.crafting;
 
 import me.shedaniel.rei.api.common.entry.EntryIngredient;
+import me.shedaniel.rei.api.common.registry.RecipeManagerContext;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.item.crafting.Recipe;
 import org.jetbrains.annotations.Nullable;
@@ -45,6 +46,12 @@ public class DefaultCustomShapedDisplay extends DefaultCraftingDisplay<Recipe<?>
         this.height = height;
     }
     
+    public static DefaultCustomShapedDisplay simple(List<EntryIngredient> input, List<EntryIngredient> output, int width, int height, Optional<ResourceLocation> location) {
+        Recipe<?> optionalRecipe = location.flatMap(resourceLocation -> RecipeManagerContext.getInstance().getRecipeManager().byKey(resourceLocation))
+                .orElse(null);
+        return new DefaultCustomShapedDisplay(location.orElse(null), optionalRecipe, input, output, width, height);
+    }
+    
     @Override
     public int getWidth() {
         return width;
diff --git a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapelessDisplay.java b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapelessDisplay.java
new file mode 100644
index 000000000..243220367
--- /dev/null
+++ b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/crafting/DefaultCustomShapelessDisplay.java
@@ -0,0 +1,84 @@
+/*
+ * 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.plugin.common.displays.crafting;
+
+import me.shedaniel.rei.api.common.entry.EntryIngredient;
+import me.shedaniel.rei.api.common.registry.RecipeManagerContext;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.world.item.crafting.Recipe;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.List;
+import java.util.Optional;
+
+public class DefaultCustomShapelessDisplay extends DefaultCraftingDisplay<Recipe<?>> {
+    public DefaultCustomShapelessDisplay(@Nullable Recipe<?> possibleRecipe, List<EntryIngredient> input, List<EntryIngredient> output) {
+        this(null, possibleRecipe, input, output);
+    }
+    
+    public DefaultCustomShapelessDisplay(@Nullable ResourceLocation location, @Nullable Recipe<?> possibleRecipe, List<EntryIngredient> input, List<EntryIngredient> output) {
+        super(input, output, Optional.ofNullable(location == null && possibleRecipe != null ? possibleRecipe.getId() : location), Optional.ofNullable(possibleRecipe));
+    }
+    
+    public static DefaultCustomShapelessDisplay simple(List<EntryIngredient> input, List<EntryIngredient> output, Optional<ResourceLocation> location) {
+        Recipe<?> optionalRecipe = location.flatMap(resourceLocation -> RecipeManagerContext.getInstance().getRecipeManager().byKey(resourceLocation))
+                .orElse(null);
+        return new DefaultCustomShapelessDisplay(location.orElse(null), optionalRecipe, input, output);
+    }
+    
+    @Override
+    public int getWidth() {
+        return getInputEntries().size() > 4 ? 3 : 2;
+    }
+    
+    @Override
+    public int getHeight() {
+        return getInputEntries().size() > 4 ? 3 : 2;
+    }
+    
+    @Override
+    public int getInputWidth() {
+        return Math.min(getInputEntries().size(), 3);
+    }
+    
+    @Override
+    public int getInputWidth(int craftingWidth, int craftingHeight) {
+        return craftingWidth * craftingHeight <= getInputEntries().size() ? craftingWidth : Math.min(getInputEntries().size(), 3);
+    }
+    
+    @Override
+    public int getInputHeight() {
+        return (int) Math.ceil(getInputEntries().size() / (double) getInputWidth());
+    }
+    
+    @Override
+    public int getInputHeight(int craftingWidth, int craftingHeight) {
+        return (int) Math.ceil(getInputEntries().size() / (double) getInputWidth(craftingWidth, craftingHeight));
+    }
+    
+    @Override
+    public boolean isShapeless() {
+        return true;
+    }
+}
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 8f500bfb9..e5234ed62 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
@@ -38,28 +38,34 @@ public class DefaultShapelessDisplay extends DefaultCraftingDisplay<ShapelessRec
         );
     }
     
-    @SuppressWarnings("OptionalGetWithoutIsPresent")
     @Override
     public int getWidth() {
-        return recipe.get().getIngredients().size() > 4 ? 3 : 2;
+        return getInputEntries().size() > 4 ? 3 : 2;
     }
     
-    @SuppressWarnings("OptionalGetWithoutIsPresent")
     @Override
     public int getHeight() {
-        return recipe.get().getIngredients().size() > 4 ? 3 : 2;
+        return getInputEntries().size() > 4 ? 3 : 2;
     }
     
-    @SuppressWarnings("OptionalGetWithoutIsPresent")
     @Override
     public int getInputWidth() {
-        return Math.min(recipe.get().getIngredients().size(), 3);
+        return Math.min(getInputEntries().size(), 3);
+    }
+    
+    @Override
+    public int getInputWidth(int craftingWidth, int craftingHeight) {
+        return craftingWidth * craftingHeight <= getInputEntries().size() ? craftingWidth : Math.min(getInputEntries().size(), 3);
     }
     
-    @SuppressWarnings("OptionalGetWithoutIsPresent")
     @Override
     public int getInputHeight() {
-        return (int) Math.ceil(recipe.get().getIngredients().size() / 3.0);
+        return (int) Math.ceil(getInputEntries().size() / (double) getInputWidth());
+    }
+    
+    @Override
+    public int getInputHeight(int craftingWidth, int craftingHeight) {
+        return (int) Math.ceil(getInputEntries().size() / (double) getInputWidth(craftingWidth, craftingHeight));
     }
     
     @Override