Async search (#245)

* Async search

* Define index when it's used

2 files changed, 81 insertions, 5 deletions

diff --git a/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java b/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java
index b36ebfb4..b283a377 100644
--- a/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java
+++ b/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java
@@ -36,12 +36,13 @@ import cuchaz.enigma.utils.search.SearchUtil;
 public class SearchDialog {
 
         private final JTextField searchField;
-        private final DefaultListModel<SearchEntryImpl> classListModel;
+        private DefaultListModel<SearchEntryImpl> classListModel;
         private final JList<SearchEntryImpl> classList;
         private final JDialog dialog;
 
         private final Gui parent;
         private final SearchUtil<SearchEntryImpl> su;
+        private SearchUtil.SearchControl currentSearch;
 
         public SearchDialog(Gui parent) {
                 this.parent = parent;
@@ -173,11 +174,13 @@ public class SearchDialog {
 
         // Updates the list of class names
         private void updateList() {
-                classListModel.clear();
+                if (currentSearch != null) currentSearch.stop();
 
-                su.search(searchField.getText())
-                                .limit(100)
-                                .forEach(classListModel::addElement);
+                DefaultListModel<SearchEntryImpl> classListModel = new DefaultListModel<>();
+                this.classListModel = classListModel;
+                classList.setModel(classListModel);
+
+                currentSearch = su.asyncSearch(searchField.getText(), (idx, e) -> SwingUtilities.invokeLater(() -> classListModel.insertElementAt(e, idx)));
         }
 
         public void dispose() {
diff --git a/src/main/java/cuchaz/enigma/utils/search/SearchUtil.java b/src/main/java/cuchaz/enigma/utils/search/SearchUtil.java
index e5ed35fd..a51afbb0 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();
+        }
+
 }