Compute the output of smithing displays on the fly

2 files changed, 89 insertions, 21 deletions

diff --git a/default-plugin/src/main/java/me/shedaniel/rei/plugin/client/categories/DefaultSmithingCategory.java b/default-plugin/src/main/java/me/shedaniel/rei/plugin/client/categories/DefaultSmithingCategory.java
index 982221600..8ca82cc14 100644
--- a/default-plugin/src/main/java/me/shedaniel/rei/plugin/client/categories/DefaultSmithingCategory.java
+++ b/default-plugin/src/main/java/me/shedaniel/rei/plugin/client/categories/DefaultSmithingCategory.java
@@ -27,15 +27,22 @@ import com.google.common.collect.Lists;
 import me.shedaniel.math.Point;
 import me.shedaniel.math.Rectangle;
 import me.shedaniel.rei.api.client.gui.Renderer;
+import me.shedaniel.rei.api.client.gui.widgets.Slot;
 import me.shedaniel.rei.api.client.gui.widgets.Widget;
 import me.shedaniel.rei.api.client.gui.widgets.Widgets;
 import me.shedaniel.rei.api.client.registry.display.DisplayCategory;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
+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.util.EntryStacks;
 import me.shedaniel.rei.plugin.common.BuiltinPlugin;
 import me.shedaniel.rei.plugin.common.displays.DefaultSmithingDisplay;
+import net.minecraft.core.RegistryAccess;
 import net.minecraft.network.chat.Component;
 import net.minecraft.world.level.block.Blocks;
+import org.apache.commons.lang3.mutable.MutableBoolean;
+import org.jetbrains.annotations.ApiStatus;
 
 import java.util.List;
 
@@ -65,14 +72,23 @@ public class DefaultSmithingCategory implements DisplayCategory<DefaultSmithingD
         widgets.add(Widgets.createArrow(new Point(startPoint.x + 27 + offsetX, startPoint.y + 4)));
         widgets.add(Widgets.createResultSlotBackground(new Point(startPoint.x + 61 + offsetX, startPoint.y + 5)));
         if (!legacy) {
-            widgets.add(Widgets.createSlot(new Point(startPoint.x + 4 - 18 * 2 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(0)).markInput());
-            widgets.add(Widgets.createSlot(new Point(startPoint.x + 4 - 18 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(1)).markInput());
-            widgets.add(Widgets.createSlot(new Point(startPoint.x + 4 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(2)).markInput());
+            Slot templateSlot, baseSlot, additionSlot, resultSlot;
+            MutableBoolean dirty = new MutableBoolean(true);
+            widgets.add(templateSlot = Widgets.createSlot(new Point(startPoint.x + 4 - 18 * 2 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(0)).withEntriesListener(slot -> dirty.setTrue()).markInput());
+            widgets.add(baseSlot = Widgets.createSlot(new Point(startPoint.x + 4 - 18 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(1)).withEntriesListener(slot -> dirty.setTrue()).markInput());
+            widgets.add(additionSlot = Widgets.createSlot(new Point(startPoint.x + 4 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(2)).withEntriesListener(slot -> dirty.setTrue()).markInput());
+            widgets.add(resultSlot = Widgets.createSlot(new Point(startPoint.x + 61 + offsetX, startPoint.y + 5)).entries(display.getOutputEntries().get(0)).disableBackground().markOutput());
+            widgets.add(Widgets.createDrawableWidget((helper, matrices, mouseX, mouseY, delta) -> {
+                if (dirty.booleanValue()) {
+                    resultSlot.clearEntries().entries(getOutput(display, BasicDisplay.registryAccess(), templateSlot.getCurrentEntry(), baseSlot.getCurrentEntry(), additionSlot.getCurrentEntry()));
+                    dirty.setFalse();
+                }
+            }));
         } else {
             widgets.add(Widgets.createSlot(new Point(startPoint.x + 4 - 22 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(0)).markInput());
             widgets.add(Widgets.createSlot(new Point(startPoint.x + 4 + offsetX, startPoint.y + 5)).entries(display.getInputEntries().get(1)).markInput());
+            widgets.add(Widgets.createSlot(new Point(startPoint.x + 61 + offsetX, startPoint.y + 5)).entries(display.getOutputEntries().get(0)).disableBackground().markOutput());
         }
-        widgets.add(Widgets.createSlot(new Point(startPoint.x + 61 + offsetX, startPoint.y + 5)).entries(display.getOutputEntries().get(0)).disableBackground().markOutput());
         return widgets;
     }
     
@@ -80,4 +96,14 @@ public class DefaultSmithingCategory implements DisplayCategory<DefaultSmithingD
     public int getDisplayHeight() {
         return 36;
     }
+    
+    @ApiStatus.Experimental
+    private static EntryIngredient getOutput(DefaultSmithingDisplay display, RegistryAccess registryAccess, EntryStack<?> template, EntryStack<?> base, EntryStack<?> addition) {
+        if (display.getType() == DefaultSmithingDisplay.SmithingRecipeType.TRIM) {
+            EntryIngredient output = DefaultSmithingDisplay.getTrimmingOutput(registryAccess, template, base, addition);
+            if (!output.isEmpty()) return output;
+        }
+        
+        return display.getOutputEntries().get(0);
+    }
 }
diff --git a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/DefaultSmithingDisplay.java b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/DefaultSmithingDisplay.java
index 5e2b2e5fc..bd38d81c7 100644
--- a/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/DefaultSmithingDisplay.java
+++ b/default-plugin/src/main/java/me/shedaniel/rei/plugin/common/displays/DefaultSmithingDisplay.java
@@ -26,6 +26,8 @@ package me.shedaniel.rei.plugin.common.displays;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 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.type.VanillaEntryTypes;
 import me.shedaniel.rei.api.common.util.EntryIngredients;
 import me.shedaniel.rei.api.common.util.EntryStacks;
 import me.shedaniel.rei.plugin.common.BuiltinPlugin;
@@ -39,12 +41,16 @@ import net.minecraft.world.item.crafting.SmithingRecipe;
 import net.minecraft.world.item.crafting.SmithingTransformRecipe;
 import net.minecraft.world.item.crafting.SmithingTrimRecipe;
 import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
 public class DefaultSmithingDisplay extends BasicDisplay {
+    @Nullable
+    private final SmithingRecipeType type;
+    
     public DefaultSmithingDisplay(LegacyUpgradeRecipe recipe) {
         this(
                 recipe,
@@ -93,25 +99,15 @@ public class DefaultSmithingDisplay extends BasicDisplay {
                         .orElse(null);
                 if (trimMaterial == null) continue;
                 
-                ArmorTrim armorTrim = new ArmorTrim(trimMaterial, trimPattern);
-                EntryIngredient.Builder baseItems = EntryIngredient.builder(), outputItems = EntryIngredient.builder();
-                for (ItemStack item : recipe.base.getItems()) {
-                    Optional<ArmorTrim> trim = ArmorTrim.getTrim(registryAccess, item);
-                    if (trim.isEmpty() || !trim.get().hasPatternAndMaterial(trimPattern, trimMaterial)) {
-                        ItemStack newItem = item.copy();
-                        newItem.setCount(1);
-                        if (ArmorTrim.setTrim(registryAccess, newItem, armorTrim)) {
-                            baseItems.add(EntryStacks.of(item.copy()));
-                            outputItems.add(EntryStacks.of(newItem));
-                        }
-                    }
-                }
+                EntryIngredient baseIngredient = EntryIngredients.ofIngredient(recipe.base);
+                EntryIngredient templateOutput = baseIngredient.isEmpty() ? EntryIngredient.empty()
+                        : getTrimmingOutput(registryAccess, EntryStacks.of(templateItem), baseIngredient.get(0), EntryStacks.of(additionStack));
+                
                 displays.add(new DefaultSmithingDisplay(List.of(
                         EntryIngredients.of(templateItem),
-                        baseItems.build(),
+                        baseIngredient,
                         EntryIngredients.of(additionStack)
-                ), List.of(outputItems.build()),
-                        Optional.ofNullable(recipe.getId())));
+                ), List.of(templateOutput), SmithingRecipeType.TRIM, Optional.ofNullable(recipe.getId())));
             }
         }
         return displays;
@@ -126,7 +122,13 @@ public class DefaultSmithingDisplay extends BasicDisplay {
     }
     
     public DefaultSmithingDisplay(List<EntryIngredient> inputs, List<EntryIngredient> outputs, Optional<ResourceLocation> location) {
+        this(inputs, outputs, null, location);
+    }
+    
+    @ApiStatus.Experimental
+    public DefaultSmithingDisplay(List<EntryIngredient> inputs, List<EntryIngredient> outputs, @Nullable SmithingRecipeType type, Optional<ResourceLocation> location) {
         super(inputs, outputs, location);
+        this.type = type;
     }
     
     @Override
@@ -134,7 +136,47 @@ public class DefaultSmithingDisplay extends BasicDisplay {
         return BuiltinPlugin.SMITHING;
     }
     
+    @ApiStatus.Experimental
+    @Nullable
+    public SmithingRecipeType getType() {
+        return type;
+    }
+    
     public static BasicDisplay.Serializer<DefaultSmithingDisplay> serializer() {
-        return BasicDisplay.Serializer.ofSimple(DefaultSmithingDisplay::new);
+        return BasicDisplay.Serializer.of((input, output, id, tag) -> {
+            SmithingRecipeType type = tag.contains("Type") ? SmithingRecipeType.valueOf(tag.getString("Type")) : null;
+            return new DefaultSmithingDisplay(input, output, type, id);
+        }, (display, tag) -> {
+            if (display.type != null) tag.putString("Type", display.type.name());
+        });
+    }
+    
+    @ApiStatus.Experimental
+    public enum SmithingRecipeType {
+        TRIM,
+        TRANSFORM
+    }
+    
+    @ApiStatus.Experimental
+    @ApiStatus.Internal
+    public static EntryIngredient getTrimmingOutput(RegistryAccess registryAccess, EntryStack<?> template, EntryStack<?> base, EntryStack<?> addition) {
+        if (template.getType() != VanillaEntryTypes.ITEM || base.getType() != VanillaEntryTypes.ITEM || addition.getType() != VanillaEntryTypes.ITEM) return EntryIngredient.empty();
+        ItemStack templateItem = template.castValue();
+        ItemStack baseItem = base.castValue();
+        ItemStack additionItem = addition.castValue();
+        Holder.Reference<TrimPattern> trimPattern = TrimPatterns.getFromTemplate(registryAccess, templateItem)
+                .orElse(null);
+        if (trimPattern == null) return EntryIngredient.empty();
+        Holder.Reference<TrimMaterial> trimMaterial = TrimMaterials.getFromIngredient(registryAccess, additionItem)
+                .orElse(null);
+        if (trimMaterial == null) return EntryIngredient.empty();
+        ArmorTrim armorTrim = new ArmorTrim(trimMaterial, trimPattern);
+        Optional<ArmorTrim> trim = ArmorTrim.getTrim(registryAccess, baseItem);
+        if (trim.isPresent() && trim.get().hasPatternAndMaterial(trimPattern, trimMaterial)) return EntryIngredient.empty();
+        ItemStack newItem = baseItem.copy();
+        newItem.setCount(1);
+        if (ArmorTrim.setTrim(registryAccess, newItem, armorTrim)) {
+            return EntryIngredients.of(newItem);
+        } else return EntryIngredient.empty();
     }
 }