Adds EntryDefinition.cheatsAs

3 files changed, 22 insertions, 20 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListEntryWidget.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListEntryWidget.java
index 4ab46773c..2da340fb0 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListEntryWidget.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListEntryWidget.java
@@ -24,16 +24,15 @@
 package me.shedaniel.rei.impl.client.gui.widget;
 
 import com.mojang.blaze3d.vertex.PoseStack;
-import dev.architectury.fluid.FluidStack;
 import me.shedaniel.math.Point;
 import me.shedaniel.rei.api.client.ClientHelper;
 import me.shedaniel.rei.api.client.config.ConfigObject;
 import me.shedaniel.rei.api.client.gui.screen.DisplayScreen;
 import me.shedaniel.rei.api.common.entry.EntryStack;
+import me.shedaniel.rei.api.common.entry.type.VanillaEntryTypes;
 import me.shedaniel.rei.api.common.util.EntryStacks;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.gui.screens.Screen;
-import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemStack;
 
 public abstract class EntryListEntryWidget extends EntryWidget {
@@ -79,12 +78,9 @@ public abstract class EntryListEntryWidget extends EntryWidget {
         if (!(ClientHelper.getInstance().isCheating() && !(Minecraft.getInstance().screen instanceof DisplayScreen))) return false;
         EntryStack<?> entry = getCurrentEntry().copy();
         if (!entry.isEmpty()) {
-            if (entry.getValueType() == FluidStack.class) {
-                FluidStack value = (FluidStack) entry.getValue();
-                Item bucketItem = value.getFluid().getBucket();
-                if (bucketItem != null) {
-                    entry = EntryStacks.of(bucketItem);
-                }
+            if (entry.getType() != VanillaEntryTypes.ITEM) {
+                EntryStack<ItemStack> cheatsAs = entry.cheatsAs();
+                entry = cheatsAs.isEmpty() ? entry : cheatsAs;
             }
             if (entry.getValueType() == ItemStack.class)
                 entry.<ItemStack>castValue().setCount(button != 1 && !Screen.hasShiftDown() ? 1 : entry.<ItemStack>castValue().getMaxStackSize());
@@ -101,12 +97,9 @@ public abstract class EntryListEntryWidget extends EntryWidget {
         if (ClientHelper.getInstance().isCheating() && !(Minecraft.getInstance().screen instanceof DisplayScreen)) {
             EntryStack<?> entry = getCurrentEntry().copy();
             if (!entry.isEmpty()) {
-                if (entry.getValueType() == FluidStack.class) {
-                    FluidStack value = (FluidStack) entry.getValue();
-                    Item bucketItem = value.getFluid().getBucket();
-                    if (bucketItem != null) {
-                        entry = EntryStacks.of(bucketItem);
-                    }
+                if (entry.getType() != VanillaEntryTypes.ITEM) {
+                    EntryStack<ItemStack> cheatsAs = entry.cheatsAs();
+                    entry = cheatsAs.isEmpty() ? entry : cheatsAs;
                 }
                 return EntryStacks.equalsExact(entry, stack);
             }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListWidget.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListWidget.java
index 2faf558d6..4796da9d5 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListWidget.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/EntryListWidget.java
@@ -29,7 +29,6 @@ import com.google.common.collect.Lists;
 import com.mojang.blaze3d.vertex.PoseStack;
 import com.mojang.blaze3d.vertex.Tesselator;
 import com.mojang.math.Matrix4f;
-import dev.architectury.fluid.FluidStack;
 import me.shedaniel.clothconfig2.ClothConfigInitializer;
 import me.shedaniel.clothconfig2.api.ScissorsHandler;
 import me.shedaniel.clothconfig2.api.ScrollingContainer;
@@ -323,11 +322,9 @@ public class EntryListWidget extends WidgetWithBounds implements OverlayListWidg
         
         if (containsChecked(mouseX, mouseY) && ClientHelper.getInstance().isCheating() && !(Minecraft.getInstance().screen instanceof DisplayScreen) && !minecraft.player.containerMenu.getCarried().isEmpty() && ClientHelperImpl.getInstance().canDeleteItems()) {
             EntryStack<?> stack = EntryStacks.of(minecraft.player.containerMenu.getCarried().copy());
-            if (stack.getValueType() == FluidStack.class) {
-                Item bucketItem = ((FluidStack) stack.getValue()).getFluid().getBucket();
-                if (bucketItem != null) {
-                    stack = EntryStacks.of(bucketItem);
-                }
+            if (stack.getType() != VanillaEntryTypes.ITEM) {
+                EntryStack<ItemStack> cheatsAs = stack.cheatsAs();
+                stack = cheatsAs.isEmpty() ? stack : cheatsAs;
             }
             for (Widget child : children()) {
                 if (child.containsMouse(mouseX, mouseY) && child instanceof EntryWidget widget) {
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/AbstractEntryStack.java b/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/AbstractEntryStack.java
index c31b54cf9..b34530db2 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/AbstractEntryStack.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/AbstractEntryStack.java
@@ -48,6 +48,7 @@ import net.minecraft.network.chat.Component;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.tags.SerializationTags;
 import net.minecraft.tags.TagContainer;
+import net.minecraft.world.item.ItemStack;
 import org.apache.commons.lang3.mutable.Mutable;
 import org.apache.commons.lang3.mutable.MutableObject;
 import org.jetbrains.annotations.ApiStatus;
@@ -158,6 +159,17 @@ public abstract class AbstractEntryStack<A> implements EntryStack<A>, Renderer {
         return wrap(getDefinition().wildcard(this, getValue()), false);
     }
     
+    @Override
+    public EntryStack<ItemStack> cheatsAs() {
+        ItemStack stack = getDefinition().cheatsAs(this, getValue());
+        
+        if (stack == null) {
+            return null;
+        }
+        
+        return EntryStacks.of(stack);
+    }
+    
     protected EntryStack<A> wrap(A value, boolean copySettings) {
         TypedEntryStack<A> stack = new TypedEntryStack<>(getDefinition(), value);
         if (copySettings) {