Add recursive flag for add and remove commands

Cargo.toml

@@ -30,6 +30,7 @@ fastrand = "2.0.0"
 glob = "0.3.0"
 ouroboros = "0.18.3"
 serde = { version = "1.0.116", features = ["derive"] }
+walkdir = "2.5.0"
 
 [target.'cfg(unix)'.dependencies]
 nix = { version = "0.30.1", default-features = false, features = [

src/cmd/add.rs

@@ -1,4 +1,4 @@
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 use anyhow::{Result, bail};
 
@@ -18,10 +18,17 @@ impl Run for Add {
 
         let mut db = Database::open()?;
 
-        for path in &self.paths {
+        // Build a unified iterator of PathBufs, whether recursive or not.
+        let paths: Box<dyn Iterator<Item = PathBuf>> = if self.recursive {
+            Box::new(self.paths.iter().flat_map(|p| util::walk_dir(p)))
+        } else {
+            Box::new(self.paths.iter().cloned())
+        };
+
+        for path in paths {
             let path =
                 if config::resolve_symlinks() { util::canonicalize } else { util::resolve_path }(
-                    path,
+                    &path,
                 )?;
             let path = util::path_to_str(&path)?;
 

src/cmd/cmd.rs

@@ -63,6 +63,10 @@ pub struct Add {
     /// doesn't
     #[clap(short, long)]
     pub score: Option<f64>,
+
+    /// Recursively add directories
+    #[clap(short, long, default_value_t = false)]
+    pub recursive: bool,
 }
 
 /// Edit the database
@@ -201,4 +205,8 @@ pub struct Query {
 pub struct Remove {
     #[clap(value_hint = ValueHint::DirPath)]
     pub paths: Vec<String>,
+
+    /// Recursively remove directories
+    #[clap(short, long, default_value_t = false)]
+    pub recursive: bool,
 }

src/cmd/edit.rs

@@ -17,7 +17,7 @@ impl Run for Edit {
                 match cmd {
                     EditCommand::Decrement { path } => db.add(path, -1.0, now),
                     EditCommand::Delete { path } => {
-                        db.remove(path);
+                        db.remove(path, false);
                     }
                     EditCommand::Increment { path } => db.add(path, 1.0, now),
                     EditCommand::Reload => {}

src/cmd/remove.rs

@@ -9,10 +9,10 @@ impl Run for Remove {
         let mut db = Database::open()?;
 
         for path in &self.paths {
-            if !db.remove(path) {
+            if !db.remove(path, self.recursive) {
                 let path_abs = util::resolve_path(path)?;
                 let path_abs = util::path_to_str(&path_abs)?;
-                if path_abs == path || !db.remove(path_abs) {
+                if path_abs == path || !db.remove(path_abs, self.recursive) {
                     bail!("path not found in database: {path}")
                 }
             }

src/db/mod.rs

@@ -101,17 +101,35 @@ impl Database {
     }
 
     /// Removes the directory with `path` from the store. This does not preserve
-    /// ordering, but is O(1).
-    pub fn remove(&mut self, path: impl AsRef<str>) -> bool {
-        match self.dirs().iter().position(|dir| dir.path == path.as_ref()) {
-            Some(idx) => {
-                self.swap_remove(idx);
-                true
+    /// ordering, but is O(1). If recursive, this will remove all directories
+    /// starting with `path`. This is O(n)
+    pub fn remove(&mut self, path: impl AsRef<str>, recursive: bool) -> bool {
+        if recursive {
+            self.remove_recursive(path)
+        } else {
+            match self.dirs().iter().position(|dir| dir.path == path.as_ref()) {
+                Some(idx) => {
+                    self.swap_remove(idx);
+                    true
+                }
+                None => false,
             }
-            None => false,
         }
     }
 
+    pub fn remove_recursive(&mut self, path: impl AsRef<str>) -> bool {
+        let mut removed = false;
+        self.with_dirs_mut(|dirs| {
+            dirs.retain(|dir| {
+                let keep = !dir.path.starts_with(path.as_ref());
+                removed |= !keep;
+                keep
+            });
+        });
+        self.with_dirty_mut(|dirty| *dirty = removed);
+        removed
+    }
+
     pub fn swap_remove(&mut self, idx: usize) {
         self.with_dirs_mut(|dirs| dirs.swap_remove(idx));
         self.with_dirty_mut(|dirty| *dirty = true);
@@ -273,14 +291,14 @@ mod tests {
 
         {
             let mut db = Database::open_dir(data_dir.path()).unwrap();
-            assert!(db.remove(path));
+            assert!(db.remove(path, false));
             db.save().unwrap();
         }
 
         {
             let mut db = Database::open_dir(data_dir.path()).unwrap();
             assert!(db.dirs().is_empty());
-            assert!(!db.remove(path));
+            assert!(!db.remove(path, false));
             db.save().unwrap();
         }
     }

src/util.rs

@@ -9,6 +9,7 @@ use std::{env, mem};
 #[cfg(windows)]
 use anyhow::anyhow;
 use anyhow::{Context, Result, bail};
+use walkdir::WalkDir;
 
 use crate::db::{Dir, Epoch};
 use crate::error::SilentExit;
@@ -382,3 +383,12 @@ pub fn to_lowercase(s: impl AsRef<str>) -> String {
     let s = s.as_ref();
     if s.is_ascii() { s.to_ascii_lowercase() } else { s.to_lowercase() }
 }
+
+/// Recursively walk a directory and yield all directories.
+pub fn walk_dir(path: &Path) -> impl Iterator<Item = PathBuf> {
+    WalkDir::new(path)
+        .into_iter()
+        .filter_map(|p| p.ok())
+        .filter(|p| p.file_type().is_dir())
+        .map(|p| p.into_path())
+}