use crate::config; use crate::dir::{Dir, Epoch, Rank}; use anyhow::{anyhow, bail, Context, Result}; use serde::{Deserialize, Serialize}; use uuid::Uuid; use std::cmp::Ordering; use std::fs::{self, File, OpenOptions}; use std::io::{self, BufRead, BufReader, BufWriter}; use std::path::{Path, PathBuf}; pub use i32 as DBVersion; pub struct DB { data: DBData, modified: bool, path: PathBuf, // FIXME: remove after next breaking version path_old: Option, } impl DB { pub fn open>(path: P) -> Result { let data = match File::open(&path) { Ok(file) => { let reader = BufReader::new(&file); bincode::config() .limit(config::DB_MAX_SIZE) .deserialize_from(reader) .context("could not deserialize database")? } Err(err) => match err.kind() { io::ErrorKind::NotFound => DBData::default(), _ => return Err(err).context("could not open database file"), }, }; if data.version != config::DB_VERSION { bail!("database version '{}' is unsupported", data.version); } Ok(DB { data, modified: false, path: path.as_ref().to_path_buf(), path_old: None, }) } // FIXME: remove after next breaking version pub fn open_and_migrate(path_old: P1, path: P2) -> Result where P1: AsRef, P2: AsRef, { let file = File::open(&path_old).context("could not open old database file")?; let reader = BufReader::new(&file); let dirs = bincode::config() .limit(config::DB_MAX_SIZE) .deserialize_from(reader) .context("could not deserialize old database")?; let data = DBData { version: config::DB_VERSION, dirs, }; Ok(DB { data, modified: true, path: path.as_ref().to_path_buf(), path_old: Some(path_old.as_ref().to_path_buf()), }) } pub fn save(&mut self) -> Result<()> { if self.modified { let path_tmp = self.get_path_tmp(); let file_tmp = OpenOptions::new() .write(true) .create_new(true) .open(&path_tmp) .context("could not open temporary database file")?; let writer = BufWriter::new(&file_tmp); bincode::serialize_into(writer, &self.data).context("could not serialize database")?; if let Err(e) = fs::rename(&path_tmp, &self.path) { fs::remove_file(&path_tmp) .context("could not move or delete temporary database file")?; return Err(e).context("could not move temporary database file"); } self.modified = false; } Ok(()) } pub fn import>(&mut self, path: P, merge: bool) -> Result<()> { if !self.data.dirs.is_empty() && !merge { bail!( "To prevent conflicts, you can only import from z with an empty zoxide database!\n\ If you wish to merge the two, specify the `--merge` flag." ); } let z_db_file = File::open(path).context("could not open z database file")?; let reader = BufReader::new(z_db_file); for (idx, read_line) in reader.lines().enumerate() { let line_number = idx + 1; let line = if let Ok(line) = read_line { line } else { eprintln!( "could not read entry at line {}: {:?}", line_number, read_line ); continue; }; let split_line = line.rsplitn(3, '|').collect::>(); match split_line.as_slice() { [epoch_str, rank_str, path_str] => { let epoch = match epoch_str.parse::() { Ok(epoch) => epoch, Err(e) => { eprintln!( "invalid epoch '{}' at line {}: {}", epoch_str, line_number, e ); continue; } }; let rank = match rank_str.parse::() { Ok(rank) => rank, Err(e) => { eprintln!("invalid rank '{}' at line {}: {}", rank_str, line_number, e); continue; } }; let path_abs = match dunce::canonicalize(path_str) { Ok(path) => path, Err(e) => { eprintln!("invalid path '{}' at line {}: {}", path_str, line_number, e); continue; } }; if merge { // If the path exists in the database, add the ranks and set the epoch to // the largest of the parsed epoch and the already present epoch. if let Some(dir) = self.data.dirs.iter_mut().find(|dir| dir.path == path_abs) { dir.rank += rank; dir.last_accessed = Epoch::max(epoch, dir.last_accessed); continue; }; } self.data.dirs.push(Dir { path: path_abs, rank, last_accessed: epoch, }); } [] | [""] => {} // ignore blank lines line => { eprintln!("invalid entry at line {}: {:?}", line_number, line); continue; } }; } self.modified = true; Ok(()) } pub fn add>(&mut self, path: P, max_age: Rank, now: Epoch) -> Result<()> { let path_abs = dunce::canonicalize(&path) .with_context(|| anyhow!("could not access directory: {}", path.as_ref().display()))?; match self.data.dirs.iter_mut().find(|dir| dir.path == path_abs) { None => self.data.dirs.push(Dir { path: path_abs, last_accessed: now, rank: 1.0, }), Some(dir) => { dir.last_accessed = now; dir.rank += 1.0; } }; let sum_age = self.data.dirs.iter().map(|dir| dir.rank).sum::(); if sum_age > max_age { let factor = 0.9 * max_age / sum_age; for dir in &mut self.data.dirs { dir.rank *= factor; } self.data.dirs.retain(|dir| dir.rank >= 1.0); } self.modified = true; Ok(()) } pub fn query(&mut self, keywords: &[String], now: Epoch) -> Option<&Dir> { let (idx, dir, _) = self .data .dirs .iter() .enumerate() .filter(|(_, dir)| dir.is_match(&keywords)) .map(|(idx, dir)| (idx, dir, dir.get_frecency(now))) .max_by(|(_, _, frecency1), (_, _, frecency2)| { frecency1.partial_cmp(frecency2).unwrap_or(Ordering::Equal) })?; if dir.is_dir() { // FIXME: change this to Some(dir) once the MIR borrow checker comes to stable Rust Some(&self.data.dirs[idx]) } else { self.data.dirs.swap_remove(idx); self.modified = true; self.query(keywords, now) } } pub fn query_many<'a>(&'a mut self, keywords: &'a [String]) -> impl Iterator { self.query_all() .iter() .filter(move |dir| dir.is_match(keywords)) } pub fn query_all(&mut self) -> &[Dir] { let orig_len = self.data.dirs.len(); self.data.dirs.retain(Dir::is_dir); if orig_len != self.data.dirs.len() { self.modified = true; } self.data.dirs.as_slice() } pub fn remove>(&mut self, path: P) -> Result<()> { if let Ok(path_abs) = dunce::canonicalize(&path) { self.remove_exact(path_abs) .or_else(|_| self.remove_exact(path)) } else { self.remove_exact(path) } } pub fn remove_exact>(&mut self, path: P) -> Result<()> { if let Some(idx) = self .data .dirs .iter() .position(|dir| dir.path == path.as_ref()) { self.data.dirs.swap_remove(idx); self.modified = true; Ok(()) } else { bail!( "could not find path in database: {}", path.as_ref().display() ) } } fn get_path_tmp(&self) -> PathBuf { let file_name = format!(".{}.zo", Uuid::new_v4()); let mut path_tmp = self.path.clone(); path_tmp.set_file_name(file_name); path_tmp } } impl Drop for DB { fn drop(&mut self) { if let Err(e) = self.save() { eprintln!("{:#}", e); } else if let Some(path_old) = &self.path_old { // FIXME: remove this branch after next breaking release if let Err(e) = fs::remove_file(path_old).context("could not remove old database") { eprintln!("{:#}", e); } } } } #[derive(Debug, Deserialize, Serialize)] struct DBData { version: DBVersion, dirs: Vec, } impl Default for DBData { fn default() -> DBData { DBData { version: config::DB_VERSION, dirs: Vec::new(), } } }