Improve Async Search

5 files changed, 176 insertions, 69 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java
index a03865208..2eabb955c 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java
@@ -77,10 +77,11 @@ public class EntryListSearchManager {
         if (ignoreLastSearch) searchManager.markDirty();
         searchManager.updateFilter(searchTerm);
         if (searchManager.isDirty()) {
-            searchManager.getAsync(list -> {
+            searchManager.getAsync((list, filter) -> {
+                InternalLogger.getInstance().log(ConfigObject.getInstance().doDebugSearchTimeRequired() ? Level.INFO : Level.TRACE, "Search \"%s\" Used [%s]: %s", filter.getFilter(), Thread.currentThread().toString(), stopwatch.toString());
                 List</*EntryStack<?> | CollapsedStack*/ Object> finalList = collapse(copyAndOrder(list));
                 
-                InternalLogger.getInstance().log(ConfigObject.getInstance().doDebugSearchTimeRequired() ? Level.INFO : Level.TRACE, "Search Used: %s", stopwatch.stop().toString());
+                InternalLogger.getInstance().log(ConfigObject.getInstance().doDebugSearchTimeRequired() ? Level.INFO : Level.TRACE, "Search \"%s\" Used and Applied [%s]: %s", filter.getFilter(), Thread.currentThread().toString(), stopwatch.stop().toString());
                 
                 Minecraft.getInstance().executeBlocking(() -> {
                     update.accept(finalList);
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java
index 610d06163..0a05c1351 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java
@@ -29,25 +29,28 @@ import me.shedaniel.rei.api.client.search.SearchFilter;
 import me.shedaniel.rei.api.client.search.SearchProvider;
 import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.util.CollectionUtils;
+import me.shedaniel.rei.impl.common.InternalLogger;
 
+import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.*;
-import java.util.function.Consumer;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 
-public class AsyncSearchManager implements Supplier<List<EntryStack<?>>> {
+public class AsyncSearchManager {
+    private static final ThreadGroup GROUP = new ThreadGroup("REI-SearchManager");
+    private static final AtomicInteger THREAD_ID = new AtomicInteger(0);
     private final Supplier<List<EntryStack<?>>> stacksProvider;
     private final Supplier<Predicate<EntryStack<?>>> additionalPredicateSupplier;
     private final UnaryOperator<EntryStack<?>> transformer;
-    private Predicate<EntryStack<?>> additionalPredicate;
+    private ExecutorTuple executor;
     private SearchFilter filter;
-    private boolean dirty = false;
-    private boolean filterDirty = false;
-    private CompletableFuture<List<EntryStack<?>>> future;
-    private List<EntryStack<?>> last;
+    private Map.Entry<List<EntryStack<?>>, SearchFilter> last;
     
     public AsyncSearchManager(Supplier<List<EntryStack<?>>> stacksProvider, Supplier<Predicate<EntryStack<?>>> additionalPredicateSupplier, UnaryOperator<EntryStack<?>> transformer) {
         this.stacksProvider = stacksProvider;
@@ -55,91 +58,179 @@ public class AsyncSearchManager implements Supplier<List<EntryStack<?>>> {
         this.transformer = transformer;
     }
     
+    private static Thread createThread(Runnable task) {
+        Thread thread = new Thread(GROUP, task, "REI-SearchManager-" + THREAD_ID.getAndIncrement());
+        thread.setDaemon(true);
+        thread.setUncaughtExceptionHandler(($, exception) -> {
+            if (!(exception instanceof InterruptedException) && !(exception instanceof CancellationException) && !(exception instanceof ThreadDeath)) {
+                InternalLogger.getInstance().throwException(exception);
+            }
+        });
+        return thread;
+    }
+    
     public void markDirty() {
-        this.dirty = true;
+        synchronized (AsyncSearchManager.this) {
+            this.last = null;
+        }
     }
     
-    public void markFilterDirty() {
-        this.filterDirty = true;
+    private record ExecutorTuple(List<Thread> threads, SearchFilter filter, CompletableFuture<Map.Entry<List<EntryStack<?>>, SearchFilter>> future) {
     }
     
     public void updateFilter(String filter) {
         if (this.filter == null || !this.filter.getFilter().equals(filter)) {
+            if (this.executor != null) {
+                for (Thread thread : this.executor.threads()) {
+                    try {
+                        thread.stop();
+                    } catch (ThreadDeath ignored) {}
+                }
+            }
+            this.executor = null;
             this.filter = SearchProvider.getInstance().createFilter(filter);
-            markDirty();
-            markFilterDirty();
         }
     }
     
     public boolean isDirty() {
-        return last == null || dirty;
+        synchronized (AsyncSearchManager.this) {
+            return this.last == null || this.last.getValue() != this.filter;
+        }
     }
     
-    public boolean isFilterDirty() {
-        return filterDirty;
+    public Future<?> getAsync(BiConsumer<List<EntryStack<?>>, SearchFilter> consumer) {
+        if (executor == null || executor.filter() != filter) {
+            if (executor != null) {
+                for (Thread thread : this.executor.threads()) {
+                    try {
+                        thread.stop();
+                    } catch (ThreadDeath ignored) {}
+                }
+                executor = null;
+            }
+            List<Thread> threads = new ArrayList<>();
+            executor = new ExecutorTuple(threads, filter, get(task -> {
+                Thread thread = createThread(task);
+                threads.add(thread);
+                thread.start();
+            }));
+        }
+        SearchFilter savedFilter = filter;
+        ExecutorTuple tuple = executor;
+        return (executor = new ExecutorTuple(tuple.threads(), executor.filter(), executor.future().thenApplyAsync(result -> {
+            if (savedFilter == filter) {
+                consumer.accept(result.getKey(), result.getValue());
+            }
+            
+            return result;
+        }, task -> {
+            Thread thread = createThread(task);
+            tuple.threads().add(thread);
+            thread.start();
+        }))).future();
     }
     
-    public Future<Void> getAsync(Consumer<List<EntryStack<?>>> consumer) {
-        if (future == null || future.isCancelled() || future.isDone() || future.isCompletedExceptionally()) {
-            if (future != null) future.cancel(true);
-            future = CompletableFuture.supplyAsync(this)
-                    .exceptionally(throwable -> {
-                        throwable.printStackTrace();
-                        return null;
-                    });
+    public List<EntryStack<?>> getNow() {
+        try {
+            return get(Runnable::run).get().getKey();
+        } catch (ExecutionException e) {
+            throw new RuntimeException(e);
+        } catch (InterruptedException | CancellationException e) {
+            return Lists.newArrayList();
         }
-        return future.thenAccept(consumer);
     }
     
-    @Override
-    public List<EntryStack<?>> get() {
+    public CompletableFuture<Map.Entry<List<EntryStack<?>>, SearchFilter>> get(Executor executor) {
         if (isDirty()) {
-            this.additionalPredicate = additionalPredicateSupplier.get();
-            int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
-            List<EntryStack<?>> stacks = stacksProvider.get();
-            last = new ArrayList<>();
+            Map.Entry<List<EntryStack<?>>, SearchFilter> last;
+            synchronized (AsyncSearchManager.this) {
+                last = this.last;
+            }
+            return get(this.filter, this.additionalPredicateSupplier.get(), this.transformer,
+                    this.stacksProvider.get(), last, this, executor)
+                    .thenApply(entry -> {
+                        synchronized (AsyncSearchManager.this) {
+                            this.last = entry;
+                        }
+                        return entry;
+                    });
+        }
+        
+        return CompletableFuture.completedFuture(last);
+    }
+    
+    public static CompletableFuture<Map.Entry<List<EntryStack<?>>, SearchFilter>> get(SearchFilter filter, Predicate<EntryStack<?>> additionalPredicate,
+            UnaryOperator<EntryStack<?>> transformer, List<EntryStack<?>> stacks, Map.Entry<List<EntryStack<?>>, SearchFilter> last,
+            AsyncSearchManager manager, Executor executor) {
+        int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
+        
+        if (!stacks.isEmpty()) {
+            CompletableFuture<Void> preparationFuture = CompletableFuture.completedFuture(null);
             
-            if (!stacks.isEmpty()) {
-                if (filterDirty) {
-                    filter.prepareFilter(stacks);
-                    filterDirty = false;
-                }
-                
-                if (ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4) {
-                    List<CompletableFuture<List<EntryStack<?>>>> futures = Lists.newArrayList();
-                    for (Iterable<EntryStack<?>> partitionStacks : CollectionUtils.partition(stacks, searchPartitionSize)) {
-                        futures.add(CompletableFuture.supplyAsync(() -> {
-                            List<EntryStack<?>> filtered = Lists.newArrayList();
-                            for (EntryStack<?> stack : partitionStacks) {
-                                if (stack != null && matches(stack) && additionalPredicate.test(stack)) {
-                                    filtered.add(transformer.apply(stack));
-                                }
-                            }
-                            return filtered;
-                        }));
-                    }
-                    try {
-                        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(30, TimeUnit.SECONDS);
-                    } catch (InterruptedException | ExecutionException | TimeoutException e) {
-                        e.printStackTrace();
-                    }
-                    for (CompletableFuture<List<EntryStack<?>>> future : futures) {
-                        List<EntryStack<?>> now = future.getNow(null);
-                        if (now != null) last.addAll(now);
+            if (last == null || last.getValue() != filter) {
+                preparationFuture = CompletableFuture.runAsync(() -> {
+                    if (manager.filter == filter) {
+                        filter.prepareFilter(stacks);
+                    } else {
+                        throw new CancellationException();
                     }
-                } else {
-                    for (EntryStack<?> stack : stacks) {
-                        if (matches(stack) && additionalPredicate.test(stack)) {
-                            last.add(transformer.apply(stack));
+                }, executor);
+            }
+            
+            if (ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4) {
+                List<CompletableFuture<List<EntryStack<?>>>> futures = Lists.newArrayList();
+                for (Iterable<EntryStack<?>> partitionStacks : CollectionUtils.partition(stacks, searchPartitionSize)) {
+                    futures.add(CompletableFuture.supplyAsync(() -> {
+                        List<EntryStack<?>> filtered = Lists.newArrayList();
+                        for (EntryStack<?> stack : partitionStacks) {
+                            if (stack != null && filter.test(stack) && additionalPredicate.test(stack)) {
+                                filtered.add(transformer.apply(stack));
+                            }
+                            if (manager.filter != filter) throw new CancellationException();
                         }
-                    }
+                        return filtered;
+                    }));
                 }
+                return preparationFuture.thenCompose($ -> CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+                                .orTimeout(30, TimeUnit.SECONDS))
+                        .thenApplyAsync($ -> {
+                            List<EntryStack<?>> list = new ArrayList<>();
+                            
+                            if (manager.filter == filter) {
+                                for (CompletableFuture<List<EntryStack<?>>> future : futures) {
+                                    List<EntryStack<?>> now = future.getNow(null);
+                                    if (now != null) list.addAll(now);
+                                }
+                            } else {
+                                throw new CancellationException();
+                            }
+                            
+                            return list;
+                        }, executor)
+                        .thenApply(result -> {
+                            return new AbstractMap.SimpleImmutableEntry<>(result, filter);
+                        });
+            } else {
+                return preparationFuture.thenApplyAsync($ -> {
+                            List<EntryStack<?>> list = new ArrayList<>();
+                            
+                            for (EntryStack<?> stack : stacks) {
+                                if (filter.test(stack) && additionalPredicate.test(stack)) {
+                                    list.add(transformer.apply(stack));
+                                }
+                                if (manager.filter != filter) throw new CancellationException();
+                            }
+                            
+                            return list;
+                        }, executor)
+                        .thenApply(result -> {
+                            return new AbstractMap.SimpleImmutableEntry<>(result, filter);
+                        });
+                
             }
-            
-            dirty = false;
         }
         
-        return last;
+        return CompletableFuture.completedFuture(new AbstractMap.SimpleImmutableEntry<>(Lists.newArrayList(), filter));
     }
     
     public boolean matches(EntryStack<?> stack) {
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 d448370e5..04ee276cd 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
@@ -329,8 +329,9 @@ 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.getRight().getNow(null);
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/type/EntryRegistryImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/type/EntryRegistryImpl.java
index 293f3bc3e..af9f4ddb3 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/type/EntryRegistryImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/common/entry/type/EntryRegistryImpl.java
@@ -204,7 +204,11 @@ public class EntryRegistryImpl implements EntryRegistry {
         
         if (afterEntry != null) {
             int index = registryList.lastIndexOf(afterEntry);
-            registryList.addAll(index, filtered, hashes);
+            if (index != -1) {
+                registryList.addAll(index, filtered, hashes);
+            } else {
+                registryList.addAll(filtered, hashes);
+            }
         } else registryList.addAll(filtered, hashes);
         
         for (EntryRegistryListener listener : listeners) {
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java b/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java
index b9908cbae..8a8cc6568 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java
@@ -27,6 +27,7 @@ import com.google.common.collect.ImmutableList;
 import me.shedaniel.rei.api.client.favorites.FavoriteEntry;
 import me.shedaniel.rei.api.client.favorites.FavoriteEntryType;
 import me.shedaniel.rei.api.client.plugins.REIClientPlugin;
+import me.shedaniel.rei.api.client.registry.entry.CollapsibleEntryRegistry;
 import me.shedaniel.rei.api.client.registry.entry.EntryRegistry;
 import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.entry.comparison.ItemComparatorRegistry;
@@ -37,6 +38,7 @@ import net.fabricmc.api.Environment;
 import net.minecraft.client.Minecraft;
 import net.minecraft.core.Registry;
 import net.minecraft.nbt.CompoundTag;
+import net.minecraft.network.chat.TextComponent;
 import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.Items;
@@ -73,6 +75,14 @@ public class REITestPlugin implements REIClientPlugin {
     }
     
     @Override
+    public void registerCollapsibleEntries(CollapsibleEntryRegistry registry) {
+        for (Item item : Registry.ITEM) {
+            registry.group(Registry.ITEM.getKey(item), new TextComponent(Registry.ITEM.getKey(item).toString()),
+                    stack -> stack.getType() == VanillaEntryTypes.ITEM && stack.<ItemStack>castValue().is(item));
+        }
+    }
+    
+    @Override
     public void registerItemComparators(ItemComparatorRegistry registry) {
         registry.registerNbt(Registry.ITEM.stream().toArray(Item[]::new));
     }