From 0058236a43fe0fe33a153c177ba90ff69e79b6bc Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Sun, 24 Apr 2022 12:13:43 -0700
Subject: [PATCH] cli: make `jj restore --to <rev>` restore from the working
 copy

---
 CHANGELOG.md                  |  3 +++
 src/commands.rs               | 22 ++++++++++++++++------
 tests/test_restore_command.rs | 26 ++++++++++++++++++++------
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ba6a0c29d..44a3003a2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -63,6 +63,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * `jj rebase` now requires at least one destination (allowing no arguments was a
   UX bug).
 
+* `jj restore --to <rev>` now restores from the working copy (it used to restore
+  from the working copy's parent).
+
 * You now get a proper error message instead of a crash when `$EDITOR` doesn't
   exist or exits with an error.
 
diff --git a/src/commands.rs b/src/commands.rs
index 73051c94d..2e4c14d35 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -1317,6 +1317,10 @@ struct UnsquashArgs {
 /// as they had in the source (`--from`). This is typically used for undoing
 /// changes to some paths in the working copy (`jj restore <paths>`).
 ///
+/// When neither `--from` nor `--to` is specified, the command restores into the
+/// working copy from its parent. If one of `--from` or `--to` is specified, the
+/// other one defaults to the working copy.
+///
 /// If you restore from a revision where the path has conflicts, then the
 /// destination revision will have the same conflict. If the destination is the
 /// working copy, then a new commit will be created on top for resolving the
@@ -1327,11 +1331,11 @@ struct UnsquashArgs {
 #[derive(clap::Args, Clone, Debug)]
 struct RestoreArgs {
     /// Revision to restore from (source)
-    #[clap(long, default_value = "@-")]
-    from: String,
+    #[clap(long)]
+    from: Option<String>,
     /// Revision to restore into (destination)
-    #[clap(long, default_value = "@")]
-    to: String,
+    #[clap(long)]
+    to: Option<String>,
     /// Interactively choose which parts to restore
     #[clap(long, short)]
     interactive: bool,
@@ -3410,8 +3414,14 @@ fn cmd_restore(
     args: &RestoreArgs,
 ) -> Result<(), CommandError> {
     let mut workspace_command = command.workspace_helper(ui)?;
-    let from_commit = workspace_command.resolve_single_rev(ui, &args.from)?;
-    let to_commit = workspace_command.resolve_single_rev(ui, &args.to)?;
+    let (from_str, to_str) = match (args.from.as_deref(), args.to.as_deref()) {
+        (None, None) => ("@-", "@"),
+        (Some(from), None) => (from, "@"),
+        (None, Some(to)) => ("@", to),
+        (Some(from), Some(to)) => (from, to),
+    };
+    let from_commit = workspace_command.resolve_single_rev(ui, from_str)?;
+    let to_commit = workspace_command.resolve_single_rev(ui, to_str)?;
     workspace_command.check_rewriteable(&to_commit)?;
     let repo = workspace_command.repo();
     let tree_id;
diff --git a/tests/test_restore_command.rs b/tests/test_restore_command.rs
index dd48f7330..74c8c55b7 100644
--- a/tests/test_restore_command.rs
+++ b/tests/test_restore_command.rs
@@ -53,11 +53,8 @@ fn test_restore() {
 ");
 
     // Can restore into other revision
-    // TODO: Maybe we should make `--from @` be the default when using `--to`. That
-    // would probably mean declaring both `--from` and `--to` optional and
-    // without default values and handle the defaults in code.
     test_env.jj_cmd_success(&repo_path, &["undo"]);
-    let stdout = test_env.jj_cmd_success(&repo_path, &["restore", "--from", "@", "--to", "@-"]);
+    let stdout = test_env.jj_cmd_success(&repo_path, &["restore", "--to", "@-"]);
     insta::assert_snapshot!(stdout, @r###"
     Created 5ed06151e039 
     Rebased 1 descendant commits
@@ -72,12 +69,29 @@ fn test_restore() {
     A file3
     "###);
 
+    // Can combine `--from` and `--to`
+    test_env.jj_cmd_success(&repo_path, &["undo"]);
+    let stdout = test_env.jj_cmd_success(&repo_path, &["restore", "--from", "@", "--to", "@-"]);
+    insta::assert_snapshot!(stdout, @r###"
+    Created c83e17dc46fd 
+    Rebased 1 descendant commits
+    Working copy now at: df9fb6892f99 
+    "###);
+    let stdout = test_env.jj_cmd_success(&repo_path, &["diff", "-s"]);
+    insta::assert_snapshot!(stdout, @"");
+    let stdout = test_env.jj_cmd_success(&repo_path, &["diff", "-s", "-r", "@-"]);
+    insta::assert_snapshot!(stdout, @r###"
+    R file1
+    A file2
+    A file3
+    "###);
+
     // Can restore only specified paths
     test_env.jj_cmd_success(&repo_path, &["undo"]);
     let stdout = test_env.jj_cmd_success(&repo_path, &["restore", "file2", "file3"]);
     insta::assert_snapshot!(stdout, @r###"
-    Created a5bd47134d7e 
-    Working copy now at: a5bd47134d7e 
+    Created 28647642d4a5 
+    Working copy now at: 28647642d4a5 
     Added 0 files, modified 1 files, removed 1 files
     "###);
     let stdout = test_env.jj_cmd_success(&repo_path, &["diff", "-s"]);