commands: add an interactive mode for jj restore

This adds an interactive mode for `jj restore`. It works by first
creating two temporary directories with the contents of the subset of
files that differ between the two trees, and then letting the user
edit the directory representing the right/after side. This has some
advantages compared to the interactive modes in Git and Mercurial:

 * It lets the user edit the final state as opposed to the diff itself
   (depending on the diff tool, of course). I think most users find it
   easier to edit the file contents than to edit the patch
   format.

 * It delegates the hard work to a tool that is already written (this
   is a big advantage for an immature tool like Jujube, but it is not
   an advantage from the user's point of view).

Almost all of the work in this commit went into adding a function that
takes two trees, lets the user edit the diff, and returns a new tree
id. I plan to reuse that function for other interactive commands. One
planned command is `jj edit`, which will let the user edit the changes
in a commit. `jj edit -r abc123` will be mostly about providing a more
intuitive name for `jj restore --source abc123^ --destination abc123`,
plus it will be different for merge commits (it will edit only the
changes in the merge commit). I also plan to add `jj split` by letting
the user edit the full diff, leaving only the parts that should go
into the first commit. Perhaps there will also be commands for moving
part of a commit out of or into a parent commit.

Cargo.lock

@@ -556,6 +556,7 @@ dependencies = [
  "serde_json",
  "tempfile",
  "test-case",
+ "thiserror",
  "uuid",
  "zstd",
 ]

Cargo.toml

@@ -28,6 +28,7 @@ protobuf = { version = "2.18.1", features = ["with-bytes"] }
 protobuf-codegen-pure = "2.18.1"
 serde_json = "1.0.60"
 tempfile = "3.1.0"
+thiserror = "1.0.22"
 uuid = { version = "0.8.1", features = ["v4"] }
 zstd = "0.6.0"
 

lib/src/trees.rs

@@ -30,6 +30,16 @@ pub enum Diff<T> {
     Removed(T),
 }
 
+impl<T> Diff<T> {
+    pub fn as_options(&self) -> (Option<&T>, Option<&T>) {
+        match self {
+            Diff::Modified(left, right) => (Some(left), Some(right)),
+            Diff::Added(right) => (None, Some(right)),
+            Diff::Removed(left) => (Some(left), None),
+        }
+    }
+}
+
 pub type TreeValueDiff<'a> = Diff<&'a TreeValue>;
 
 fn diff_entries<'a, E>(

lib/src/working_copy.rs

@@ -71,7 +71,7 @@ impl FileState {
     }
 }
 
-struct TreeState {
+pub struct TreeState {
     store: Arc<StoreWrapper>,
     working_copy_path: PathBuf,
     state_path: PathBuf,

src/commands.rs

@@ -51,6 +51,8 @@ use jj_lib::trees::TreeValueDiff;
 use jj_lib::working_copy::{CheckoutStats, WorkingCopy};
 
 use self::chrono::{FixedOffset, TimeZone, Utc};
+use crate::commands::CommandError::UserError;
+use crate::diff_edit::DiffEditError;
 use crate::graphlog::{AsciiGraphDrawer, Edge};
 use crate::styler::{ColorStyler, Styler};
 use crate::template_parser::TemplateParser;
@@ -66,6 +68,12 @@ enum CommandError {
     InternalError(String),
 }
 
+impl From<DiffEditError> for CommandError {
+    fn from(err: DiffEditError) -> Self {
+        CommandError::UserError(format!("Failed to edit diff: {}", err))
+    }
+}
+
 fn get_repo(ui: &Ui, matches: &ArgMatches) -> Result<Arc<ReadonlyRepo>, CommandError> {
     let repo_path_str = matches.value_of("repository").unwrap();
     let repo_path = ui.cwd().join(repo_path_str);
@@ -392,10 +400,14 @@ fn get_app<'a, 'b>() -> App<'a, 'b> {
                         .takes_value(true)
                         .default_value("@"),
                 )
+                .arg(
+                    Arg::with_name("interactive")
+                        .long("interactive")
+                        .short("i"),
+                )
                 .arg(
                     Arg::with_name("paths")
                         .index(1)
-                        .required(true)
                         .multiple(true),
                 ),
         )
@@ -1319,22 +1331,32 @@ fn cmd_restore(
     let source_commit = resolve_single_rev(ui, mut_repo, sub_matches.value_of("source").unwrap())?;
     let destination_commit =
         resolve_single_rev(ui, mut_repo, sub_matches.value_of("destination").unwrap())?;
-    let paths = sub_matches.values_of("paths").unwrap();
-    let mut tree_builder = repo
-        .store()
-        .tree_builder(destination_commit.tree().id().clone());
-    for path in paths {
-        let repo_path = RepoPath::from(path);
-        match source_commit.tree().path_value(&repo_path) {
-            Some(value) => {
-                tree_builder.set(repo_path, value);
-            }
-            None => {
-                tree_builder.remove(repo_path);
+    let tree_id;
+    if sub_matches.is_present("interactive") {
+        if sub_matches.is_present("paths") {
+            return Err(UserError(
+                "restore with --interactive and path is not yet supported".to_string(),
+            ));
+        }
+        tree_id = crate::diff_edit::edit_diff(&source_commit.tree(), &destination_commit.tree())?;
+    } else {
+        let paths = sub_matches.values_of("paths").unwrap();
+        let mut tree_builder = repo
+            .store()
+            .tree_builder(destination_commit.tree().id().clone());
+        for path in paths {
+            let repo_path = RepoPath::from(path);
+            match source_commit.tree().path_value(&repo_path) {
+                Some(value) => {
+                    tree_builder.set(repo_path, value);
+                }
+                None => {
+                    tree_builder.remove(repo_path);
+                }
             }
         }
+        tree_id = tree_builder.write_tree();
     }
-    let tree_id = tree_builder.write_tree();
     if &tree_id == destination_commit.tree().id() {
         ui.write("Nothing changed.\n");
     } else {

src/diff_edit.rs

@@ -0,0 +1,127 @@
+use jj_lib::repo_path::{DirRepoPath, RepoPath};
+use jj_lib::store::{StoreError, TreeId, TreeValue};
+use jj_lib::store_wrapper::StoreWrapper;
+use jj_lib::tree::Tree;
+use jj_lib::tree_builder::TreeBuilder;
+use jj_lib::trees::merge_trees;
+use jj_lib::working_copy::{CheckoutError, TreeState};
+use std::path::PathBuf;
+use std::process::Command;
+use std::sync::Arc;
+use tempfile::tempdir;
+use thiserror::Error;
+
+#[derive(Debug, Error, PartialEq, Eq)]
+pub enum DiffEditError {
+    #[error("The diff tool exited with a non-zero code")]
+    DifftoolAborted,
+    #[error("Failed to write directories to diff: {0:?}")]
+    CheckoutError(CheckoutError),
+    #[error("Internal error: {0:?}")]
+    InternalStoreError(StoreError),
+}
+
+impl From<CheckoutError> for DiffEditError {
+    fn from(err: CheckoutError) -> Self {
+        DiffEditError::CheckoutError(err)
+    }
+}
+
+impl From<StoreError> for DiffEditError {
+    fn from(err: StoreError) -> Self {
+        DiffEditError::InternalStoreError(err)
+    }
+}
+
+fn add_to_tree(
+    store: &StoreWrapper,
+    tree_builder: &mut TreeBuilder,
+    repo_path: &RepoPath,
+    value: &TreeValue,
+) -> Result<(), StoreError> {
+    match value {
+        TreeValue::Conflict(conflict_id) => {
+            let conflict = store.read_conflict(conflict_id)?;
+            let materialized_value =
+                jj_lib::conflicts::conflict_to_materialized_value(store, repo_path, &conflict);
+            tree_builder.set(repo_path.clone(), materialized_value);
+        }
+        _ => {
+            tree_builder.set(repo_path.clone(), (*value).clone());
+        }
+    }
+    Ok(())
+}
+
+fn check_out(
+    store: Arc<StoreWrapper>,
+    wc_dir: PathBuf,
+    state_dir: PathBuf,
+    tree_id: TreeId,
+) -> Result<TreeState, DiffEditError> {
+    std::fs::create_dir(&wc_dir).unwrap();
+    std::fs::create_dir(&state_dir).unwrap();
+    let mut tree_state = TreeState::init(store, wc_dir, state_dir);
+    tree_state.check_out(tree_id)?;
+    Ok(tree_state)
+}
+
+pub fn edit_diff(left_tree: &Tree, right_tree: &Tree) -> Result<TreeId, DiffEditError> {
+    // First create partial Trees of only the subset of the left and right trees
+    // that affect files changed between them.
+    let store = left_tree.store();
+    let mut left_tree_builder = store.tree_builder(store.empty_tree_id().clone());
+    let mut right_tree_builder = store.tree_builder(store.empty_tree_id().clone());
+    left_tree.diff(&right_tree, &mut |file_path, diff| {
+        let (left_value, right_value) = diff.as_options();
+        let repo_path = file_path.to_repo_path();
+        if let Some(value) = left_value {
+            add_to_tree(store, &mut left_tree_builder, &repo_path, value).unwrap();
+        }
+        if let Some(value) = right_value {
+            add_to_tree(store, &mut right_tree_builder, &repo_path, value).unwrap();
+        }
+    });
+    let left_partial_tree_id = left_tree_builder.write_tree();
+    let right_partial_tree_id = right_tree_builder.write_tree();
+    let right_partial_tree = store.get_tree(&DirRepoPath::root(), &right_partial_tree_id)?;
+
+    // Check out the two partial trees in temporary directories.
+    let temp_dir = tempdir().unwrap();
+    let left_wc_dir = temp_dir.path().join("left");
+    let left_state_dir = temp_dir.path().join("left_state");
+    let right_wc_dir = temp_dir.path().join("right");
+    let right_state_dir = temp_dir.path().join("right_state");
+    check_out(
+        store.clone(),
+        left_wc_dir.clone(),
+        left_state_dir,
+        left_partial_tree_id,
+    )?;
+    // TODO: mark left dir readonly
+    let mut right_tree_state = check_out(
+        store.clone(),
+        right_wc_dir.clone(),
+        right_state_dir,
+        right_partial_tree_id,
+    )?;
+
+    // Start a diff editor on the two directories.
+    let exit_status = Command::new("meld")
+        .arg(&left_wc_dir)
+        .arg(&right_wc_dir)
+        .status()
+        .expect("failed to run diff editor");
+    if !exit_status.success() {
+        return Err(DiffEditError::DifftoolAborted);
+    }
+
+    // Create a Tree based on the initial right tree, applying the changes made to
+    // that directory by the diff editor.
+    let new_right_partial_tree_id = right_tree_state.write_tree();
+    let new_right_partial_tree =
+        store.get_tree(&DirRepoPath::root(), &new_right_partial_tree_id)?;
+    let new_tree_id = merge_trees(right_tree, &right_partial_tree, &new_right_partial_tree)?;
+
+    Ok(new_tree_id)
+}

src/lib.rs

@@ -18,6 +18,7 @@
 extern crate pest_derive;
 
 pub mod commands;
+pub mod diff_edit;
 pub mod graphlog;
 pub mod styler;
 pub mod template_parser;