feat: fuzzy search in candidate selection menus

Cargo.toml

@@ -11,6 +11,7 @@ generational-arena = "0.2"
 crossterm = "0.27.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+regex = "1.10"
 ron = "0.8.1"
 # TODO: opt-out of uneeded Rhai features
 [dependencies.rhai]

src/keymap/menu.rs

@@ -2,7 +2,7 @@ use super::keymap::{Candidate, KeyProg, Keymap};
 use crate::frontends::Key;
 use crate::language::Storage;
 use crate::tree::Node;
-use crate::util::{bug_assert, SynlessBug};
+use crate::util::{bug_assert, fuzzy_search, SynlessBug};
 
 const SELECTION_LANGUAGE_NAME: &str = "selection_menu";
 
@@ -27,6 +27,7 @@ pub struct Menu {
 struct MenuSelection {
     custom_candidate: Option<Candidate>,
     candidates: Vec<Candidate>,
+    filtered_candidates: Vec<Candidate>,
     input: String,
     index: usize,
     default_index: usize,
@@ -40,56 +41,59 @@ impl MenuSelection {
             return None;
         }
         let default_index = if custom_candidate.is_some() { 1 } else { 0 };
-        Some(MenuSelection {
+        let mut menu = MenuSelection {
             custom_candidate,
             candidates,
+            filtered_candidates: Vec::new(),
             input: String::new(),
             index: default_index,
             default_index,
-        })
+        };
+        menu.update_filtered_candidates();
+        Some(menu)
     }
 
     fn execute(&mut self, cmd: MenuSelectionCmd) {
         use MenuSelectionCmd::{Backspace, Down, Insert, Up};
 
         match cmd {
             Up => self.index = self.index.saturating_sub(1),
-            Down => self.index += 1,
+            Down => {
+                if self.index + 1 < self.filtered_candidates.len() {
+                    self.index += 1;
+                }
+            }
             Backspace => {
                 self.input.pop();
                 if let Some(Candidate::Custom { input }) = &mut self.custom_candidate {
                     input.pop();
                 }
                 self.index = self.default_index;
+                self.update_filtered_candidates();
             }
             Insert(ch) => {
                 self.input.push(ch);
                 if let Some(Candidate::Custom { input }) = &mut self.custom_candidate {
                     input.push(ch);
                 }
                 self.index = self.default_index;
+                self.update_filtered_candidates();
             }
         }
     }
 
-    // TODO: Implement fuzzy search
-    fn filtered_candidates(&self) -> Vec<&Candidate> {
-        let mut filtered = Vec::new();
+    fn update_filtered_candidates(&mut self) {
+        self.filtered_candidates =
+            fuzzy_search(&self.input, self.candidates.clone(), |candidate| {
+                candidate.display_str()
+            });
         if let Some(candidate) = &self.custom_candidate {
-            filtered.push(candidate);
-        }
-        for candidate in &self.candidates {
-            if candidate.display_str().contains(&self.input) {
-                filtered.push(candidate);
-            }
+            self.filtered_candidates.insert(0, candidate.to_owned());
         }
-        filtered
     }
 
     fn selected_candidate(&self) -> Option<&Candidate> {
-        let candidates = self.filtered_candidates();
-        let index = self.index.min(candidates.len().saturating_sub(1));
-        candidates.get(index).copied()
+        self.filtered_candidates.get(self.index)
     }
 
     fn make_candidate_selection_doc(&self, s: &mut Storage) -> Node {
@@ -114,16 +118,14 @@ impl MenuSelection {
         bug_assert!(root.insert_last_child(s, input_node));
 
         // Add candidate entries, highlighting the one at self.index
-        let candidates = self.filtered_candidates();
-        let index = self.index.min(candidates.len().saturating_sub(1));
-        for (i, candidate) in candidates.iter().enumerate() {
+        for (i, candidate) in self.filtered_candidates.iter().enumerate() {
             let construct = match candidate {
                 Custom { .. } => c_custom,
                 Regular { .. } => c_regular,
                 Special { .. } => c_special,
             };
             let mut node = Node::with_text(s, construct, candidate.display_str().to_owned()).bug();
-            if i == index {
+            if i == self.index {
                 node = Node::with_children(s, c_selected, [node]).bug();
             }
             bug_assert!(root.insert_last_child(s, node));

src/util/fuzzy_search.rs

@@ -0,0 +1,130 @@
+use crate::util::SynlessBug;
+use regex::{self, Regex, RegexBuilder};
+
+/// Return only the `items` that match the `input` search string, with the best matches first.
+/// `get_str` returns a string representation for each item (to be compared with `input`).
+pub fn fuzzy_search<T>(input: &str, items: Vec<T>, get_str: impl Fn(&T) -> &str) -> Vec<T> {
+    if input.is_empty() {
+        return items;
+    }
+    let searcher = Searcher::new(input).bug_msg("fuzzy_search: bad regex construction");
+    let mut scored_items = items
+        .into_iter()
+        .filter_map(|item| searcher.score(get_str(&item)).map(|score| (score, item)))
+        .collect::<Vec<_>>();
+    scored_items.sort_by_key(|&(score, _)| score);
+    scored_items
+        .into_iter()
+        .map(|(_, item)| item)
+        .collect::<Vec<_>>()
+}
+
+/// Searches using regexes. Searching for `Baz` will return, in priority order:
+///
+///     Baz
+///     baZ
+///     Baz.js
+///     baZ.js
+///     FooBaz
+///     foObaZ
+///     FooBaz.js
+///     foObaZ.js
+struct Searcher(Vec<Regex>);
+
+impl Searcher {
+    fn new(input: &str) -> Result<Searcher, regex::Error> {
+        let pattern = input
+            .split(" ")
+            .map(regex::escape)
+            .collect::<Vec<_>>()
+            .join(".*");
+        let regex_strs = vec![
+            format!("^{}$", pattern),
+            format!("^{}", pattern),
+            format!("{}$", pattern),
+            pattern.clone(),
+        ];
+        let mut regexes = Vec::new();
+        for regex_str in regex_strs {
+            regexes.push(Regex::new(&regex_str)?);
+            regexes.push(
+                RegexBuilder::new(&regex_str)
+                    .case_insensitive(true)
+                    .build()?,
+            );
+        }
+        Ok(Searcher(regexes))
+    }
+
+    /// A score between for how well `item` matches the input used to construct `regexes`, where
+    /// 0.0 is a perfect match and larger numbers are worse matches. Assumes that `regexes` are in
+    /// order from most specific (good match) to least specific (bad match).
+    fn score(&self, item: &str) -> Option<[usize; 3]> {
+        for (i, regex) in self.0.iter().enumerate() {
+            if let Some(matched) = regex.find(item) {
+                return Some([i, matched.start(), item.len()]);
+            }
+        }
+        None
+    }
+}
+
+#[test]
+fn test_fuzzy_search() {
+    let items = vec!["foo", "bar", "foobarz"];
+    let sorted = fuzzy_search("ba", |x| x, items);
+    assert_eq!(sorted, vec!["bar", "foobarz"]);
+
+    let items = vec!["foobarz", "bar", "foo"];
+    let sorted = fuzzy_search("fo", |x| x, items);
+    assert_eq!(sorted, vec!["foo", "foobarz"]);
+
+    let items = vec![
+        "bare.js",
+        "foobear.js",
+        "foobare.js",
+        "foo.js",
+        "foo.rs",
+        "xfoobear.js",
+        "fooo",
+        "foo",
+        "fo",
+        "bar.rs.rs",
+    ];
+
+    let sorted = fuzzy_search("foo", |x| x, items.clone());
+    assert_eq!(
+        sorted,
+        vec![
+            "foo",
+            "fooo",
+            "foo.js",
+            "foo.rs",
+            "foobear.js",
+            "foobare.js",
+            "xfoobear.js",
+        ]
+    );
+
+    let sorted = fuzzy_search("", |x| x, items.clone());
+    assert_eq!(
+        sorted,
+        vec![
+            "bare.js",
+            "foobear.js",
+            "foobare.js",
+            "foo.js",
+            "foo.rs",
+            "xfoobear.js",
+            "fooo",
+            "foo",
+            "fo",
+            "bar.rs.rs",
+        ]
+    );
+
+    assert_eq!(
+        fuzzy_search("json", |x| x, vec!["foo.json", "json.cpp"]),
+        vec!["json.cpp", "foo.json"]
+    );
+}

src/util/mod.rs

@@ -1,11 +1,13 @@
 mod bug;
 mod error;
+mod fuzzy_search;
 mod indexed_map;
 mod log;
 mod ordered_map;
 
 pub use bug::{bug, bug_assert, format_bug, SynlessBug};
 pub use error::{error, ErrorCategory, SynlessError};
+pub use fuzzy_search::fuzzy_search;
 pub use indexed_map::IndexedMap;
 pub use log::{log, Log, LogEntry, LogLevel};
 pub use ordered_map::OrderedMap;