Merge commit '41180dd40ac5214da245cfa7956dc662c4d95bea' into modularity

2 files changed, 29 insertions, 14 deletions

diff --git a/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java b/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
index 8724af7d1..d4bd7b6e1 100644
--- a/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
+++ b/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
@@ -47,6 +47,8 @@ import me.shedaniel.rei.api.common.util.EntryStacks;
 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 me.shedaniel.rei.impl.client.util.ThreadCreator;
+import me.shedaniel.rei.impl.common.InternalLogger;
 import me.shedaniel.rei.impl.common.util.HashedEntryStackWrapper;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
@@ -60,10 +62,7 @@ import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -71,6 +70,7 @@ import java.util.regex.Pattern;
 @ApiStatus.Internal
 @Environment(EnvType.CLIENT)
 public class Argument<T, R> {
+    private static final ExecutorService EXECUTOR_SERVICE = new ThreadCreator("REI-ArgumentCache").asService();
     public static final Short2ObjectMap<Long2ObjectMap<Object>> SEARCH_CACHE = Short2ObjectMaps.synchronize(new Short2ObjectOpenHashMap<>());
     private static final Object NO_CACHE = new Object();
     private static final AtomicReference<String> lastLanguage = new AtomicReference<>();
@@ -264,7 +264,7 @@ public class Argument<T, R> {
     }
     
     public static Long prepareStart = null;
-    public static Collection<EntryStack<?>> prepareStacks = null;
+    public static List<HashedEntryStackWrapper> prepareStacks = null;
     public static IntIntPair prepareStage = null;
     public static IntIntPair[] currentStages = null;
     
@@ -272,22 +272,34 @@ public class Argument<T, R> {
         if (prepareStage != null || currentStages != null) return;
         try {
             prepareStart = Util.getEpochMillis();
-            prepareStacks = stacks;
+            Long2ObjectMap<Object>[] caches = CollectionUtils.map(argumentTypes, Argument::getSearchCache).toArray(Long2ObjectMap[]::new);
+            prepareStacks = CollectionUtils.mapAndFilter(stacks, stack -> {
+                for (Long2ObjectMap<Object> cache : caches) {
+                    if (!cache.containsKey(stack.hashExact())) {
+                        return true;
+                    }
+                }
+                
+                return false;
+            }, HashedEntryStackWrapper::new);
+            if (prepareStacks.isEmpty()) {
+                return;
+            }
+            InternalLogger.getInstance().trace("Preparing " + prepareStacks.size() + " stacks for search arguments");
             prepareStage = new IntIntMutablePair(0, argumentTypes.size());
             currentStages = new IntIntPair[argumentTypes.size()];
-            List<HashedEntryStackWrapper> hashedStacks = CollectionUtils.map(stacks, HashedEntryStackWrapper::new);
             int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
-            boolean async = ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4;
+            boolean async = ConfigObject.getInstance().shouldAsyncSearch() && prepareStacks.size() > searchPartitionSize * 4;
             List<CompletableFuture<Long2ObjectMap<Object>>> futures = Lists.newArrayList();
             List<Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>>> pairs = Lists.newArrayList();
             
             for (ArgumentType<?, ?> argumentType : argumentTypes) {
                 prepareStage.first(prepareStage.firstInt() + 1);
                 Long2ObjectMap<Object> map = getSearchCache(argumentType);
-                IntIntPair currentStage = currentStages[prepareStage.firstInt() - 1] = new IntIntMutablePair(0, hashedStacks.size());
+                IntIntPair currentStage = currentStages[prepareStage.firstInt() - 1] = new IntIntMutablePair(0, prepareStacks.size());
                 
                 if (async) {
-                    for (Collection<HashedEntryStackWrapper> partitionStacks : CollectionUtils.partition(hashedStacks, searchPartitionSize)) {
+                    for (Collection<HashedEntryStackWrapper> partitionStacks : CollectionUtils.partition(prepareStacks, searchPartitionSize)) {
                         CompletableFuture<Long2ObjectMap<Object>> future = CompletableFuture.supplyAsync(() -> {
                             Long2ObjectMap<Object> out = new Long2ObjectArrayMap<>(searchPartitionSize + 1);
                             for (HashedEntryStackWrapper stack : partitionStacks) {
@@ -300,14 +312,14 @@ public class Argument<T, R> {
                                 }
                             }
                             return out;
-                        }).whenComplete((objectLong2ObjectMap, throwable) -> {
+                        }, EXECUTOR_SERVICE).whenComplete((objectLong2ObjectMap, throwable) -> {
                             currentStage.first(currentStage.firstInt() + partitionStacks.size());
                         });
                         futures.add(future);
                         pairs.add(Pair.of(argumentType, future));
                     }
                 } else {
-                    for (HashedEntryStackWrapper stack : hashedStacks) {
+                    for (HashedEntryStackWrapper stack : prepareStacks) {
                         currentStage.first(currentStage.firstInt() + 1);
                         
                         if (map.get(stack.hashExact()) == null) {
@@ -324,14 +336,17 @@ public class Argument<T, R> {
             if (async) {
                 try {
                     CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(30, TimeUnit.SECONDS);
-                } catch (InterruptedException | ExecutionException | TimeoutException e) {
+                } catch (ExecutionException | TimeoutException e) {
                     e.printStackTrace();
+                } catch (InterruptedException ignore) {
                 }
                 for (Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>> pair : pairs) {
                     Long2ObjectMap<Object> now = pair.second().getNow(null);
                     if (now != null) getSearchCache(pair.left()).putAll(now);
                 }
             }
+    
+            InternalLogger.getInstance().debug("Prepared " + prepareStacks.size() + " stacks for search arguments in " + (Util.getEpochMillis() - prepareStart) + "ms");
         } finally {
             prepareStart = null;
             prepareStacks = null;
diff --git a/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/method/InputMethodRegistryImpl.java b/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/method/InputMethodRegistryImpl.java
index 532792d8e..0a04608be 100644
--- a/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/method/InputMethodRegistryImpl.java
+++ b/runtime-engine/search/src/main/java/me/shedaniel/rei/impl/client/search/method/InputMethodRegistryImpl.java
@@ -99,7 +99,7 @@ public class InputMethodRegistryImpl implements InputMethodRegistry {
         }).join();
         service.shutdown();
         
-        InternalLogger.getInstance().debug("Registered %d input methods: ", inputMethods.size(),
+        InternalLogger.getInstance().debug("Registered %d input methods: %s", inputMethods.size(),
                 inputMethods.values().stream().map(inputMethod -> inputMethod.getName().getString()).collect(Collectors.joining(", ")));
     }