Refactor DB architecture

src/config.rs

@@ -1,4 +1,5 @@
-use crate::types::Rank;
+use crate::db::DBVersion;
+use crate::dir::Rank;
 
 use anyhow::{bail, Context, Result};
 
@@ -6,6 +7,10 @@ use std::env;
 use std::fs;
 use std::path::PathBuf;
 
+pub const DB_MAX_SIZE: u64 = 8 * 1024 * 1024; // 8 MiB
+
+pub const DB_VERSION: DBVersion = 3;
+
 pub fn zo_data() -> Result<PathBuf> {
     let path = match env::var_os("_ZO_DATA") {
         Some(data_osstr) => PathBuf::from(data_osstr),

src/db.rs

@@ -1,40 +1,75 @@
-use crate::dir::Dir;
+use crate::config;
-use crate::types::{Epoch, Rank};
+use crate::dir::{Dir, Epoch, Rank};
 
 use anyhow::{anyhow, bail, Context, Result};
 use indoc::indoc;
+use serde::{Deserialize, Serialize};
+
+use std::cmp::Ordering;
 use std::fs::{self, File};
 use std::io::{self, BufRead, BufReader, BufWriter};
 use std::path::{Path, PathBuf};
 
+pub use i32 as DBVersion;
+
 pub struct DB {
-    path: PathBuf,
+    data: DBData,
-    dirs: Vec<Dir>,
     modified: bool,
+    path: PathBuf,
+    path_old: Option<PathBuf>,
 }
 
 impl DB {
     pub fn open<P: AsRef<Path>>(path: P) -> Result<DB> {
-        let path = path.as_ref().to_path_buf();
+        let data = match File::open(&path) {
-
-        let dirs = match File::open(&path) {
             Ok(file) => {
                 let reader = BufReader::new(&file);
                 bincode::config()
-                    .limit(8 * 1024 * 1024) // only databases upto 8 MiB are supported
+                    .limit(config::DB_MAX_SIZE)
                     .deserialize_from(reader)
                     .context("could not deserialize database")?
             }
             Err(err) => match err.kind() {
-                io::ErrorKind::NotFound => Vec::<Dir>::new(),
+                io::ErrorKind::NotFound => DBData::default(),
                 _ => return Err(err).context("could not open database file"),
             },
         };
 
+        if data.version != config::DB_VERSION {
+            bail!("this database version ({}) is unsupported", data.version);
+        }
+
         Ok(DB {
-            path,
+            data,
-            dirs,
             modified: false,
+            path: path.as_ref().to_path_buf(),
+            path_old: None,
+        })
+    }
+
+    pub fn open_old<P1, P2>(path_old: P1, path: P2) -> Result<DB>
+    where
+        P1: AsRef<Path>,
+        P2: AsRef<Path>,
+    {
+        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()),
         })
     }
 
@@ -46,7 +81,8 @@ impl DB {
                 File::create(&path_tmp).context("could not open temporary database file")?;
 
             let writer = BufWriter::new(&file_tmp);
-            bincode::serialize_into(writer, &self.dirs).context("could not serialize database")?;
+            bincode::serialize_into(writer, &self.data)
+                .context("could not serialize database")?;
 
             fs::rename(&path_tmp, &self.path).context("could not move temporary database file")?;
         }
@@ -55,7 +91,7 @@ impl DB {
     }
 
     pub fn import<P: AsRef<Path>>(&mut self, path: P, merge: bool) -> Result<()> {
-        if !self.dirs.is_empty() && !merge {
+        if !self.data.dirs.is_empty() && !merge {
             bail!(indoc!(
                 "To prevent conflicts, you can only import from z with an empty zoxide database!
                 If you wish to merge the two, specify the `--merge` flag."
@@ -109,7 +145,9 @@ impl DB {
                     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.dirs.iter_mut().find(|dir| dir.path == path_abs) {
+                        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);
 
@@ -117,7 +155,7 @@ impl DB {
                         };
                     }
 
-                    self.dirs.push(Dir {
+                    self.data.dirs.push(Dir {
                         path: path_abs,
                         rank,
                         last_accessed: epoch,
@@ -142,8 +180,8 @@ impl DB {
             .canonicalize()
             .with_context(|| anyhow!("could not access directory: {}", path.as_ref().display()))?;
 
-        match self.dirs.iter_mut().find(|dir| dir.path == path_abs) {
+        match self.data.dirs.iter_mut().find(|dir| dir.path == path_abs) {
-            None => self.dirs.push(Dir {
+            None => self.data.dirs.push(Dir {
                 path: path_abs,
                 last_accessed: now,
                 rank: 1.0,
@@ -154,15 +192,15 @@ impl DB {
             }
         };
 
-        let sum_age = self.dirs.iter().map(|dir| dir.rank).sum::<Rank>();
+        let sum_age = self.data.dirs.iter().map(|dir| dir.rank).sum::<Rank>();
 
         if sum_age > max_age {
             let factor = 0.9 * max_age / sum_age;
-            for dir in &mut self.dirs {
+            for dir in &mut self.data.dirs {
                 dir.rank *= factor;
             }
 
-            self.dirs.retain(|dir| dir.rank >= 1.0);
+            self.data.dirs.retain(|dir| dir.rank >= 1.0);
         }
 
         self.modified = true;
@@ -170,26 +208,36 @@ impl DB {
     }
 
     pub fn query(&mut self, keywords: &[String], now: Epoch) -> Option<Dir> {
-        let (idx, dir) = self
+        let (idx, dir, _) = self
+            .data
             .dirs
             .iter()
             .enumerate()
             .filter(|(_, dir)| dir.is_match(&keywords))
-            .max_by_key(|(_, dir)| dir.get_frecency(now) as i64)?;
+            .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() {
             Some(dir.to_owned())
         } else {
-            self.dirs.swap_remove(idx);
+            self.data.dirs.swap_remove(idx);
             self.modified = true;
             self.query(keywords, now)
         }
     }
 
     pub fn query_all(&mut self, keywords: &[String]) -> Vec<Dir> {
-        self.remove_invalid();
+        let orig_len = self.data.dirs.len();
+        self.data.dirs.retain(Dir::is_dir);
 
-        self.dirs
+        if orig_len != self.data.dirs.len() {
+            self.modified = true;
+        }
+
+        self.data
+            .dirs
             .iter()
             .filter(|dir| dir.is_match(&keywords))
             .cloned()
@@ -202,8 +250,8 @@ impl DB {
             Err(_) => path.as_ref().to_path_buf(),
         };
 
-        if let Some(idx) = self.dirs.iter().position(|dir| dir.path == path_abs) {
+        if let Some(idx) = self.data.dirs.iter().position(|dir| dir.path == path_abs) {
-            self.dirs.swap_remove(idx);
+            self.data.dirs.swap_remove(idx);
             self.modified = true;
         }
 
@@ -215,21 +263,31 @@ impl DB {
         path_tmp.set_file_name("db.zo.tmp");
         path_tmp
     }
-
-    fn remove_invalid(&mut self) {
-        let orig_len = self.dirs.len();
-        self.dirs.retain(Dir::is_dir);
-
-        if orig_len != self.dirs.len() {
-            self.modified = true;
-        }
-    }
 }
 
 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 {
+            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<Dir>,
+}
+
+impl Default for DBData {
+    fn default() -> DBData {
+        DBData {
+            version: config::DB_VERSION,
+            dirs: Vec::new(),
         }
     }
 }

src/dir.rs

@@ -1,8 +1,9 @@
-use crate::types::{Epoch, Rank};
-
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
 
+pub use f64 as Rank;
+pub use i64 as Epoch; // use a signed integer so subtraction can be performed on it
+
 #[derive(Clone, Debug, Default, Deserialize, Serialize)]
 pub struct Dir {
     pub path: PathBuf,

src/main.rs

@@ -2,7 +2,6 @@ mod config;
 mod db;
 mod dir;
 mod subcommand;
-mod types;
 mod util;
 
 use anyhow::Result;

src/subcommand/import.rs

@@ -3,10 +3,12 @@ use crate::util;
 use anyhow::Result;
 use structopt::StructOpt;
 
+use std::path::PathBuf;
+
 #[derive(Debug, StructOpt)]
 #[structopt(about = "Import from z database")]
 pub struct Import {
-    path: String,
+    path: PathBuf,
 
     #[structopt(long, help = "Merge entries into existing database")]
     merge: bool,

src/subcommand/query.rs

@@ -1,9 +1,10 @@
 use crate::util;
 
 use anyhow::{bail, Result};
+use structopt::StructOpt;
+
 use std::io::{self, Write};
 use std::path::Path;
-use structopt::StructOpt;
 
 #[derive(Debug, StructOpt)]
 #[structopt(about = "Search for a directory")]

src/subcommand/remove.rs

@@ -3,10 +3,12 @@ use crate::util;
 use anyhow::Result;
 use structopt::StructOpt;
 
+use std::path::PathBuf;
+
 #[derive(Debug, StructOpt)]
 #[structopt(about = "Remove a directory")]
 pub struct Remove {
-    path: String,
+    path: PathBuf,
 }
 
 impl Remove {

src/types.rs

@@ -1,2 +0,0 @@
-pub use f64 as Rank;
-pub use i64 as Epoch; // use a signed integer so subtraction can be performed on it

src/util.rs

@@ -1,9 +1,9 @@
 use crate::config;
 use crate::db::DB;
-use crate::dir::Dir;
+use crate::dir::{Dir, Epoch};
-use crate::types::Epoch;
 
 use anyhow::{anyhow, bail, Context, Result};
+
 use std::cmp::{Ordering, PartialOrd};
 use std::io::{Read, Write};
 use std::path::Path;