diff --git a/src/cmd/query.rs b/src/cmd/query.rs index 6539c2e..17690b7 100644 --- a/src/cmd/query.rs +++ b/src/cmd/query.rs @@ -86,6 +86,10 @@ impl Query { options = options.with_exists(true).with_resolve_symlinks(resolve_symlinks); } + if config::prefer_exact_match() { + options = options.with_prefer_exact_match(true); + } + let stream = Stream::new(db, options); Ok(stream) } diff --git a/src/config.rs b/src/config.rs index 0aeda5c..7f32458 100644 --- a/src/config.rs +++ b/src/config.rs @@ -60,3 +60,7 @@ pub fn maxage() -> Result { pub fn resolve_symlinks() -> bool { env::var_os("_ZO_RESOLVE_SYMLINKS").is_some_and(|var| var == "1") } + +pub fn prefer_exact_match() -> bool { + env::var_os("_ZO_PREFER_EXACT_MATCH").is_some_and(|var| var == "1") +} diff --git a/src/db/dir.rs b/src/db/dir.rs index 5d6d62c..d3a1ac2 100644 --- a/src/db/dir.rs +++ b/src/db/dir.rs @@ -1,9 +1,10 @@ use std::borrow::Cow; use std::fmt::{self, Display, Formatter}; +use std::path::Path; use serde::{Deserialize, Serialize}; -use crate::util::{DAY, HOUR, WEEK}; +use crate::util::{DAY, HOUR, WEEK, to_lowercase}; #[derive(Clone, Debug, Deserialize, Serialize)] pub struct Dir<'a> { @@ -31,6 +32,11 @@ impl Dir<'_> { self.rank * 0.25 } } + + pub fn is_exact_match(&self, keyword: &str) -> bool { + let last = Path::new(self.path.as_ref()).file_name().and_then(|s| s.to_str()).unwrap_or(""); + to_lowercase(last) == keyword + } } pub struct DirDisplay<'a> { diff --git a/src/db/mod.rs b/src/db/mod.rs index 1856fda..91be1bd 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -180,6 +180,24 @@ impl Database { self.with_dirty_mut(|dirty| *dirty = true); } + pub fn sort_by_score_with_keywords(&mut self, keywords: &[String], now: Epoch) { + self.with_dirs_mut(|dirs| { + dirs.sort_unstable_by(|dir1: &Dir, dir2: &Dir| { + let score = |dir: &Dir| { + if keywords.last().is_some_and(|kw| dir.is_exact_match(kw)) { + f64::MAX + } else { + dir.score(now) + } + }; + score(dir1) + .total_cmp(&score(dir2)) + .then(dir1.score(now).total_cmp(&dir2.score(now))) + }) + }); + self.with_dirty_mut(|dirty| *dirty = true); + } + pub fn dirty(&self) -> bool { *self.borrow_dirty() } @@ -260,6 +278,56 @@ mod tests { } } + #[test] + fn sort_by_score_with_keywords_exact_match_wins() { + let now = util::current_time().unwrap(); + let mut db = Database::new(PathBuf::default(), Vec::default(), |_| Vec::new(), false); + + db.add_unchecked("/foo/baz", 100.0, now); + db.add_unchecked("/foo/bar", 1.0, now); + + let keywords = vec!["bar".to_string()]; + + let baz = db.dirs()[0].score(now); + let bar = db.dirs()[1].score(now); + assert!(baz > bar); + + db.sort_by_score_with_keywords(&keywords, now); + + let dirs = db.dirs(); + assert_eq!(dirs.last().unwrap().path.as_ref(), "/foo/bar"); + assert!((dirs.last().unwrap().rank - 1.0).abs() < 0.01); + } + + #[test] + fn sort_by_score_with_keywords_highest_score_wins_without_exact_match() { + let now = util::current_time().unwrap(); + let mut db = Database::new(PathBuf::default(), Vec::default(), |_| Vec::new(), false); + + db.add_unchecked("/foo/baz", 100.0, now); + db.add_unchecked("/foo/bar", 1.0, now); + + let keywords = vec!["foo".to_string()]; + db.sort_by_score_with_keywords(&keywords, now); + + let dirs = db.dirs(); + assert_eq!(dirs.last().unwrap().path.as_ref(), "/foo/baz"); + } + + #[test] + fn sort_by_score_with_keywords_both_exact_match_frecency_breaks_tie() { + let now = util::current_time().unwrap(); + let mut db = Database::new(PathBuf::default(), Vec::default(), |_| Vec::new(), false); + + db.add_unchecked("/foo/bar", 1.0, now); + db.add_unchecked("/baz/bar", 100.0, now); + + let keywords = vec!["bar".to_string()]; + db.sort_by_score_with_keywords(&keywords, now); + + assert_eq!(db.dirs().last().unwrap().path.as_ref(), "/baz/bar"); + } + #[test] fn remove() { let data_dir = tempfile::tempdir().unwrap(); diff --git a/src/db/stream.rs b/src/db/stream.rs index 24c84e0..febd43c 100644 --- a/src/db/stream.rs +++ b/src/db/stream.rs @@ -16,7 +16,14 @@ pub struct Stream<'a> { impl<'a> Stream<'a> { pub fn new(db: &'a mut Database, options: StreamOptions) -> Self { - db.sort_by_score(options.now); + let now = options.now; + let keywords = &options.keywords; + + if options.prefer_exact_match { + db.sort_by_score_with_keywords(keywords, now); + } else { + db.sort_by_score(now); + } let idxs = (0..db.dirs().len()).rev(); Stream { db, idxs, options } } @@ -129,6 +136,9 @@ pub struct StreamOptions { /// Only return directories within this parent directory /// Does not check if the path exists base_dir: Option, + + /// Whether to prefer exact matches over the match with higher score. + prefer_exact_match: bool, } impl StreamOptions { @@ -141,6 +151,7 @@ impl StreamOptions { resolve_symlinks: false, ttl: now.saturating_sub(3 * MONTH), base_dir: None, + prefer_exact_match: false, } } @@ -172,6 +183,11 @@ impl StreamOptions { self.base_dir = base_dir; self } + + pub fn with_prefer_exact_match(mut self, prefer_exact_match: bool) -> Self { + self.prefer_exact_match = prefer_exact_match; + self + } } #[cfg(test)]