From 3419adf4f2de3adb14e2ed5a628f9eb34fd53b11 Mon Sep 17 00:00:00 2001
From: Ilya Grigoriev <ilyagr@users.noreply.github.com>
Date: Wed, 31 May 2023 20:55:46 -0700
Subject: [PATCH] undo: preserve git-tracking refs in colocated repos by
 default

---
 src/commands/operation.rs   | 22 ++++++++++++++--------
 tests/test_git_colocated.rs |  6 ++----
 tests/test_undo.rs          |  6 ++++--
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/commands/operation.rs b/src/commands/operation.rs
index 254cf8491e..1aaa89ad2c 100644
--- a/src/commands/operation.rs
+++ b/src/commands/operation.rs
@@ -206,15 +206,19 @@ fn view_with_desired_portions_restored(
     new_view
 }
 
-fn process_what_arg(what_arg: &[UndoWhatToRestore]) -> BTreeSet<UndoWhatToRestore> {
+fn process_what_arg(
+    what_arg: &[UndoWhatToRestore],
+    colocated: bool,
+) -> BTreeSet<UndoWhatToRestore> {
     if !what_arg.is_empty() {
         what_arg.iter().cloned().collect()
     } else {
-        btreeset!(
-            UndoWhatToRestore::Repo,
-            UndoWhatToRestore::RemoteTracking,
-            UndoWhatToRestore::GitTracking
-        )
+        let mut default_what =
+            btreeset!(UndoWhatToRestore::Repo, UndoWhatToRestore::RemoteTracking);
+        if !colocated {
+            default_what.insert(UndoWhatToRestore::GitTracking);
+        }
+        default_what
     }
 }
 
@@ -225,6 +229,7 @@ pub fn cmd_op_undo(
 ) -> Result<(), CommandError> {
     let mut workspace_command = command.workspace_helper(ui)?;
     let bad_op = workspace_command.resolve_single_op(&args.operation)?;
+    let repo_is_colocated = workspace_command.working_copy_shared_with_git();
     let parent_ops = bad_op.parents();
     if parent_ops.len() > 1 {
         return Err(user_error("Cannot undo a merge operation"));
@@ -242,7 +247,7 @@ pub fn cmd_op_undo(
     let new_view = view_with_desired_portions_restored(
         tx.repo().view().store_view(),
         tx.base_repo().view().store_view(),
-        process_what_arg(&args.what),
+        process_what_arg(&args.what, repo_is_colocated),
     );
     tx.mut_repo().set_view(new_view);
     tx.finish(ui)?;
@@ -257,12 +262,13 @@ fn cmd_op_restore(
 ) -> Result<(), CommandError> {
     let mut workspace_command = command.workspace_helper(ui)?;
     let target_op = workspace_command.resolve_single_op(&args.operation)?;
+    let repo_is_colocated = workspace_command.working_copy_shared_with_git();
     let mut tx = workspace_command
         .start_transaction(&format!("restore to operation {}", target_op.id().hex()));
     let new_view = view_with_desired_portions_restored(
         target_op.view().store_view(),
         tx.base_repo().view().store_view(),
-        process_what_arg(&args.what),
+        process_what_arg(&args.what, repo_is_colocated),
     );
     tx.mut_repo().set_view(new_view);
     tx.finish(ui)?;
diff --git a/tests/test_git_colocated.rs b/tests/test_git_colocated.rs
index c481dc6f8c..16405ff990 100644
--- a/tests/test_git_colocated.rs
+++ b/tests/test_git_colocated.rs
@@ -374,10 +374,8 @@ fn test_git_colocated_squash_undo() {
     // TODO: There should be no divergence here; 2f376ea1478c should be hidden
     // (#922)
     insta::assert_snapshot!(get_log_output_divergence(&test_env, &repo_path), @r###"
-    ◉  qpvuntsmwlqt 2f376ea1478c A master !divergence!
-    │ @  rlvkpnrzqnoo 8f71e3b6a3be
-    │ ◉  qpvuntsmwlqt a86754f975f9 A HEAD@git !divergence!
-    ├─╯
+    @  rlvkpnrzqnoo 8f71e3b6a3be
+    ◉  qpvuntsmwlqt a86754f975f9 A master HEAD@git
     ◉  zzzzzzzzzzzz 000000000000
     "###);
 }
diff --git a/tests/test_undo.rs b/tests/test_undo.rs
index 68ae1732e1..aa6033b1b7 100644
--- a/tests/test_undo.rs
+++ b/tests/test_undo.rs
@@ -145,10 +145,12 @@ fn test_git_push_undo_colocated() {
     test_env.advance_test_rng_seed_to_multiple_of(100_000);
     test_env.jj_cmd_success(&repo_path, &["describe", "-m", "CC"]);
     test_env.jj_cmd_success(&repo_path, &["git", "fetch"]);
-    // This currently gives an identical result to `test_git_push_undo_import`
+    // This seems like an OK behavior: `jj` forgot that the remote-tracking branch
+    // was undone, but keeps track of the latest local version. TODO: Would this
+    // be improved if `jj git export` exported remote-tracking branches?
     insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
     main: 0a3e99f08a48 CC
-      @origin (ahead by 1 commits, behind by 1 commits): 8c05de152218 BB
+      @origin (ahead by 1 commits, behind by 1 commits): 0cffb6146141 AA
     "###);
 }