From 083a83431042aed338e9dee7855d1041eea54801 Mon Sep 17 00:00:00 2001
From: Cole Helbling <cole.e.helbling@outlook.com>
Date: Sun, 15 Mar 2020 11:18:59 -0700
Subject: [PATCH] Allow merging when migrating from `z` (#33)

If the user passes the `--merge` flag to the `migrate` subcommand, all
duplicate entries will have their ranks and epochs updated: the rank
will be the sum of the stored rank and the newly-parsed rank, while the
epoch will be the maximum of the stored epoch and the newly-parsed
epoch.

This allows one to import from the `z` database even after having used
`zoxide` for any amount of time. This also permits a user who has
already sourced the init script to import their old database without
needing to do something like `rm ~/.zo && zoxide migrate ~/.z`.
---
 src/db.rs                 | 24 +++++++++++++++++++-----
 src/subcommand/migrate.rs |  5 ++++-
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index 84f3cff..0e9c4a8 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -70,11 +70,11 @@ impl DB {
         Ok(())
     }
 
-    pub fn migrate<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
-        if !self.dirs.is_empty() {
+    pub fn migrate<P: AsRef<Path>>(&mut self, path: P, merge: bool) -> Result<()> {
+        if !self.dirs.is_empty() && !merge {
             bail!(
-                "To prevent conflicts, you can only migrate from z with an empty \
-                 zoxide database!"
+                "To prevent conflicts, you can only migrate from z with an empty zoxide database!
+If you wish to merge the two, specify the `--merge` flag."
             );
         }
 
@@ -87,7 +87,10 @@ impl DB {
             let line = if let Ok(line) = read_line {
                 line
             } else {
-                eprintln!("could not read line {}: {:?}", line_number, read_line);
+                eprintln!(
+                    "could not read entry at line {}: {:?}",
+                    line_number, read_line
+                );
                 continue;
             };
 
@@ -131,6 +134,17 @@ 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_str) {
+                            dir.rank += rank;
+                            dir.last_accessed = Epoch::max(epoch, dir.last_accessed);
+
+                            continue;
+                        };
+                    }
+
                     // FIXME: When we switch to PathBuf for storing directories inside Dir, just
                     // pass `PathBuf::from(path_str)`
                     self.dirs.push(Dir {
diff --git a/src/subcommand/migrate.rs b/src/subcommand/migrate.rs
index 3ffd017..c3e67ec 100644
--- a/src/subcommand/migrate.rs
+++ b/src/subcommand/migrate.rs
@@ -7,10 +7,13 @@ use structopt::StructOpt;
 #[structopt(about = "Migrate from z database")]
 pub struct Migrate {
     path: String,
+
+    #[structopt(long, help = "Merge entries into existing database")]
+    merge: bool,
 }
 
 impl Migrate {
     pub fn run(&self, env: &Env) -> Result<()> {
-        util::get_db(env)?.migrate(&self.path)
+        util::get_db(env)?.migrate(&self.path, self.merge)
     }
 }