Merge pull request #83 from diniamo/cache-choices

feat: cache choices

Cargo.toml

@@ -11,3 +11,4 @@ license = "MIT"
 [dependencies]
 clap = { version = "4.3.12", features = ["derive", "cargo", "env"] }
 xdg = "2.5.0"
+bitcode = "0.6.3"

src/cache.rs

@@ -0,0 +1,59 @@
+use std::{collections::HashMap, error::Error, fs, path::PathBuf};
+
+use bitcode::{Decode, Encode};
+
+#[derive(Encode, Decode)]
+struct CacheData(HashMap<String, String>);
+
+pub struct Cache {
+    path: PathBuf,
+    data: CacheData,
+    update: bool,
+}
+
+impl Cache {
+    pub fn new() -> Result<Self, Box<dyn Error>> {
+        let path = xdg::BaseDirectories::new()?.place_state_file("comma-choices")?;
+
+        Ok(Self {
+            data: if path.exists() {
+                let bytes = fs::read(&path)?;
+                bitcode::decode(&bytes)?
+            } else {
+                CacheData(HashMap::new())
+            },
+            path,
+            update: false,
+        })
+    }
+
+    pub fn query(&self, command: &str) -> Option<String> {
+        self.data.0.get(command).cloned()
+    }
+
+    pub fn update(&mut self, command: &str, derivation: &str) {
+        self.data.0.insert(command.into(), derivation.into());
+        self.update = true;
+    }
+
+    pub fn delete(&mut self, command: &str) {
+        self.data.0.remove(command);
+        self.update = true;
+    }
+
+    pub fn empty(&mut self) {
+        self.data.0.clear();
+        self.update = true;
+    }
+}
+
+impl Drop for Cache {
+    fn drop(&mut self) {
+        if self.update {
+            let bytes = bitcode::encode(&self.data.0);
+            if let Err(e) = fs::write(&self.path, bytes) {
+                eprintln!("failed to write cache: {e}");
+            }
+        }
+    }
+}

src/main.rs

@@ -1,12 +1,15 @@
+mod cache;
 mod index;
 mod shell;
+
 use std::{
     env,
     io::Write,
     os::unix::prelude::CommandExt,
     process::{self, Command, ExitCode, Stdio},
 };
 
+use cache::Cache;
 use clap::crate_version;
 use clap::Parser;
 
@@ -36,6 +39,43 @@ fn pick(picker: &str, derivations: &[&str]) -> Option<String> {
     )
 }
 
+fn index_database(command: &str, picker: &str) -> Option<String> {
+    index::check_database_updated();
+
+    let nix_locate_output = Command::new("nix-locate")
+        .args(["--top-level", "--minimal", "--at-root", "--whole-name"])
+        .arg(format!("/bin/{command}"))
+        .output()
+        .expect("failed to execute nix-locate");
+
+    if !nix_locate_output.status.success() {
+        match std::str::from_utf8(&nix_locate_output.stderr) {
+            Ok(stderr) => eprintln!("nix-locate failed with: {stderr}"),
+            Err(_) => eprintln!("nix-locate failed"),
+        }
+        return None;
+    }
+
+    let attrs = nix_locate_output.stdout;
+
+    if attrs.is_empty() {
+        eprintln!("No executable `{command}` found in nix-index database.");
+        return None;
+    }
+
+    let attrs = std::str::from_utf8(&attrs)
+        .expect("fail")
+        .trim()
+        .split('\n')
+        .collect::<Box<[&str]>>();
+
+    if attrs.len() > 1 {
+        pick(picker, &attrs)
+    } else {
+        attrs.first().map(|s| s.trim().to_owned())
+    }
+}
+
 fn run_command_or_open_shell(
     use_channel: bool,
     choice: &str,
@@ -70,57 +110,57 @@ fn run_command_or_open_shell(
 fn main() -> ExitCode {
     let args = Opt::parse();
 
+    let mut cache = Cache::new();
+    if let Err(ref e) = cache {
+        eprintln!("failed to initialize cache, disabling related functionality: {e}");
+    }
+
     if args.update {
         eprintln!("\"comma --update\" has been deprecated. either obtain a prebuilt database from https://github.com/Mic92/nix-index-database or use \"nix run 'nixpkgs#nix-index' --extra-experimental-features 'nix-command flakes'\"");
         index::update_database();
     }
 
+    if args.empty_cache {
+        if let Ok(ref mut cache) = cache {
+            cache.empty();
+        }
+    }
+
     // The command may not be given if `--update` was specified.
     if args.cmd.is_empty() {
-        return ExitCode::FAILURE;
+        return if args.update || args.empty_cache {
+            ExitCode::SUCCESS
+        } else {
+            ExitCode::FAILURE
+        };
     }
 
     let command = &args.cmd[0];
     let trail = &args.cmd[1..];
 
-    index::check_database_updated();
-
-    let nix_locate_output = Command::new("nix-locate")
-        .args(["--top-level", "--minimal", "--at-root", "--whole-name"])
-        .arg(format!("/bin/{command}"))
-        .output()
-        .expect("failed to execute nix-locate");
-
-    if !nix_locate_output.status.success() {
-        match std::str::from_utf8(&nix_locate_output.stderr) {
-            Ok(stderr) => eprintln!("nix-locate failed with: {stderr}"),
-            Err(_) => eprint!("nix-locate failed"),
+    if args.delete_entry {
+        if let Ok(ref mut cache) = cache {
+            cache.delete(command);
         }
-        return ExitCode::FAILURE;
     }
 
-    let attrs = nix_locate_output.stdout;
-
-    if attrs.is_empty() {
-        eprintln!("No executable `{command}` found in nix-index database.");
-        return ExitCode::FAILURE;
-    }
-
-    let attrs: Vec<_> = std::str::from_utf8(&attrs)
-        .expect("fail")
-        .trim()
-        .split('\n')
-        .collect();
+    let derivation = match cache {
+        Ok(mut cache) => cache.query(command).or_else(|| {
+            index_database(command, &args.picker).map(|derivation| {
+                cache.update(command, &derivation);
+                derivation
+            })
+        }),
+        Err(_) => index_database(command, &args.picker),
+    };
 
-    let choice = if attrs.len() > 1 {
-        match pick(&args.picker, &attrs) {
-            Some(x) => x,
-            None => return ExitCode::FAILURE,
-        }
-    } else {
-        attrs.first().unwrap().trim().to_owned()
+    let derivation = match derivation {
+        Some(d) => d,
+        None => return ExitCode::FAILURE,
     };
 
+    let basename = derivation.rsplit('.').last().unwrap();
+
     let use_channel = match env::var("NIX_PATH") {
         Ok(val) => val,
         Err(_) => String::new(),
@@ -130,22 +170,27 @@ fn main() -> ExitCode {
     if args.print_package {
         println!(
             "Package that contains executable /bin/{}: {}",
-            command,
-            &choice.rsplit('.').last().unwrap()
+            command, basename
         );
     };
 
     if args.install {
         Command::new("nix-env")
-            .args(["-f", "<nixpkgs>", "-iA", choice.rsplit('.').last().unwrap()])
+            .args(["-f", "<nixpkgs>", "-iA", basename])
             .exec();
     } else if args.shell {
         let shell_cmd = shell::select_shell_from_pid(process::id()).unwrap_or("bash".into());
-        run_command_or_open_shell(use_channel, &choice, &shell_cmd, &[], &args.nixpkgs_flake);
+        run_command_or_open_shell(
+            use_channel,
+            &derivation,
+            &shell_cmd,
+            &[],
+            &args.nixpkgs_flake,
+        );
     } else if args.print_path {
         run_command_or_open_shell(
             use_channel,
-            &choice,
+            &derivation,
             "sh",
             &[
                 String::from("-c"),
@@ -154,7 +199,13 @@ fn main() -> ExitCode {
             &args.nixpkgs_flake,
         );
     } else {
-        run_command_or_open_shell(use_channel, &choice, command, trail, &args.nixpkgs_flake);
+        run_command_or_open_shell(
+            use_channel,
+            &derivation,
+            command,
+            trail,
+            &args.nixpkgs_flake,
+        );
     }
 
     ExitCode::SUCCESS
@@ -195,7 +246,16 @@ struct Opt {
     #[clap(short = 'x', long = "print-path")]
     print_path: bool,
 
+    /// Empty the cache
+    #[clap(short, long = "empty-cache")]
+    empty_cache: bool,
+
+    /// Overwrite the cache entry for the specified command. This is achieved by first deleting it
+    /// from the cache, then running comma as normal.
+    #[clap(short, long = "delete-entry")]
+    delete_entry: bool,
+
     /// Command to run
-    #[clap(required_unless_present = "update", name = "cmd")]
+    #[clap(required_unless_present_any = ["update", "empty_cache"], name = "cmd")]
     cmd: Vec<String>,
 }