Async search (#245)

* Async search

* Define index when it's used

1 files changed, 73 insertions, 0 deletions

diff --git a/src/main/java/cuchaz/enigma/utils/search/SearchUtil.java b/src/main/java/cuchaz/enigma/utils/search/SearchUtil.java
index e5ed35f..a51afbb 100644
--- a/src/main/java/cuchaz/enigma/utils/search/SearchUtil.java
+++ b/src/main/java/cuchaz/enigma/utils/search/SearchUtil.java
@@ -1,6 +1,12 @@
 package cuchaz.enigma.utils.search;
 
 import java.util.*;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -11,6 +17,7 @@ public class SearchUtil<T extends SearchEntry> {
 
         private final Map<T, Entry<T>> entries = new HashMap<>();
         private final Map<String, Integer> hitCount = new HashMap<>();
+        private final Executor searchExecutor = Executors.newWorkStealingPool();
 
         public void add(T entry) {
                 Entry<T> e = Entry.from(entry);
@@ -46,6 +53,59 @@ public class SearchUtil<T extends SearchEntry> {
                                 .sequential();
         }
 
+        public SearchControl asyncSearch(String term, SearchResultConsumer<T> consumer) {
+                Map<String, Integer> hitCount = new HashMap<>(this.hitCount);
+                Map<T, Entry<T>> entries = new HashMap<>(this.entries);
+                float[] scores = new float[entries.size()];
+                Lock scoresLock = new ReentrantLock();
+                AtomicInteger size = new AtomicInteger();
+                AtomicBoolean control = new AtomicBoolean(false);
+                AtomicInteger elapsed = new AtomicInteger();
+                for (Entry<T> value : entries.values()) {
+                        searchExecutor.execute(() -> {
+                                try {
+                                        if (control.get()) return;
+                                        float score = value.getScore(term, hitCount.getOrDefault(value.searchEntry.getIdentifier(), 0));
+                                        if (score <= 0) return;
+                                        score = -score; // sort descending
+                                        try {
+                                                scoresLock.lock();
+                                                if (control.get()) return;
+                                                int dataSize = size.getAndIncrement();
+                                                int index = Arrays.binarySearch(scores, 0, dataSize, score);
+                                                if (index < 0) {
+                                                        index = ~index;
+                                                }
+                                                System.arraycopy(scores, index, scores, index + 1, dataSize - index);
+                                                scores[index] = score;
+                                                consumer.add(index, value.searchEntry);
+                                        } finally {
+                                                scoresLock.unlock();
+                                        }
+                                } finally {
+                                        elapsed.incrementAndGet();
+                                }
+                        });
+                }
+
+                return new SearchControl() {
+                        @Override
+                        public void stop() {
+                                control.set(true);
+                        }
+
+                        @Override
+                        public boolean isFinished() {
+                                return entries.size() == elapsed.get();
+                        }
+
+                        @Override
+                        public float getProgress() {
+                                return (float) elapsed.get() / entries.size();
+                        }
+                };
+        }
+
         public void hit(T entry) {
                 if (entries.containsKey(entry)) {
                         hitCount.compute(entry.getIdentifier(), (_id, i) -> i == null ? 1 : i + 1);
@@ -192,4 +252,17 @@ public class SearchUtil<T extends SearchEntry> {
 
         }
 
+        @FunctionalInterface
+        public interface SearchResultConsumer<T extends SearchEntry> {
+                void add(int index, T entry);
+        }
+
+        public interface SearchControl {
+                void stop();
+
+                boolean isFinished();
+
+                float getProgress();
+        }
+
 }