rebase: update bookkeeping for branches as we rebase descendants

The `DescendantRebaser` keeps a map of branches from the source
commit, so it gets efficient lookup of branches to update when a
commit has been rebased. This map was not kept up to date as we
rebased. That could lead to branches getting left on hidden
intermediate commits. Specifically, if a commit with a branch was
rewritten by some command, and an ancestor of it was also rewritten,
then we'd only update the branch only the first step and not update it
again when rebasing onto the rewritten ancestor.

CHANGELOG.md

@@ -53,6 +53,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   directory with different files in. They are now merged as if the missing
   directory had been empty.
 
+* When using `jj move` to move part of a commit into an ancestor, any branches
+  pointing to the source commit used to be left on a hidden intermediate commit.
+  They are now correctly updated.
+
 * `jj untrack` now requires at least one path (allowing no arguments was a UX
   bug).
 

lib/src/rewrite.rs

@@ -134,7 +134,7 @@ pub struct DescendantRebaser<'settings, 'repo> {
     new_commits: HashSet<CommitId>,
     rebased: HashMap<CommitId, CommitId>,
     // Names of branches where local target includes the commit id in the key.
-    branches: HashMap<CommitId, Vec<String>>,
+    branches: HashMap<CommitId, HashSet<String>>,
     // Parents of rebased/abandoned commit that should become new heads once their descendants
     // have been rebased.
     heads_to_add: HashSet<CommitId>,
@@ -215,20 +215,20 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
 
         // Build a map from commit to branches pointing to it, so we don't need to scan
         // all branches each time we rebase a commit.
-        let mut branches: HashMap<_, Vec<_>> = HashMap::new();
+        let mut branches: HashMap<_, HashSet<_>> = HashMap::new();
         for (branch_name, branch_target) in mut_repo.view().branches() {
             if let Some(local_target) = &branch_target.local_target {
                 for commit in local_target.adds() {
                     branches
                         .entry(commit)
                         .or_default()
-                        .push(branch_name.clone());
+                        .insert(branch_name.clone());
                 }
                 for commit in local_target.removes() {
                     branches
                         .entry(commit)
                         .or_default()
-                        .push(branch_name.clone());
+                        .insert(branch_name.clone());
                 }
             }
         }
@@ -293,9 +293,15 @@ impl<'settings, 'repo> DescendantRebaser<'settings, 'repo> {
     fn update_references(&mut self, old_commit_id: CommitId, new_commit_ids: Vec<CommitId>) {
         self.update_checkouts(&old_commit_id, &new_commit_ids);
 
-        if let Some(branch_names) = self.branches.get(&old_commit_id) {
+        if let Some(branch_names) = self.branches.get(&old_commit_id).cloned() {
             let mut branch_updates = vec![];
-            for branch_name in branch_names {
+            for branch_name in &branch_names {
+                for new_commit_id in &new_commit_ids {
+                    self.branches
+                        .entry(new_commit_id.clone())
+                        .or_default()
+                        .insert(branch_name.clone());
+                }
                 let local_target = self.mut_repo.get_local_branch(branch_name).unwrap();
                 for old_add in local_target.adds() {
                     if old_add == old_commit_id {

lib/tests/test_rewrite.rs

@@ -899,7 +899,7 @@ fn test_rebase_descendants_basic_branch_update() {
     let test_repo = testutils::init_repo(&settings, false);
     let repo = &test_repo.repo;
 
-    // Branch "main" points to branch B. B gets rewritten as B2. Branch main should
+    // Branch "main" points to commit B. B gets rewritten as B2. Branch main should
     // be updated to point to B2.
     //
     // B main         B2 main
@@ -928,13 +928,56 @@ fn test_rebase_descendants_basic_branch_update() {
     );
 }
 
+#[test]
+fn test_rebase_descendants_branch_move_two_steps() {
+    let settings = testutils::user_settings();
+    let test_repo = testutils::init_repo(&settings, false);
+    let repo = &test_repo.repo;
+
+    // Branch "main" points to branch C. C gets rewritten as C2 and B gets rewritten
+    // as B2. C2 should be rebased onto B2, creating C3, and main should be
+    // updated to point to C3.
+    //
+    // C2 C main      C3 main
+    // | /            |
+    // |/        =>   |
+    // B B2           B2
+    // |/             |
+    // 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]);
+    let commit_c = graph_builder.commit_with_parents(&[&commit_b]);
+    tx.mut_repo()
+        .set_local_branch("main".to_string(), RefTarget::Normal(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)
+        .write_to_repo(tx.mut_repo());
+    let commit_c2 = CommitBuilder::for_rewrite_from(&settings, repo.store(), &commit_c)
+        .write_to_repo(tx.mut_repo());
+    tx.mut_repo().rebase_descendants(&settings);
+    let heads = tx.mut_repo().view().heads();
+    assert_eq!(heads.len(), 1);
+    let c3_id = heads.iter().next().unwrap().clone();
+    let commit_c3 = repo.store().get_commit(&c3_id).unwrap();
+    assert_ne!(commit_c3.id(), commit_c2.id());
+    assert_eq!(commit_c3.parent_ids(), vec![commit_b2.id().clone()]);
+    assert_eq!(
+        tx.mut_repo().get_local_branch("main"),
+        Some(RefTarget::Normal(commit_c3.id().clone()))
+    );
+}
+
 #[test]
 fn test_rebase_descendants_basic_branch_update_with_non_local_branch() {
     let settings = testutils::user_settings();
     let test_repo = testutils::init_repo(&settings, false);
     let repo = &test_repo.repo;
 
-    // Branch "main" points to branch B. B gets rewritten as B2. Branch main should
+    // Branch "main" points to commit B. B gets rewritten as B2. Branch main should
     // be updated to point to B2. Remote branch main@origin and tag v1 should not
     // get updated.
     //
@@ -989,7 +1032,7 @@ fn test_rebase_descendants_update_branch_after_abandon() {
     let test_repo = testutils::init_repo(&settings, false);
     let repo = &test_repo.repo;
 
-    // Branch "main" points to branch B. B is then abandoned. Branch main should
+    // Branch "main" points to commit B. B is then abandoned. Branch main should
     // be updated to point to A.
     //
     // B main

tests/test_move_command.rs

@@ -286,9 +286,8 @@ fn test_move_partial() {
         test_env.jj_cmd_success(&repo_path, &["move", "--from", "c", "--to", "b", "file1"]);
     insta::assert_snapshot!(stdout, @"Rebased 1 descendant commits
 ");
-    // TODO: The 'c' branch got lost
     insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
-    o 21253406d416 
+    o 21253406d416 c
     o e1cf08aae711 b
     | @ bdd835cae844 d
     |/