cli: move logic for updating branches after rewrite to lib crate

This patch moves the function for updating branches after rewrite from
`commands.rs` into `rewrite.rs`.

It also changes the function to update branches even if they were
conflicted or become conflicted. I think that seems better than
leaving branches on old commits. For example, let's say you have start
with this:

```
C main
|
B origin@main
|
A
```

You now pull from origin, which has updated the main branch from B to
B'. We apply that change to both the remote branch and the local
branch, which results in a conflict in the local branch:

```
C main?
|
B B' main? origin@main
|/
A
```

If you now rewrite C to C', the conflicted main branch will still
point to C, which is just weird. This patch changes that so the
conflicted side of main gets repointed to C'.

I also refactored the code to reuse our existing
`MutableRepo::merge_single_ref()`, which improves the behavior in
several cases, such as the conflict-resolution case in the last test
case.

lib/src/rewrite.rs

@@ -19,11 +19,13 @@ use itertools::Itertools;
 use crate::backend::CommitId;
 use crate::commit::Commit;
 use crate::commit_builder::CommitBuilder;
+use crate::op_store::RefTarget;
 use crate::repo::{MutableRepo, RepoRef};
 use crate::repo_path::RepoPath;
 use crate::revset::RevsetExpression;
 use crate::settings::UserSettings;
 use crate::tree::{merge_trees, Tree};
+use crate::view::RefName;
 
 pub fn merge_commit_trees(repo: RepoRef, commits: &[Commit]) -> Tree {
     let store = repo.store();
@@ -228,3 +230,55 @@ pub enum RebasedDescendant {
         new_commit: Commit,
     },
 }
+
+pub fn update_branches_after_rewrite(mut_repo: &mut MutableRepo) {
+    let new_evolution = mut_repo.evolution();
+    let base_repo = mut_repo.base_repo();
+    let old_evolution = base_repo.evolution();
+    let mut updates = vec![];
+    let index = mut_repo.index().as_index_ref();
+
+    let ref_target_update = |old_id: CommitId| -> Option<(RefTarget, RefTarget)> {
+        if new_evolution.is_obsolete(&old_id) && !old_evolution.is_obsolete(&old_id) {
+            // The call to index.heads() is mostly to get a predictable order
+            let new_ids = index.heads(&new_evolution.new_parent(mut_repo.as_repo_ref(), &old_id));
+            let old_ids = std::iter::repeat(old_id).take(new_ids.len()).collect_vec();
+            Some((
+                RefTarget::Conflict {
+                    removes: vec![],
+                    adds: old_ids,
+                },
+                RefTarget::Conflict {
+                    removes: vec![],
+                    adds: new_ids,
+                },
+            ))
+        } else {
+            None
+        }
+    };
+
+    for (branch_name, branch_target) in mut_repo.view().branches() {
+        if let Some(old_target) = &branch_target.local_target {
+            for old_add in old_target.adds() {
+                if let Some((old_target, new_target)) = ref_target_update(old_add) {
+                    updates.push((branch_name.clone(), old_target, new_target));
+                }
+            }
+            for old_remove in old_target.removes() {
+                if let Some((old_target, new_target)) = ref_target_update(old_remove) {
+                    // Arguments reversed for removes
+                    updates.push((branch_name.clone(), new_target, old_target));
+                }
+            }
+        }
+    }
+
+    for (branch_name, old_target, new_target) in updates {
+        mut_repo.merge_single_ref(
+            &RefName::LocalBranch(branch_name),
+            Some(&old_target),
+            Some(&new_target),
+        );
+    }
+}

lib/tests/test_rewrite.rs

@@ -15,8 +15,9 @@
 use jujutsu_lib::backend::CommitId;
 use jujutsu_lib::commit::Commit;
 use jujutsu_lib::commit_builder::CommitBuilder;
+use jujutsu_lib::op_store::RefTarget;
 use jujutsu_lib::repo_path::RepoPath;
-use jujutsu_lib::rewrite::{DescendantRebaser, RebasedDescendant};
+use jujutsu_lib::rewrite::{update_branches_after_rewrite, DescendantRebaser, RebasedDescendant};
 use jujutsu_lib::testutils;
 use jujutsu_lib::testutils::CommitGraphBuilder;
 use maplit::hashmap;
@@ -539,3 +540,204 @@ fn test_rebase_descendants_contents(use_git: bool) {
 
     tx.discard();
 }
+
+#[test]
+fn test_update_branches_after_rewrite_basic() {
+    let settings = testutils::user_settings();
+    let (_temp_dir, repo) = testutils::init_repo(&settings, false);
+
+    // Branch "main" points to branch B. B gets rewritten as B2. Branch main should
+    // be updated to point to B2.
+    //
+    // B main         B2 main
+    // |         =>   |
+    // A              A
+    let mut tx = repo.start_transaction("test");
+    let mut graph_builder = CommitGraphBuilder::new(&settings, tx.mut_repo());
+    let commit_a = graph_builder.initial_commit();
+    let commit_b = graph_builder.commit_with_parents(&[&commit_a]);
+    tx.mut_repo()
+        .set_local_branch("main".to_string(), RefTarget::Normal(commit_b.id().clone()));
+    tx.mut_repo().set_remote_branch(
+        "main".to_string(),
+        "origin".to_string(),
+        RefTarget::Normal(commit_b.id().clone()),
+    );
+    tx.mut_repo()
+        .set_tag("v1".to_string(), RefTarget::Normal(commit_b.id().clone()));
+    let repo = tx.commit();
+
+    let mut tx = repo.start_transaction("test");
+    let commit_b2 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .write_to_repo(tx.mut_repo());
+    update_branches_after_rewrite(tx.mut_repo());
+    assert_eq!(
+        tx.mut_repo().get_local_branch("main"),
+        Some(RefTarget::Normal(commit_b2.id().clone()))
+    );
+    // The remote branch and tag should not get updated
+    assert_eq!(
+        tx.mut_repo().get_remote_branch("main", "origin"),
+        Some(RefTarget::Normal(commit_b.id().clone()))
+    );
+    assert_eq!(
+        tx.mut_repo().get_tag("v1"),
+        Some(RefTarget::Normal(commit_b.id().clone()))
+    );
+
+    tx.discard();
+}
+
+#[test]
+fn test_update_branches_after_rewrite_to_conflict() {
+    let settings = testutils::user_settings();
+    let (_temp_dir, repo) = testutils::init_repo(&settings, false);
+
+    // Branch "main" points to commit B. B gets rewritten as B2, B3, B4. Branch main
+    // should become a conflict pointing to all of them.
+    //
+    //                B4 main?
+    //                | B3 main?
+    // B main         |/B2 main?
+    // |         =>   |/
+    // A              A
+    let mut tx = repo.start_transaction("test");
+    let mut graph_builder = CommitGraphBuilder::new(&settings, tx.mut_repo());
+    let commit_a = graph_builder.initial_commit();
+    let commit_b = graph_builder.commit_with_parents(&[&commit_a]);
+    tx.mut_repo()
+        .set_local_branch("main".to_string(), RefTarget::Normal(commit_b.id().clone()));
+    let repo = tx.commit();
+
+    let mut tx = repo.start_transaction("test");
+    let commit_b2 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .write_to_repo(tx.mut_repo());
+    // Different description so they're not the same commit
+    let commit_b3 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .set_description("different".to_string())
+        .write_to_repo(tx.mut_repo());
+    // Different description so they're not the same commit
+    let commit_b4 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .set_description("more different".to_string())
+        .write_to_repo(tx.mut_repo());
+    update_branches_after_rewrite(tx.mut_repo());
+    assert_eq!(
+        tx.mut_repo().get_local_branch("main"),
+        Some(RefTarget::Conflict {
+            removes: vec![commit_b.id().clone(), commit_b.id().clone()],
+            adds: vec![
+                commit_b2.id().clone(),
+                commit_b3.id().clone(),
+                commit_b4.id().clone()
+            ]
+        })
+    );
+
+    tx.discard();
+}
+
+#[test]
+fn test_update_branches_after_rewrite_update_conflict() {
+    let settings = testutils::user_settings();
+    let (_temp_dir, repo) = testutils::init_repo(&settings, false);
+
+    // Branch "main" is a conflict removing commit A and adding commit B and C.
+    // A gets rewritten as A2 and A3. B gets rewritten as B2 and B2. The branch
+    // should become a conflict removing A2, A3, and B, and adding A, B2, B3, C.
+    let mut tx = repo.start_transaction("test");
+    let mut graph_builder = CommitGraphBuilder::new(&settings, tx.mut_repo());
+    let commit_a = graph_builder.initial_commit();
+    let commit_b = graph_builder.initial_commit();
+    let commit_c = graph_builder.initial_commit();
+    tx.mut_repo().set_local_branch(
+        "main".to_string(),
+        RefTarget::Conflict {
+            removes: vec![commit_a.id().clone()],
+            adds: vec![commit_b.id().clone(), commit_c.id().clone()],
+        },
+    );
+    let repo = tx.commit();
+
+    let mut tx = repo.start_transaction("test");
+    let commit_a2 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_a)
+        .write_to_repo(tx.mut_repo());
+    // Different description so they're not the same commit
+    let commit_a3 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_a)
+        .set_description("different".to_string())
+        .write_to_repo(tx.mut_repo());
+    let commit_b2 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .write_to_repo(tx.mut_repo());
+    // Different description so they're not the same commit
+    let commit_b3 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .set_description("different".to_string())
+        .write_to_repo(tx.mut_repo());
+    update_branches_after_rewrite(tx.mut_repo());
+    assert_eq!(
+        tx.mut_repo().get_local_branch("main"),
+        Some(RefTarget::Conflict {
+            removes: vec![
+                commit_b.id().clone(),
+                commit_a2.id().clone(),
+                commit_a3.id().clone()
+            ],
+            adds: vec![
+                commit_c.id().clone(),
+                commit_b2.id().clone(),
+                commit_b3.id().clone(),
+                commit_a.id().clone()
+            ]
+        })
+    );
+
+    tx.discard();
+}
+
+#[test]
+fn test_update_branches_after_rewrite_resolves_conflict() {
+    let settings = testutils::user_settings();
+    let (_temp_dir, repo) = testutils::init_repo(&settings, false);
+
+    // Branch "main" is a conflict removing ancestor commit A and adding commit B
+    // and C (maybe it moved forward to B locally and moved forward to C
+    // remotely). Now B gets rewritten as B2, which is a descendant of C (maybe
+    // B was automatically rebased on top of the updated remote). That
+    // would result in a conflict removing A and adding B2 and C. However, since C
+    // is a descendant of A, and B2 is a descendant of C, the conflict gets
+    // resolved to B2.
+    let mut tx = repo.start_transaction("test");
+    let mut graph_builder = CommitGraphBuilder::new(&settings, tx.mut_repo());
+    let commit_a = graph_builder.initial_commit();
+    let commit_b = graph_builder.commit_with_parents(&[&commit_a]);
+    let commit_c = graph_builder.commit_with_parents(&[&commit_a]);
+    tx.mut_repo().set_local_branch(
+        "main".to_string(),
+        RefTarget::Conflict {
+            removes: vec![commit_a.id().clone()],
+            adds: vec![commit_b.id().clone(), commit_c.id().clone()],
+        },
+    );
+    let repo = tx.commit();
+
+    let mut tx = repo.start_transaction("test");
+    let commit_b2 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_b)
+        .set_parents(vec![commit_c.id().clone()])
+        .write_to_repo(tx.mut_repo());
+    update_branches_after_rewrite(tx.mut_repo());
+    assert_eq!(
+        tx.mut_repo().get_local_branch("main"),
+        Some(RefTarget::Normal(commit_b2.id().clone()))
+    );
+
+    tx.discard();
+}
+
+// TODO: Add a test for the following case, which can't happen with our current
+// evolution-based rewriting.
+//
+// 1. Operation 1 points a branch to commit A
+// 2. Operation 2 repoints the branch to commit B
+// 3. Operation 3, which is concurrent with operation 2, deletes the branch
+// 4. Resolved state (operation 4) will have a "-A+B" state for the branch
+//
+// Now we hide B and make A visible instead. When that diff is applied to the
+// branch, the branch state becomes empty and is thus deleted.

src/commands.rs

@@ -52,7 +52,10 @@ use jujutsu_lib::repo::{
 };
 use jujutsu_lib::revset::{RevsetError, RevsetExpression, RevsetParseError};
 use jujutsu_lib::revset_graph_iterator::RevsetGraphEdgeType;
-use jujutsu_lib::rewrite::{back_out_commit, merge_commit_trees, rebase_commit, DescendantRebaser};
+use jujutsu_lib::rewrite::{
+    back_out_commit, merge_commit_trees, rebase_commit, update_branches_after_rewrite,
+    DescendantRebaser,
+};
 use jujutsu_lib::settings::UserSettings;
 use jujutsu_lib::store::Store;
 use jujutsu_lib::transaction::Transaction;
@@ -494,7 +497,7 @@ impl RepoCommandHelper {
             update_checkout_after_rewrite(ui, mut_repo)?;
         }
         if self.auto_update_branches {
-            update_branches_after_rewrite(ui, mut_repo)?;
+            update_branches_after_rewrite(mut_repo);
         }
         self.repo = tx.commit();
         update_working_copy(ui, &self.repo, &self.repo.working_copy_locked())
@@ -709,60 +712,6 @@ fn update_checkout_after_rewrite(ui: &mut Ui, mut_repo: &mut MutableRepo) -> io:
     Ok(())
 }
 
-fn update_branches_after_rewrite(ui: &mut Ui, mut_repo: &mut MutableRepo) -> io::Result<()> {
-    // TODO: Perhaps this method should be in MutableRepo.
-    let new_evolution = mut_repo.evolution();
-    let base_repo = mut_repo.base_repo();
-    let old_evolution = base_repo.evolution();
-    let mut updates = HashMap::new();
-    for (branch_name, branch_target) in mut_repo.view().branches() {
-        match &branch_target.local_target {
-            None => {
-                // nothing to do (a deleted branch doesn't need updating)
-            }
-            Some(RefTarget::Normal(current_target)) => {
-                if new_evolution.is_obsolete(current_target)
-                    && !old_evolution.is_obsolete(current_target)
-                {
-                    let new_targets =
-                        new_evolution.new_parent(mut_repo.as_repo_ref(), current_target);
-                    if new_targets.len() == 1 {
-                        updates.insert(
-                            branch_name.clone(),
-                            RefTarget::Normal(new_targets.iter().next().unwrap().clone()),
-                        );
-                    } else {
-                        writeln!(
-                            ui,
-                            "Branch {}'s target was obsoleted, but the new target is unclear",
-                            branch_name
-                        )?;
-                    }
-                }
-            }
-            Some(RefTarget::Conflict { adds, .. }) => {
-                for current_target in adds {
-                    if new_evolution.is_obsolete(current_target)
-                        && !old_evolution.is_obsolete(current_target)
-                    {
-                        writeln!(
-                            ui,
-                            "Branch {}'s target was obsoleted, but not updating it since it's \
-                             conflicted",
-                            branch_name
-                        )?;
-                    }
-                }
-            }
-        }
-    }
-    for (branch_name, new_local_target) in updates {
-        mut_repo.set_local_branch(branch_name, new_local_target);
-    }
-
-    Ok(())
-}
-
 fn get_app<'a, 'b>() -> App<'a, 'b> {
     let init_command = SubCommand::with_name("init")
         .about("Create a new repo in the given directory")