Cache search results

Signed-off-by: shedaniel <daniel@shedaniel.me>

4 files changed, 33 insertions, 12 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/FilteringRule.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/FilteringRule.java
index dd8cd6340..69bfaca27 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/FilteringRule.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/FilteringRule.java
@@ -29,6 +29,7 @@ import me.shedaniel.rei.impl.client.entry.filtering.rules.ManualFilteringRule;
 import me.shedaniel.rei.impl.client.entry.filtering.rules.SearchFilteringRule;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
+import net.minecraft.Util;
 import net.minecraft.client.gui.screens.Screen;
 import net.minecraft.core.MappedRegistry;
 import net.minecraft.core.Registry;
@@ -46,15 +47,10 @@ import java.util.function.BiFunction;
 @Environment(EnvType.CLIENT)
 public interface FilteringRule<T extends FilteringRule<?>> {
     ResourceKey<Registry<FilteringRule<?>>> REGISTRY_KEY = ResourceKey.createRegistryKey(new ResourceLocation("roughlyenoughitems", "filtering_rule"));
-    Registry<FilteringRule<?>> REGISTRY = createRegistry();
-    
-    @ApiStatus.Internal
-    static Registry<FilteringRule<?>> createRegistry() {
-        MappedRegistry<FilteringRule<?>> registry = new MappedRegistry<>(REGISTRY_KEY, Lifecycle.stable());
+    Registry<FilteringRule<?>> REGISTRY = Util.make(new MappedRegistry<>(REGISTRY_KEY, Lifecycle.stable()), registry -> {
         Registry.register(registry, new ResourceLocation("roughlyenoughitems", "search"), new SearchFilteringRule());
         Registry.register(registry, new ResourceLocation("roughlyenoughitems", "manual"), new ManualFilteringRule());
-        return registry;
-    }
+    });
     
     static CompoundTag save(FilteringRule<?> rule, CompoundTag tag) {
         tag.putString("id", REGISTRY.getKey(rule).toString());
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/ContainerScreenOverlay.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/ContainerScreenOverlay.java
index 8abd4ea35..80881a3da 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/ContainerScreenOverlay.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/ContainerScreenOverlay.java
@@ -68,6 +68,7 @@ import me.shedaniel.rei.impl.client.gui.widget.FavoritesListWidget;
 import me.shedaniel.rei.impl.client.gui.widget.InternalWidgets;
 import me.shedaniel.rei.impl.client.gui.widget.LateRenderable;
 import me.shedaniel.rei.impl.client.gui.widget.search.OverlaySearchField;
+import me.shedaniel.rei.impl.client.search.argument.Argument;
 import me.shedaniel.rei.impl.common.util.Weather;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.gui.chat.NarratorChatListener;
@@ -209,6 +210,7 @@ public class ContainerScreenOverlay extends REIOverlay {
     }
     
     public void init() {
+        Argument.SEARCH_CACHE.clear();
         draggingStack.set(DraggableStackProvider.from(() -> ScreenRegistry.getInstance().getDraggableProviders()), 
                 DraggableStackVisitor.from(() -> ScreenRegistry.getInstance().getDraggableVisitors()));
         
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
index dcabfbdf2..7fd96e637 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
@@ -25,8 +25,14 @@ package me.shedaniel.rei.impl.client.search.argument;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.Lists;
+import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import it.unimi.dsi.fastutil.longs.Long2ObjectMaps;
+import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
+import it.unimi.dsi.fastutil.shorts.Short2ObjectMap;
+import it.unimi.dsi.fastutil.shorts.Short2ObjectOpenHashMap;
 import me.shedaniel.rei.api.client.gui.config.SearchMode;
 import me.shedaniel.rei.api.common.entry.EntryStack;
+import me.shedaniel.rei.api.common.util.EntryStacks;
 import me.shedaniel.rei.impl.client.search.IntRange;
 import me.shedaniel.rei.impl.client.search.argument.type.AlwaysMatchingArgumentType;
 import me.shedaniel.rei.impl.client.search.argument.type.ArgumentType;
@@ -51,6 +57,7 @@ import java.util.regex.Pattern;
 @Environment(EnvType.CLIENT)
 public class Argument<T, R> {
     public static final String SPACE = " ", EMPTY = "";
+    public static final Short2ObjectMap<Long2ObjectMap<Object>> SEARCH_CACHE = new Short2ObjectOpenHashMap<>();
     static final Argument<Unit, Unit> ALWAYS = new Argument<>(AlwaysMatchingArgumentType.INSTANCE, EMPTY, true, -1, -1, true);
     private ArgumentType<T, R> argumentType;
     private String text;
@@ -180,11 +187,10 @@ public class Argument<T, R> {
     
     private static <T, R, Z, B> boolean matches(EntryStack<?> stack, AlternativeArgument alternativeArgument, Mutable<?> mutable) {
         if (alternativeArgument.isEmpty()) return true;
+        long hashExact = EntryStacks.hashExact(stack);
         
         for (Argument<?, ?> argument : alternativeArgument) {
-            mutable.setValue(null);
-            
-            if (matches(argument.getArgument(), mutable, stack, argument.getText(), argument.filterData) == argument.isRegular()) {
+            if (matches((short) argument.getArgument().getIndex(), argument.getArgument(), mutable, stack, hashExact, argument.getText(), argument.filterData) == argument.isRegular()) {
                 return true;
             }
         }
@@ -192,8 +198,18 @@ public class Argument<T, R> {
         return false;
     }
     
-    private static <T, R, Z, B> boolean matches(ArgumentType<T, B> argumentType, Mutable<Z> data, EntryStack<?> stack, String filter, R filterData) {
-        return argumentType.matches((Mutable<B>) data, stack, filter, (T) filterData);
+    private static <T, R, Z, B> boolean matches(short argumentIndex, ArgumentType<T, B> argumentType, Mutable<Z> data, EntryStack<?> stack, long hashExact, String filter, R filterData) {
+        Long2ObjectMap<Object> map = SEARCH_CACHE.get(argumentIndex);
+        if (map == null) {
+            SEARCH_CACHE.put(argumentIndex, map = Long2ObjectMaps.synchronize(new Long2ObjectOpenHashMap<>()));
+        }
+        Z value = (Z) map.get(hashExact);
+        data.setValue(value);
+        boolean matches = argumentType.matches((Mutable<B>) data, stack, filter, (T) filterData);
+        if (value == null) {
+            map.put(hashExact, data.getValue());
+        }
+        return matches;
     }
     
     public ArgumentType<?, ?> getArgument() {
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/ArgumentType.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/ArgumentType.java
index 9aa9882f4..d46eacaf6 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/ArgumentType.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/ArgumentType.java
@@ -36,6 +36,8 @@ import org.jetbrains.annotations.Nullable;
 @ApiStatus.Internal
 @Environment(EnvType.CLIENT)
 public abstract class ArgumentType<T, R> {
+    private int index = -1;
+    
     public ArgumentType() {
     }
     
@@ -83,4 +85,9 @@ public abstract class ArgumentType<T, R> {
     public abstract boolean matches(Mutable<R> data, EntryStack<?> stack, String searchText, T filterData);
     
     public abstract T prepareSearchFilter(String searchText);
+    
+    public int getIndex() {
+        if (index >= 0) return index;
+        return index = ArgumentTypesRegistry.ARGUMENT_TYPE_LIST.indexOf(this);
+    }
 }