2 files changed, 16 insertions, 10 deletions
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 f27b30c15..90af91175 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
@@ -62,6 +62,7 @@ import net.minecraft.network.chat.Component;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.tags.TagKey;
 import net.minecraft.world.inventory.tooltip.TooltipComponent;
+import net.minecraft.world.item.BlockItem;
 import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemDisplayContext;
 import net.minecraft.world.item.ItemStack;
@@ -216,7 +217,11 @@ public class ItemEntryDefinition implements EntryDefinition<ItemStack>, EntrySer
     
     @Override
     public Stream<? extends TagKey<?>> getTagsFor(EntryStack<ItemStack> entry, ItemStack value) {
-        return value.getTags();
+        Stream<? extends TagKey<?>> tags = value.getTags();
+        if (value.getItem() instanceof BlockItem blockItem) {
+            tags = Stream.concat(tags, blockItem.getBlock().builtInRegistryHolder().tags());
+        }
+        return tags;
     }
     
     @Environment(EnvType.CLIENT)
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/SearchFilterPrepareWatcher.java b/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/SearchFilterPrepareWatcher.java
index ea1db5d72..f2acdcc80 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/SearchFilterPrepareWatcher.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/client/runtime/SearchFilterPrepareWatcher.java
@@ -24,12 +24,12 @@
 package me.shedaniel.rei.plugin.client.runtime;
 
 import com.google.common.collect.ImmutableList;
-import it.unimi.dsi.fastutil.ints.IntIntPair;
 import me.shedaniel.math.Color;
 import me.shedaniel.math.Point;
 import me.shedaniel.rei.api.client.gui.widgets.Tooltip;
 import me.shedaniel.rei.impl.client.gui.hints.HintProvider;
 import me.shedaniel.rei.impl.client.search.argument.Argument;
+import me.shedaniel.rei.impl.client.search.argument.ArgumentCache;
 import net.minecraft.Util;
 import net.minecraft.network.chat.Component;
 import org.apache.commons.lang3.ArrayUtils;
@@ -45,14 +45,15 @@ public class SearchFilterPrepareWatcher implements HintProvider {
     public List<Component> provide() {
         lastProcess = null;
         try {
-            if (Argument.prepareStage != null && Argument.currentStages != null && Argument.prepareStacks != null && Argument.prepareStacks.size() > 100
-                && Argument.prepareStart != null) {
-                if (Util.getEpochMillis() - Argument.prepareStart < 100) return Collections.emptyList();
-                int prepareStageCurrent = Argument.prepareStage.firstInt();
-                int prepareStageTotal = Argument.prepareStage.secondInt();
-                IntIntPair currentStage = ArrayUtils.get(Argument.currentStages, prepareStageCurrent - 1);
-                int currentStageCurrent = currentStage == null ? 0 : currentStage.firstInt();
-                int currentStageTotal = currentStage == null ? 0 : currentStage.secondInt();
+            ArgumentCache cache = Argument.cache;
+            if (cache.currentStep != null && cache.prepareStacks != null && cache.prepareStacks.size() > 100
+                    && cache.prepareStart != null) {
+                if (Util.getEpochMillis() - cache.prepareStart < 100) return Collections.emptyList();
+                int prepareStageCurrent = cache.currentStep.step;
+                int prepareStageTotal = cache.currentStep.totalSteps;
+                ArgumentCache.CurrentStep.Step currentStage = ArrayUtils.get(cache.currentStep.steps, prepareStageCurrent - 1);
+                int currentStageCurrent = currentStage == null ? 0 : currentStage.stacks;
+                int currentStageTotal = currentStage == null ? 0 : currentStage.totalStacks;
                 double prepareStageProgress = prepareStageTotal == 0 ? 0 : prepareStageCurrent / (double) prepareStageTotal;
                 double currentStageProgress = currentStageTotal == 0 ? 0 : currentStageCurrent / (double) currentStageTotal;
                 double lastProcess = prepareStageTotal == 0 ? 0 : Math.max(0, prepareStageProgress - (1 - currentStageProgress) / prepareStageTotal);