Compute the output of smithing displays on the fly

1 files changed, 59 insertions, 17 deletions

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();
     }
 }