Improve Search Argument Matching

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

11 files changed, 93 insertions, 72 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/FavoritesListWidget.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/FavoritesListWidget.java
index fbdee9dda..7e85ff752 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/FavoritesListWidget.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/FavoritesListWidget.java
@@ -462,7 +462,7 @@ public class FavoritesListWidget extends WidgetWithBounds implements DraggableSt
         entry.y.setAs(y - 8);
         
         boolean contains = currentBounds.contains(PointHelper.ofMouse());
-        int newIndex = contains ? getReleaseIndex() : Iterables.indexOf(entries.values(), e -> e == entry);
+        int newIndex = contains ? getReleaseIndex() : Math.max(0, Iterables.indexOf(entries.values(), e -> e == entry));
         
         if (entries.size() - 1 <= newIndex) {
             Entry remove = this.entries.remove(entry.hashIgnoreAmount());
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/AlternativeArgument.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/AlternativeArgument.java
index 772e02633..9c1f24008 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/AlternativeArgument.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/AlternativeArgument.java
@@ -59,6 +59,10 @@ public class AlternativeArgument extends ForwardingList<Argument<?, ?>> {
             return this;
         }
         
+        public boolean isEmpty() {
+            return arguments == null;
+        }
+        
         public AlternativeArgument build() {
             if (arguments == null) return AlternativeArgument.EMPTY;
             if (arguments.size() == 1) return new AlternativeArgument(Collections.singletonList(arguments.get(0)));
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 8d5cbb59e..f654ce076 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,18 +25,14 @@ package me.shedaniel.rei.impl.client.search.argument;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.Lists;
-import me.shedaniel.math.Point;
 import me.shedaniel.rei.api.client.gui.config.SearchMode;
-import me.shedaniel.rei.api.client.gui.widgets.Tooltip;
 import me.shedaniel.rei.api.common.entry.EntryStack;
-import me.shedaniel.rei.api.common.util.CollectionUtils;
 import me.shedaniel.rei.impl.client.search.argument.type.AlwaysMatchingArgumentType;
 import me.shedaniel.rei.impl.client.search.argument.type.ArgumentType;
 import me.shedaniel.rei.impl.client.search.argument.type.ArgumentTypesRegistry;
 import me.shedaniel.rei.impl.client.search.result.ArgumentApplicableResult;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
-import net.minecraft.network.chat.Component;
 import net.minecraft.util.IntRange;
 import net.minecraft.util.Unit;
 import org.apache.commons.lang3.StringUtils;
@@ -102,33 +98,23 @@ public class Argument<T, R> {
     public static List<CompoundArgument> bakeArguments(String searchTerm, @Nullable ProcessedSink sink) {
         List<CompoundArgument> compoundArguments = Lists.newArrayList();
         int tokenStartIndex = 0;
-        String[] allTokens = StringUtils.splitByWholeSeparatorPreserveAllTokens(searchTerm, "|");
         
-        for (String token : allTokens) {
+        for (String token : StringUtils.splitByWholeSeparatorPreserveAllTokens(searchTerm, "|")) {
             Matcher terms = SPLIT_PATTERN.matcher(token);
             CompoundArgument.Builder builder = CompoundArgument.builder();
             while (terms.find()) {
-                String term = MoreObjects.firstNonNull(terms.group(1), terms.group(2));
                 AlternativeArgument.Builder alternativeBuilder = AlternativeArgument.builder();
                 
                 for (ArgumentType<?, ?> type : ArgumentTypesRegistry.ARGUMENT_TYPE_LIST) {
-                    if (type.getSearchMode() == SearchMode.NEVER) continue;
-                    ArgumentApplicableResult result = type.checkApplicable(term);
-                    
-                    if (result.isApplicable()) {
-                        int group = terms.group(1) != null ? 1 : 2;
-                        Argument<?, ?> argument = new Argument<>(type, result.getText(), !result.isInverted(),
-                                terms.start(group) + tokenStartIndex, terms.end(group) + tokenStartIndex, !result.shouldPreserveCasing());
-                        alternativeBuilder.add(argument);
-                        if (sink != null) {
-                            if (group == 1) {
-                                sink.addQuote(terms.start() + tokenStartIndex);
-                                if (terms.end() - 1 + tokenStartIndex < searchTerm.length()) {
-                                    sink.addQuote(terms.end() - 1 + tokenStartIndex);
-                                }
-                            }
-                            sink.addPart(argument, result.isUsingGrammar(), result.grammarRanges(), terms.start() + tokenStartIndex);
-                        }
+                    applyArgument(type, searchTerm, terms, tokenStartIndex, alternativeBuilder, true, sink);
+                    if (!alternativeBuilder.isEmpty()) {
+                        break;
+                    }
+                }
+                
+                if (alternativeBuilder.isEmpty()) {
+                    for (ArgumentType<?, ?> type : ArgumentTypesRegistry.ARGUMENT_TYPE_LIST) {
+                        applyArgument(type, searchTerm, terms, tokenStartIndex, alternativeBuilder, false, sink);
                     }
                 }
                 
@@ -140,6 +126,11 @@ public class Argument<T, R> {
                 sink.addSplitter(tokenStartIndex - 1);
             }
         }
+        prepareSearchFilter(compoundArguments);
+        return compoundArguments;
+    }
+    
+    private static void prepareSearchFilter(List<CompoundArgument> compoundArguments) {
         for (CompoundArgument arguments : compoundArguments) {
             for (AlternativeArgument alternativeArgument : arguments) {
                 for (Argument<?, ?> argument : alternativeArgument) {
@@ -148,14 +139,35 @@ public class Argument<T, R> {
                 }
             }
         }
-        return compoundArguments;
+    }
+    
+    private static void applyArgument(ArgumentType<?, ?> type, String searchTerm, Matcher terms, int tokenStartIndex, AlternativeArgument.Builder alternativeBuilder, boolean forceGrammar, @Nullable ProcessedSink sink) {
+        String term = MoreObjects.firstNonNull(terms.group(1), terms.group(2));
+        if (type.getSearchMode() == SearchMode.NEVER) return;
+        ArgumentApplicableResult result = type.checkApplicable(term, forceGrammar);
+        
+        if (result.isApplicable()) {
+            int group = terms.group(1) != null ? 1 : 2;
+            Argument<?, ?> argument = new Argument<>(type, result.getText(), !result.isInverted(),
+                    terms.start(group) + tokenStartIndex, terms.end(group) + tokenStartIndex, !result.shouldPreserveCasing());
+            alternativeBuilder.add(argument);
+            if (sink != null) {
+                if (group == 1) {
+                    sink.addQuote(terms.start() + tokenStartIndex);
+                    if (terms.end() - 1 + tokenStartIndex < searchTerm.length()) {
+                        sink.addQuote(terms.end() - 1 + tokenStartIndex);
+                    }
+                }
+                sink.addPart(argument, result.isUsingGrammar(), result.grammarRanges(), terms.start() + tokenStartIndex);
+            }
+        }
     }
     
     @ApiStatus.Internal
     public static boolean matches(EntryStack<?> stack, List<CompoundArgument> compoundArguments) {
         if (compoundArguments.isEmpty()) return true;
         Mutable<?> mutable = new MutableObject<>();
-    
+        
         a:
         for (CompoundArgument arguments : compoundArguments) {
             for (AlternativeArgument argument : arguments) {
@@ -188,13 +200,6 @@ public class Argument<T, R> {
         return argumentType.matches((Mutable<B>) data, stack, filter, (T) filterData);
     }
     
-    public static String tryGetEntryStackTooltip(EntryStack<?> stack) {
-        Tooltip tooltip = stack.getTooltip(new Point());
-        if (tooltip != null)
-            return CollectionUtils.mapAndJoinToString(tooltip.getText(), Component::getString, "\n");
-        return "";
-    }
-    
     public ArgumentType<?, ?> getArgument() {
         return argumentType;
     }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/AlwaysMatchingArgumentType.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/AlwaysMatchingArgumentType.java
index 8c0a3af0d..d64b0d6cb 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/AlwaysMatchingArgumentType.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/AlwaysMatchingArgumentType.java
@@ -52,7 +52,7 @@ public final class AlwaysMatchingArgumentType extends ArgumentType<Unit, Unit> {
     }
     
     @Override
-    public ArgumentApplicableResult checkApplicable(String text) {
+    public ArgumentApplicableResult checkApplicable(String text, boolean forceGrammar) {
         return ArgumentApplicableResult.notApplicable();
     }
     
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 aa8368967..9aa9882f4 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
@@ -54,12 +54,16 @@ public abstract class ArgumentType<T, R> {
         return SearchMode.PREFIX;
     }
     
-    public ArgumentApplicableResult checkApplicable(String text) {
-        ArgumentApplicableResult status = checkApplicable(text, getPrefix());
-        if (status.isApplicable()) {
-            return status;
-        } else if (getSearchMode() == SearchMode.ALWAYS) {
-            status = checkApplicable(text, "");
+    public ArgumentApplicableResult checkApplicable(String text, boolean forceGrammar) {
+        String prefix = getPrefix();
+        if (forceGrammar && !prefix.isEmpty()) {
+            ArgumentApplicableResult status = checkApplicable(text, prefix);
+            if (status.isApplicable()) {
+                return status;
+            }
+        }
+        if (!forceGrammar && getSearchMode() == SearchMode.ALWAYS) {
+            ArgumentApplicableResult status = checkApplicable(text, "");
             if (status.isApplicable()) {
                 status.notUsingGrammar();
             }
@@ -71,7 +75,7 @@ public abstract class ArgumentType<T, R> {
     private ArgumentApplicableResult checkApplicable(String text, String prefix) {
         if (prefix == null) return ArgumentApplicableResult.notApplicable();
         if (text.startsWith("-" + prefix)) return ArgumentApplicableResult.applyInverted(text.substring(1 + prefix.length())).grammar(0, prefix.length() + 1);
-        if (text.startsWith(prefix + "-")) return ArgumentApplicableResult.applyInverted(text.substring(1 + prefix.length())).grammar(0, prefix.length() + 1);
+        if (!prefix.isEmpty() && text.startsWith(prefix + "-")) return ArgumentApplicableResult.applyInverted(text.substring(1 + prefix.length())).grammar(0, prefix.length() + 1);
         if (text.startsWith(prefix)) return ArgumentApplicableResult.apply(text.substring(prefix.length())).grammar(0, prefix.length());
         return ArgumentApplicableResult.notApplicable();
     }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/IdentifierArgumentType.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/IdentifierArgumentType.java
index a17bfbcf9..1fd0a8249 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/IdentifierArgumentType.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/IdentifierArgumentType.java
@@ -71,7 +71,7 @@ public final class IdentifierArgumentType extends ArgumentType<Unit, String> {
             if (identifier == null) {
                 data.setValue(EMPTY);
             } else {
-                String s = identifier.toString();
+                String s = identifier.getPath();
                 if (s.isEmpty()) {
                     data.setValue(EMPTY);
                 } else {
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/RegexArgumentType.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/RegexArgumentType.java
index d8c13007f..566698faa 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/RegexArgumentType.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/RegexArgumentType.java
@@ -50,7 +50,7 @@ public final class RegexArgumentType extends ArgumentType<@Nullable Pattern, Str
     }
     
     @Override
-    public ArgumentApplicableResult checkApplicable(String text) {
+    public ArgumentApplicableResult checkApplicable(String text, boolean forceGrammar) {
         boolean inverted = false;
         String matchText = text;
         if (matchText.startsWith("-")) {
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TextArgumentType.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TextArgumentType.java
index 347862cd5..a0661c8c0 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TextArgumentType.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TextArgumentType.java
@@ -44,7 +44,8 @@ public final class TextArgumentType extends ArgumentType<Unit, String> {
     }
     
     @Override
-    public @Nullable String getPrefix() {
+    @Nullable
+    public String getPrefix() {
         return "";
     }
     
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TooltipArgumentType.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TooltipArgumentType.java
index 759e6cdeb..f200c7afa 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TooltipArgumentType.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/type/TooltipArgumentType.java
@@ -23,12 +23,15 @@
 
 package me.shedaniel.rei.impl.client.search.argument.type;
 
+import me.shedaniel.math.Point;
 import me.shedaniel.rei.api.client.config.ConfigObject;
 import me.shedaniel.rei.api.client.gui.config.SearchMode;
+import me.shedaniel.rei.api.client.gui.widgets.Tooltip;
 import me.shedaniel.rei.api.common.entry.EntryStack;
-import me.shedaniel.rei.impl.client.search.argument.Argument;
+import me.shedaniel.rei.api.common.util.CollectionUtils;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
+import net.minecraft.network.chat.Component;
 import net.minecraft.network.chat.Style;
 import net.minecraft.network.chat.TextColor;
 import net.minecraft.util.Unit;
@@ -68,12 +71,20 @@ public final class TooltipArgumentType extends ArgumentType<Unit, String> {
     @Override
     public boolean matches(Mutable<String> data, EntryStack<?> stack, String searchText, Unit filterData) {
         if (data.getValue() == null) {
-            data.setValue(Argument.tryGetEntryStackTooltip(stack).toLowerCase(Locale.ROOT));
+            data.setValue(tryGetEntryStackTooltip(stack).toLowerCase(Locale.ROOT));
         }
         String tooltip = data.getValue();
         return tooltip.isEmpty() || tooltip.contains(searchText);
     }
     
+    public static String tryGetEntryStackTooltip(EntryStack<?> stack) {
+        Tooltip tooltip = stack.getTooltip(new Point(), false);
+        if (tooltip != null) {
+            return CollectionUtils.mapAndJoinToString(tooltip.getText(), Component::getString, "\n");
+        }
+        return "";
+    }
+    
     @Override
     public Unit prepareSearchFilter(String searchText) {
         return Unit.INSTANCE;
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 3d5b3794f..5a0a5db08 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
@@ -30,13 +30,14 @@ import it.unimi.dsi.fastutil.shorts.Short2ObjectOpenHashMap;
 import me.shedaniel.math.Point;
 import me.shedaniel.math.Rectangle;
 import me.shedaniel.rei.api.client.ClientHelper;
-import me.shedaniel.rei.api.client.config.ConfigObject;
 import me.shedaniel.rei.api.client.gui.AbstractRenderer;
 import me.shedaniel.rei.api.client.gui.widgets.Tooltip;
 import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.entry.comparison.ComparisonContext;
 import me.shedaniel.rei.api.common.util.EntryStacks;
 import net.minecraft.resources.ResourceLocation;
+import org.apache.commons.lang3.mutable.Mutable;
+import org.apache.commons.lang3.mutable.MutableObject;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
 
@@ -107,19 +108,19 @@ public abstract class AbstractEntryStack<A> extends AbstractRenderer implements
     
     @Override
     @Nullable
-    public Tooltip getTooltip(Point mouse) {
-        Tooltip[] tooltip = {this.get(Settings.RENDER).apply(this).<A>cast().getTooltip(this, mouse)};
-        if (tooltip[0] == null) return null;
-        tooltip[0].getText().addAll(get(EntryStack.Settings.TOOLTIP_APPEND_EXTRA).apply(this));
-        tooltip[0] = get(EntryStack.Settings.TOOLTIP_PROCESSOR).apply(this, tooltip[0]);
-        if (tooltip[0] == null) return null;
-        if (ConfigObject.getInstance().shouldAppendModNames()) {
+    public Tooltip getTooltip(Point mouse, boolean appendModName) {
+        Mutable<Tooltip> tooltip = new MutableObject<>(this.get(Settings.RENDER).apply(this).<A>cast().getTooltip(this, mouse));
+        if (tooltip.getValue() == null) return null;
+        tooltip.getValue().getText().addAll(get(EntryStack.Settings.TOOLTIP_APPEND_EXTRA).apply(this));
+        tooltip.setValue(get(Settings.TOOLTIP_PROCESSOR).apply(this, tooltip.getValue()));
+        if (tooltip.getValue() == null) return null;
+        if (appendModName) {
             ResourceLocation location = getIdentifier();
             if (location != null) {
-                ClientHelper.getInstance().appendModIdToTooltips(tooltip[0].getText(), location.getNamespace());
+                ClientHelper.getInstance().appendModIdToTooltips(tooltip.getValue().getText(), location.getNamespace());
             }
         }
-        return tooltip[0];
+        return tooltip.getValue();
     }
     
     @Override
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java
index 6a7314253..218622def 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java
@@ -23,6 +23,7 @@
 
 package me.shedaniel.rei.impl.common.transfer;
 
+import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
 import me.shedaniel.rei.api.common.display.Display;
@@ -31,23 +32,22 @@ import me.shedaniel.rei.api.common.transfer.info.MenuInfo;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfoRegistry;
 import net.minecraft.world.inventory.AbstractContainerMenu;
 
+import java.util.List;
 import java.util.Map;
 import java.util.function.Predicate;
 
 public class MenuInfoRegistryImpl implements MenuInfoRegistry {
     private final Map<CategoryIdentifier<?>, Map<Class<? extends AbstractContainerMenu>, MenuInfo<?, ?>>> map = Maps.newLinkedHashMap();
-    private final Map<Predicate<CategoryIdentifier<?>>, Map<Class<? extends AbstractContainerMenu>, MenuInfo<?, ?>>> mapGeneric = Maps.newLinkedHashMap();
+    private final Map<Predicate<CategoryIdentifier<?>>, List<MenuInfo<?, ?>>> mapGeneric = Maps.newLinkedHashMap();
     
     @Override
-    public <D extends Display> void register(CategoryIdentifier<D> category, MenuInfo<?, D> menuInfo) {
-        map.computeIfAbsent(category, id -> Maps.newLinkedHashMap())
-                .put(menuInfo.getMenuClass(), menuInfo);
+    public <C extends AbstractContainerMenu, D extends Display> void register(CategoryIdentifier<D> category, Class<C> menuClass, MenuInfo<C, D> menuInfo) {
+        map.computeIfAbsent(category, id -> Maps.newLinkedHashMap()).put(menuClass, menuInfo);
     }
     
     @Override
     public <D extends Display> void registerGeneric(Predicate<CategoryIdentifier<?>> categoryPredicate, MenuInfo<?, D> menuInfo) {
-        mapGeneric.computeIfAbsent(categoryPredicate, id -> Maps.newLinkedHashMap())
-                .put(menuInfo.getMenuClass(), menuInfo);
+        mapGeneric.computeIfAbsent(categoryPredicate, id -> Lists.newArrayList()).add(menuInfo);
     }
     
     @Override
@@ -64,16 +64,11 @@ public class MenuInfoRegistryImpl implements MenuInfoRegistry {
             }
         }
         
-        for (Map.Entry<Predicate<CategoryIdentifier<?>>, Map<Class<? extends AbstractContainerMenu>, MenuInfo<?, ?>>> entry : mapGeneric.entrySet()) {
+        for (Map.Entry<Predicate<CategoryIdentifier<?>>, List<MenuInfo<?, ?>>> entry : mapGeneric.entrySet()) {
             if (entry.getKey().test(category) && !entry.getValue().isEmpty()) {
-                infoMap = entry.getValue();
-                if (infoMap.containsKey(menuClass)) {
-                    return (MenuInfo<T, D>) infoMap.get(menuClass);
-                }
-                for (Map.Entry<Class<? extends AbstractContainerMenu>, MenuInfo<?, ?>> infoEntry : infoMap.entrySet()) {
-                    if (infoEntry.getKey().isAssignableFrom(menuClass)) {
-                        return (MenuInfo<T, D>) infoEntry.getValue();
-                    }
+                List<MenuInfo<?, ?>> infoList = entry.getValue();
+                if (!infoList.isEmpty()) {
+                    return (MenuInfo<T, D>) infoList.get(0);
                 }
             }
         }