cmd: Make jj restore work if @ is a merge commit

To be clear, this applies to `jj restore` without any arguments.

Fixes #1228

CHANGELOG.md

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * `jj git push --deleted` will remove all locally deleted branches from the remote.
 
+* `jj restore` without `--from` works correctly even if `@` is a merge
+  commit.
+
 ### Fixed bugs
 
 * Modify/delete conflicts now include context lines

src/commands/mod.rs

@@ -2492,24 +2492,27 @@ fn cmd_restore(
     args: &RestoreArgs,
 ) -> Result<(), CommandError> {
     let mut workspace_command = command.workspace_helper(ui)?;
-    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(from_str)?;
-    let to_commit = workspace_command.resolve_single_rev(to_str)?;
+    let (from_tree, to_commit);
+    if args.from.is_some() || args.to.is_some() {
+        to_commit = workspace_command.resolve_single_rev(args.to.as_deref().unwrap_or("@"))?;
+        from_tree = workspace_command
+            .resolve_single_rev(args.from.as_deref().unwrap_or("@"))?
+            .tree();
+    } else {
+        to_commit = workspace_command.resolve_single_rev("@")?;
+        from_tree = merge_commit_trees(workspace_command.repo(), &to_commit.parents());
+    }
     workspace_command.check_rewritable(&to_commit)?;
-    let tree_id = if args.paths.is_empty() {
-        from_commit.tree_id().clone()
+
+    let new_tree_id = if args.paths.is_empty() {
+        from_tree.id().clone()
     } else {
         let matcher = workspace_command.matcher_from_values(&args.paths)?;
         let mut tree_builder = workspace_command
             .repo()
             .store()
             .tree_builder(to_commit.tree_id().clone());
-        for (repo_path, diff) in from_commit.tree().diff(&to_commit.tree(), matcher.as_ref()) {
+        for (repo_path, diff) in from_tree.diff(&to_commit.tree(), matcher.as_ref()) {
             match diff.into_options().0 {
                 Some(value) => {
                     tree_builder.set(repo_path, value);
@@ -2521,7 +2524,7 @@ fn cmd_restore(
         }
         tree_builder.write_tree()
     };
-    if &tree_id == to_commit.tree_id() {
+    if &new_tree_id == to_commit.tree_id() {
         ui.write("Nothing changed.\n")?;
     } else {
         let mut tx = workspace_command
@@ -2529,7 +2532,7 @@ fn cmd_restore(
         let mut_repo = tx.mut_repo();
         let new_commit = mut_repo
             .rewrite_commit(command.settings(), &to_commit)
-            .set_tree(tree_id)
+            .set_tree(new_tree_id)
             .write()?;
         ui.write("Created ")?;
         tx.write_commit_summary(ui.stdout_formatter().as_mut(), &new_commit)?;

tests/test_restore_command.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::path::Path;
+
 use crate::common::TestEnvironment;
 
 pub mod common;
@@ -100,3 +102,131 @@ fn test_restore() {
     R file1
     "###);
 }
+
+// Much of this test is copied from test_resolve_command
+#[test]
+fn test_restore_conflicted_merge() {
+    let test_env = TestEnvironment::default();
+    test_env.jj_cmd_success(test_env.env_root(), &["init", "repo", "--git"]);
+    let repo_path = test_env.env_root().join("repo");
+
+    create_commit(&test_env, &repo_path, "base", &[], &[("file", "base\n")]);
+    create_commit(&test_env, &repo_path, "a", &["base"], &[("file", "a\n")]);
+    create_commit(&test_env, &repo_path, "b", &["base"], &[("file", "b\n")]);
+    create_commit(&test_env, &repo_path, "conflict", &["a", "b"], &[]);
+    // Test the setup
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    @    conflict
+    ├─╮
+    o │  b
+    │ o  a
+    ├─╯
+    o  base
+    o
+    "###);
+    insta::assert_snapshot!(
+    std::fs::read_to_string(repo_path.join("file")).unwrap()
+        , @r###"
+            <<<<<<<
+            %%%%%%%
+            -base
+            +a
+            +++++++
+            b
+            >>>>>>>
+            "###);
+
+    // Overwrite the file...
+    std::fs::write(repo_path.join("file"), "resolution").unwrap();
+    insta::assert_snapshot!(test_env.jj_cmd_success(&repo_path, &["diff"]), 
+    @r###"
+    Resolved conflict in file:
+       1     : <<<<<<<
+       2     : %%%%%%%
+       3     : -base
+       4     : +a
+       5     : +++++++
+       6     : b
+       7     : >>>>>>>
+            1: resolution
+    "###);
+
+    // ...and restore it back again.
+    let stdout = test_env.jj_cmd_success(&repo_path, &["restore", "file"]);
+    insta::assert_snapshot!(stdout, @r###"
+    Created 63198ca2e4aa conflict
+    Working copy now at: 63198ca2e4aa conflict
+    Added 0 files, modified 1 files, removed 0 files
+    "###);
+    insta::assert_snapshot!(
+    std::fs::read_to_string(repo_path.join("file")).unwrap()
+        , @r###"
+    <<<<<<<
+    %%%%%%%
+    -base
+    +a
+    +++++++
+    b
+    >>>>>>>
+    "###);
+    let stdout = test_env.jj_cmd_success(&repo_path, &["diff"]);
+    insta::assert_snapshot!(stdout, @"");
+
+    // The same, but without the `file` argument. Overwrite the file...
+    std::fs::write(repo_path.join("file"), "resolution").unwrap();
+    insta::assert_snapshot!(test_env.jj_cmd_success(&repo_path, &["diff"]), 
+    @r###"
+    Resolved conflict in file:
+       1     : <<<<<<<
+       2     : %%%%%%%
+       3     : -base
+       4     : +a
+       5     : +++++++
+       6     : b
+       7     : >>>>>>>
+            1: resolution
+    "###);
+
+    // ... and restore it back again.
+    let stdout = test_env.jj_cmd_success(&repo_path, &["restore"]);
+    insta::assert_snapshot!(stdout, @r###"
+    Created d955febceac1 conflict
+    Working copy now at: d955febceac1 conflict
+    Added 0 files, modified 1 files, removed 0 files
+    "###);
+    insta::assert_snapshot!(
+    std::fs::read_to_string(repo_path.join("file")).unwrap()
+        , @r###"
+    <<<<<<<
+    %%%%%%%
+    -base
+    +a
+    +++++++
+    b
+    >>>>>>>
+    "###);
+}
+
+fn create_commit(
+    test_env: &TestEnvironment,
+    repo_path: &Path,
+    name: &str,
+    parents: &[&str],
+    files: &[(&str, &str)],
+) {
+    if parents.is_empty() {
+        test_env.jj_cmd_success(repo_path, &["new", "root", "-m", name]);
+    } else {
+        let mut args = vec!["new", "-m", name];
+        args.extend(parents);
+        test_env.jj_cmd_success(repo_path, &args);
+    }
+    for (name, content) in files {
+        std::fs::write(repo_path.join(name), content).unwrap();
+    }
+    test_env.jj_cmd_success(repo_path, &["branch", "create", name]);
+}
+
+fn get_log_output(test_env: &TestEnvironment, repo_path: &Path) -> String {
+    test_env.jj_cmd_success(repo_path, &["log", "-T", "branches"])
+}