undo: options to preserve git refs and/or remote-tracking branches on…

… undo or restore

src/commands/operation.rs

@@ -1,5 +1,10 @@
+use std::collections::{BTreeMap, BTreeSet};
+
 use clap::Subcommand;
+use jujutsu_lib::op_store::BranchTarget;
 use jujutsu_lib::operation;
+use jujutsu_lib::repo::Repo;
+use maplit::btreeset;
 
 use crate::cli_util::{user_error, CommandError, CommandHelper, LogContentFormat};
 use crate::graphlog::{get_graphlog, Edge};
@@ -40,6 +45,17 @@ pub struct OperationRestoreArgs {
     /// --at-op=<operation ID> log` before restoring to an operation to see the
     /// state of the repo at that operation.
     operation: String,
+
+    /// What portions of the local state to restore (can be repeated)
+    ///
+    /// Defaults to everything for non-colocated repos.
+    ///
+    /// Defaults to `repo` and `remote-tracking` for colocated repos. This
+    /// ensures that the automatic `jj git export` succeeds.
+    ///
+    /// This option is EXPERIMENTAL.
+    #[arg(long)]
+    what: Vec<UndoWhatToRestore>,
 }
 
 /// Create a new operation that undoes an earlier operation
@@ -53,6 +69,28 @@ pub struct OperationUndoArgs {
     /// Use `jj op log` to find an operation to undo.
     #[arg(default_value = "@")]
     operation: String,
+
+    /// What portions of the local state to restore (can be repeated)
+    ///
+    /// Defaults to everything for non-colocated repos.
+    ///
+    /// Defaults to `repo` and `remote-tracking` for colocated repos. This
+    /// ensures that the automatic `jj git export` succeeds.
+    ///
+    /// This option is EXPERIMENTAL.
+    #[arg(long)]
+    what: Vec<UndoWhatToRestore>,
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, clap::ValueEnum)]
+enum UndoWhatToRestore {
+    /// The jj repo state and local branches
+    Repo,
+    /// The remote-tracking branches. Do not restore these if you'd like to push
+    /// after the undo
+    RemoteTracking,
+    /// Remembered git repo state from the last `jj git import`
+    GitTracking,
 }
 
 fn cmd_op_log(
@@ -112,6 +150,70 @@ fn cmd_op_log(
     Ok(())
 }
 
+/// Restore only the portions of the view specified by the `what` argument
+fn view_with_desired_portions_restored(
+    view_being_restored: &jujutsu_lib::op_store::View,
+    current_view: &jujutsu_lib::op_store::View,
+    what: BTreeSet<UndoWhatToRestore>,
+) -> jujutsu_lib::op_store::View {
+    let mut new_view = if what.contains(&UndoWhatToRestore::Repo) {
+        view_being_restored.clone()
+    } else {
+        current_view.clone()
+    };
+    new_view.git_refs = if what.contains(&UndoWhatToRestore::GitTracking) {
+        view_being_restored.git_refs.clone()
+    } else {
+        current_view.git_refs.clone()
+    };
+
+    let mut all_branch_names: BTreeSet<_> = view_being_restored.branches.keys().collect();
+    all_branch_names.append(&mut current_view.branches.keys().collect());
+
+    let branch_source_view = if what.contains(&UndoWhatToRestore::RemoteTracking) {
+        view_being_restored
+    } else {
+        current_view
+    };
+    // Short-term TODO: we will optimize this to avoid recreating branches if they
+    // are already correct in `new_view`
+    let mut new_branches = BTreeMap::default();
+    for branch_name in all_branch_names {
+        let local_target = new_view
+            .branches
+            .get(branch_name)
+            .and_then(|br| br.local_target.clone());
+        let remote_targets = branch_source_view
+            .branches
+            .get(branch_name)
+            .map(|br| br.remote_targets.clone())
+            .unwrap_or_default();
+        if local_target.is_some() || !remote_targets.is_empty() {
+            new_branches.insert(
+                branch_name.to_string(),
+                BranchTarget {
+                    local_target,
+                    remote_targets,
+                },
+            );
+        }
+    }
+    new_view.branches = new_branches;
+    new_view
+}
+
+fn process_what_arg(what_arg: &[UndoWhatToRestore]) -> BTreeSet<UndoWhatToRestore> {
+    if !what_arg.is_empty() {
+        what_arg.iter().cloned().collect()
+    } else {
+        btreeset!(
+            UndoWhatToRestore::Repo,
+            UndoWhatToRestore::RemoteTracking,
+            UndoWhatToRestore::GitTracking
+        )
+    }
+}
+
 pub fn cmd_op_undo(
     ui: &mut Ui,
     command: &CommandHelper,
@@ -133,6 +235,12 @@ pub fn cmd_op_undo(
     let bad_repo = repo_loader.load_at(&bad_op);
     let parent_repo = repo_loader.load_at(&parent_ops[0]);
     tx.mut_repo().merge(&bad_repo, &parent_repo);
+    let new_view = view_with_desired_portions_restored(
+        tx.repo().view().store_view(),
+        tx.base_repo().view().store_view(),
+        process_what_arg(&args.what),
+    );
+    tx.mut_repo().set_view(new_view);
     tx.finish(ui)?;
 
     Ok(())
@@ -147,7 +255,12 @@ fn cmd_op_restore(
     let target_op = workspace_command.resolve_single_op(&args.operation)?;
     let mut tx = workspace_command
         .start_transaction(&format!("restore to operation {}", target_op.id().hex()));
-    tx.mut_repo().set_view(target_op.view().take_store_view());
+    let new_view = view_with_desired_portions_restored(
+        target_op.view().store_view(),
+        tx.base_repo().view().store_view(),
+        process_what_arg(&args.what),
+    );
+    tx.mut_repo().set_view(new_view);
     tx.finish(ui)?;
 
     Ok(())

tests/test_git_fetch.rs

@@ -50,6 +50,13 @@ fn get_branch_output(test_env: &TestEnvironment, repo_path: &Path) -> String {
     test_env.jj_cmd_success(repo_path, &["branch", "list"])
 }
 
+fn current_operation_id(test_env: &TestEnvironment, repo_path: &Path) -> String {
+    let mut id = test_env.jj_cmd_success(repo_path, &["debug", "operation", "--display=id"]);
+    let len_trimmed = id.trim_end().len();
+    id.truncate(len_trimmed);
+    id
+}
+
 fn create_commit(test_env: &TestEnvironment, repo_path: &Path, name: &str, parents: &[&str]) {
     let descr = format!("descr_for_{name}");
     if parents.is_empty() {
@@ -608,6 +615,8 @@ fn test_git_fetch_some_of_many_branches() {
     "###);
 }
 
+// See `test_undo_restore_commands.rs` for fetch-undo-push and fetch-undo-fetch
+// of the same branches for various kinds of undo.
 #[test]
 fn test_git_fetch_undo() {
     let test_env = TestEnvironment::default();
@@ -669,6 +678,96 @@ fn test_git_fetch_undo() {
     "###);
 }
 
+// Compare to `test_git_import_undo` in test_git_import_export
+// TODO: Explain why these behaviors are useful
+#[test]
+fn test_fetch_undo_what() {
+    let test_env = TestEnvironment::default();
+    let source_git_repo_path = test_env.env_root().join("source");
+    let _git_repo = git2::Repository::init(source_git_repo_path.clone()).unwrap();
+
+    // Clone an empty repo. The target repo is a normal `jj` repo, *not* colocated
+    let stdout =
+        test_env.jj_cmd_success(test_env.env_root(), &["git", "clone", "source", "target"]);
+    insta::assert_snapshot!(stdout, @r###"
+    Fetching into new repo in "$TEST_ENV/target"
+    Nothing changed.
+    "###);
+    let repo_path = test_env.env_root().join("target");
+
+    let source_log =
+        create_colocated_repo_and_branches_from_trunk1(&test_env, &source_git_repo_path);
+    insta::assert_snapshot!(source_log, @r###"
+       ===== Source git repo contents =====
+    @  c7d4bdcbc215 descr_for_b b
+    │ ◉  decaa3966c83 descr_for_a2 a2
+    ├─╯
+    │ ◉  359a9a02457d descr_for_a1 a1
+    ├─╯
+    ◉  ff36dc55760e descr_for_trunk1 master trunk1
+    ◉  000000000000
+    "###);
+
+    // Initial state we will try to return to after `op restore`. There are no
+    // branches.
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @"");
+    let base_operation_id = current_operation_id(&test_env, &repo_path);
+
+    // Fetch a branch
+    let stdout = test_env.jj_cmd_success(&repo_path, &["git", "fetch", "--branch", "b"]);
+    insta::assert_snapshot!(stdout, @"");
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    ◉  c7d4bdcbc215 descr_for_b b
+    ◉  ff36dc55760e descr_for_trunk1
+    │ @  230dd059e1b0
+    ├─╯
+    ◉  000000000000
+    "###);
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    b: c7d4bdcbc215 descr_for_b
+    "###);
+
+    // We can undo the change in the repo without moving the remote-tracking branch
+    let stdout = test_env.jj_cmd_success(
+        &repo_path,
+        &["op", "restore", "--what", "repo", &base_operation_id],
+    );
+    insta::assert_snapshot!(stdout, @"");
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    b (deleted)
+      @origin: c7d4bdcbc215 descr_for_b
+      (this branch will be *deleted permanently* on the remote on the
+       next `jj git push`. Use `jj branch forget` to prevent this)
+    "###);
+
+    // Now, let's demo restoring just the remote-tracking branch. First, let's
+    // change our local repo state...
+    test_env.jj_cmd_success(&repo_path, &["branch", "c", "newbranch"]);
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    b (deleted)
+      @origin: c7d4bdcbc215 descr_for_b
+      (this branch will be *deleted permanently* on the remote on the
+       next `jj git push`. Use `jj branch forget` to prevent this)
+    newbranch: 230dd059e1b0 (no description set)
+    "###);
+    // Restoring just the remote-tracking state will not affect `newbranch`, but
+    // will eliminate `b@origin`.
+    let stdout = test_env.jj_cmd_success(
+        &repo_path,
+        &[
+            "op",
+            "restore",
+            "--what",
+            "remote-tracking",
+            &base_operation_id,
+        ],
+    );
+    insta::assert_snapshot!(stdout, @"");
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    newbranch: 230dd059e1b0 (no description set)
+    "###);
+}
+
 #[test]
 fn test_git_fetch_remove_fetch() {
     let test_env = TestEnvironment::default();

tests/test_git_import_export.rs

@@ -145,7 +145,7 @@ fn test_git_import_undo() {
     a: 230dd059e1b0 (no description set)
     "###);
 
-    // "git import" can be undone.
+    // "git import" can be undone by default in non-colocated repositories.
     let stdout = test_env.jj_cmd_success(&repo_path, &["op", "restore", &base_operation_id]);
     insta::assert_snapshot!(stdout, @r###"
     "###);
@@ -155,10 +155,53 @@ fn test_git_import_undo() {
     insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
     a: 230dd059e1b0 (no description set)
     "###);
+
+    // If we don't restore the git_refs, undoing the import removes the local branch
+    // but makes a following import a no-op.
+    let stdout = test_env.jj_cmd_success(
+        &repo_path,
+        &[
+            "op",
+            "restore",
+            &base_operation_id,
+            "--what=repo",
+            "--what=remote-tracking",
+        ],
+    );
+    insta::assert_snapshot!(stdout, @r###"
+    "###);
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    a (deleted)
+      @git: 230dd059e1b0 (no description set)
+    "###);
+    // Try "git import" again, which should *not* re-import the branch "a" and be a
+    // no-op.
+    insta::assert_snapshot!(test_env.jj_cmd_success(&repo_path, &["git", "import"]), @r###"
+    Nothing changed.
+    "###);
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    a (deleted)
+      @git: 230dd059e1b0 (no description set)
+    "###);
+
+    // We can restore *only* the git refs to make an import re-import the branch
+    let stdout = test_env.jj_cmd_success(
+        &repo_path,
+        &["op", "restore", &base_operation_id, "--what=git-tracking"],
+    );
+    insta::assert_snapshot!(stdout, @r###"
+    "###);
+    // The git-tracking branch disappears.
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @"");
+    // Try "git import" again, which should again re-import the branch "a".
+    insta::assert_snapshot!(test_env.jj_cmd_success(&repo_path, &["git", "import"]), @"");
+    insta::assert_snapshot!(get_branch_output(&test_env, &repo_path), @r###"
+    a: 230dd059e1b0 (no description set)
+    "###);
 }
 
 #[test]
-fn test_git_import_move_export_undo() {
+fn test_git_import_move_export_with_default_undo() {
     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");

tests/test_undo.rs

@@ -145,9 +145,42 @@ 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` (NOT
-    // `test_git_push_undo` because of the automatic import). However, a
-    // follow-up commit will make the two tests behave differently.
+    // This currently gives an identical result to `test_git_push_undo_import`
+    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
+    "###);
+}
+
+// This test is currently *identical* to `test_git_push_undo` except
+// both the git_refs and the remote-tracking branches are preserved by undo.
+// TODO: Investigate the different outcome
+#[test]
+fn test_git_push_undo_repo_only() {
+    let test_env = TestEnvironment::default();
+    let git_repo_path = test_env.env_root().join("git-repo");
+    git2::Repository::init_bare(git_repo_path).unwrap();
+    test_env.jj_cmd_success(test_env.env_root(), &["git", "clone", "git-repo", "repo"]);
+    let repo_path = test_env.env_root().join("repo");
+
+    test_env.advance_test_rng_seed_to_multiple_of(100_000);
+    test_env.jj_cmd_success(&repo_path, &["branch", "create", "main"]);
+    test_env.jj_cmd_success(&repo_path, &["describe", "-m", "AA"]);
+    test_env.jj_cmd_success(&repo_path, &["git", "push"]);
+    test_env.advance_test_rng_seed_to_multiple_of(100_000);
+    test_env.jj_cmd_success(&repo_path, &["describe", "-m", "BB"]);
+    let pre_push_opid = current_operation_id(&test_env, &repo_path);
+    test_env.jj_cmd_success(&repo_path, &["git", "push"]);
+
+    // Undo the push, but keep both the git_refs and the remote-tracking branches
+    test_env.jj_cmd_success(
+        &repo_path,
+        &["op", "restore", "--what=repo", &pre_push_opid],
+    );
+    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`.
     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