Adds EntryDefinition.cheatsAs

3 files changed, 20 insertions, 7 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/FluidEntryDefinition.java b/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/FluidEntryDefinition.java
index afb0f3613..6382f79b3 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/FluidEntryDefinition.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/FluidEntryDefinition.java
@@ -67,6 +67,8 @@ import net.minecraft.tags.TagCollection;
 import net.minecraft.tags.TagContainer;
 import net.minecraft.util.Mth;
 import net.minecraft.world.inventory.InventoryMenu;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.level.material.FlowingFluid;
 import net.minecraft.world.level.material.Fluid;
 import org.jetbrains.annotations.Nullable;
@@ -141,6 +143,15 @@ public class FluidEntryDefinition implements EntryDefinition<FluidStack>, EntryS
     }
     
     @Override
+    @Nullable
+    public ItemStack cheatsAs(EntryStack<FluidStack> entry, FluidStack value) {
+        if (value.isEmpty()) return ItemStack.EMPTY;
+        Item bucket = value.getFluid().getBucket();
+        if (bucket == null) return ItemStack.EMPTY;
+        return new ItemStack(bucket);
+    }
+    
+    @Override
     public long hash(EntryStack<FluidStack> entry, FluidStack value, ComparisonContext context) {
         int code = 1;
         code = 31 * code + value.getFluid().hashCode();
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/ItemEntryDefinition.java b/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/ItemEntryDefinition.java
index 9b34f65d8..ef36c8117 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/ItemEntryDefinition.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/client/entry/ItemEntryDefinition.java
@@ -134,6 +134,12 @@ public class ItemEntryDefinition implements EntryDefinition<ItemStack>, EntrySer
     }
     
     @Override
+    @Nullable
+    public ItemStack cheatsAs(EntryStack<ItemStack> entry, ItemStack value) {
+        return value.copy();
+    }
+    
+    @Override
     public long hash(EntryStack<ItemStack> entry, ItemStack value, ComparisonContext context) {
         int code = 1;
         code = 31 * code + System.identityHashCode(value.getItem());
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/DefaultClientRuntimePlugin.java b/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/DefaultClientRuntimePlugin.java
index 2e25cf0bf..d67d8e99e 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/DefaultClientRuntimePlugin.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/DefaultClientRuntimePlugin.java
@@ -27,7 +27,6 @@ import com.mojang.blaze3d.systems.RenderSystem;
 import com.mojang.blaze3d.vertex.PoseStack;
 import com.mojang.serialization.DataResult;
 import com.mojang.serialization.Lifecycle;
-import dev.architectury.fluid.FluidStack;
 import me.shedaniel.math.Point;
 import me.shedaniel.math.Rectangle;
 import me.shedaniel.rei.RoughlyEnoughItemsCoreClient;
@@ -73,7 +72,6 @@ import net.minecraft.nbt.CompoundTag;
 import net.minecraft.network.chat.TextComponent;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.InteractionResult;
-import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemStack;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
@@ -242,11 +240,9 @@ public class DefaultClientRuntimePlugin implements REIClientPlugin {
             if (!(ClientHelper.getInstance().isCheating() && !(Minecraft.getInstance().screen instanceof DisplayScreen))) return false;
             EntryStack<?> entry = stack.copy();
             if (!entry.isEmpty()) {
-                if (entry.getValueType() == FluidStack.class) {
-                    Item bucketItem = ((FluidStack) entry.getValue()).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.getType() == VanillaEntryTypes.ITEM)
                     entry.<ItemStack>castValue().setCount(button != 1 && !Screen.hasShiftDown() ? 1 : entry.<ItemStack>castValue().getMaxStackSize());