Undo commit: Files stay in the branch that owned the undone commit

Undoing a commit would move the files to the default branch, instead of keeping them in the source branch.
Now, the files stay in the branch the commit was undone from.

crates/gitbutler-branch-actions/src/undo_commit.rs

@@ -1,8 +1,11 @@
+use std::collections::HashMap;
+
 use anyhow::{bail, Context as _, Result};
 use gitbutler_command_context::CommandContext;
 use gitbutler_commit::commit_ext::CommitExt as _;
+use gitbutler_diff::Hunk;
 use gitbutler_repo::{rebase::cherry_rebase_group, LogUntil, RepositoryExt as _};
-use gitbutler_stack::{Stack, StackId};
+use gitbutler_stack::{OwnershipClaim, Stack, StackId};
 
 use crate::VirtualBranchesExt as _;
 
@@ -26,7 +29,14 @@ pub(crate) fn undo_commit(
 
     let mut branch = vb_state.get_branch_in_workspace(branch_id)?;
 
-    let new_head_commit = inner_undo_commit(ctx.repository(), branch.head(), commit_oid)?;
+    let UndoResult {
+        new_head: new_head_commit,
+        ownership_update,
+    } = inner_undo_commit(ctx.repository(), branch.head(), commit_oid)?;
+
+    for ownership in ownership_update {
+        branch.ownership.put(ownership);
+    }
 
     branch.set_stack_head(ctx, new_head_commit, None)?;
 
@@ -39,20 +49,54 @@ pub(crate) fn undo_commit(
     Ok(branch)
 }
 
+struct UndoResult {
+    new_head: git2::Oid,
+    ownership_update: Vec<OwnershipClaim>,
+}
+
 fn inner_undo_commit(
     repository: &git2::Repository,
     branch_head_commit: git2::Oid,
     commit_to_remove: git2::Oid,
-) -> Result<git2::Oid> {
+) -> Result<UndoResult> {
     let commit_to_remove = repository.find_commit(commit_to_remove)?;
 
     if commit_to_remove.is_conflicted() {
         bail!("Can not undo a conflicted commit");
     }
+    let commit_tree = commit_to_remove
+        .tree()
+        .context("failed to get commit tree")?;
+    let commit_to_remove_parent = commit_to_remove.parent(0)?;
+    let commit_parent_tree = commit_to_remove_parent
+        .tree()
+        .context("failed to get parent tree")?;
+
+    let diff = gitbutler_diff::trees(repository, &commit_parent_tree, &commit_tree, true)?;
+    let diff: HashMap<_, _> = gitbutler_diff::diff_files_into_hunks(diff).collect();
+    let ownership_update = diff
+        .iter()
+        .filter_map(|(file_path, hunks)| {
+            let file_path = file_path.clone();
+            let hunks = hunks
+                .iter()
+                .map(Into::into)
+                .filter(|hunk: &Hunk| hunk.start != 0 && hunk.end != 0)
+                .collect::<Vec<_>>();
+            if hunks.is_empty() {
+                return None;
+            }
+            Some((file_path, hunks))
+        })
+        .map(|(file_path, hunks)| OwnershipClaim { file_path, hunks })
+        .collect::<Vec<_>>();
 
     // if commit is the head, just set head to the parent
     if branch_head_commit == commit_to_remove.id() {
-        return Ok(commit_to_remove.parent(0)?.id());
+        return Ok(UndoResult {
+            new_head: commit_to_remove_parent.id(),
+            ownership_update,
+        });
     };
 
     let commits_to_rebase = repository.l(
@@ -67,20 +111,25 @@ fn inner_undo_commit(
         &commits_to_rebase,
     )?;
 
-    Ok(new_head)
+    Ok(UndoResult {
+        new_head,
+        ownership_update,
+    })
 }
 
 #[cfg(test)]
 mod test {
     #[cfg(test)]
     mod inner_undo_commit {
+        use std::path::PathBuf;
+
         use gitbutler_commit::commit_ext::CommitExt as _;
         use gitbutler_repo::rebase::gitbutler_merge_commits;
         use gitbutler_testsupport::testing_repository::{
             assert_commit_tree_matches, TestingRepository,
         };
 
-        use crate::undo_commit::inner_undo_commit;
+        use crate::undo_commit::{inner_undo_commit, UndoResult};
 
         #[test]
         fn undoing_conflicted_commit_errors() {
@@ -115,9 +164,27 @@ mod test {
             let b = test_repository.commit_tree(Some(&a), &[("bar.txt", "bar")]);
             let c = test_repository.commit_tree(Some(&b), &[("baz.txt", "baz")]);
 
-            let new_head = inner_undo_commit(&test_repository.repository, c.id(), c.id()).unwrap();
+            let UndoResult {
+                new_head,
+                ownership_update,
+            } = inner_undo_commit(&test_repository.repository, c.id(), c.id()).unwrap();
 
             assert_eq!(new_head, b.id(), "The new head should be C's parent");
+            assert_eq!(
+                ownership_update.len(),
+                1,
+                "Should have one ownership update"
+            );
+            assert_eq!(
+                ownership_update[0].file_path,
+                PathBuf::from("baz.txt"),
+                "Ownership update should be for baz.txt"
+            );
+            assert_eq!(
+                ownership_update[0].hunks.len(),
+                1,
+                "Ownership update should have one hunk"
+            );
         }
 
         #[test]
@@ -136,7 +203,10 @@ mod test {
             //
             // As the theirs and ours both are different to the base, it ends up
             // conflicted.
-            let new_head = inner_undo_commit(&test_repository.repository, c.id(), b.id()).unwrap();
+            let UndoResult {
+                new_head,
+                ownership_update,
+            } = inner_undo_commit(&test_repository.repository, c.id(), b.id()).unwrap();
 
             let new_head_commit: git2::Commit =
                 test_repository.repository.find_commit(new_head).unwrap();
@@ -159,6 +229,21 @@ mod test {
                 a.id(),
                 "A should be C prime's parent"
             );
+            assert_eq!(
+                ownership_update.len(),
+                1,
+                "Should have one ownership update"
+            );
+            assert_eq!(
+                ownership_update[0].file_path,
+                PathBuf::from("foo.txt"),
+                "Ownership update should be for foo.txt"
+            );
+            assert_eq!(
+                ownership_update[0].hunks.len(),
+                1,
+                "Ownership update should have one hunk"
+            );
         }
     }
 }

crates/gitbutler-branch-actions/tests/virtual_branches/undo_commit.rs

@@ -62,3 +62,78 @@ fn undo_commit_simple() {
 
     assert_eq!(descriptions, vec!["commit three", "commit one"]);
 }
+
+#[test]
+fn undo_commit_in_non_default_branch() {
+    let Test {
+        repository,
+        project,
+        ..
+    } = &Test::default();
+
+    gitbutler_branch_actions::set_base_branch(
+        project,
+        &"refs/remotes/origin/master".parse().unwrap(),
+    )
+    .unwrap();
+
+    let branch_id =
+        gitbutler_branch_actions::create_virtual_branch(project, &BranchCreateRequest::default())
+            .unwrap();
+
+    // create commit
+    fs::write(repository.path().join("file.txt"), "content").unwrap();
+    let _commit1_id =
+        gitbutler_branch_actions::create_commit(project, branch_id, "commit one", None, false)
+            .unwrap();
+
+    // create commit
+    fs::write(repository.path().join("file2.txt"), "content2").unwrap();
+    fs::write(repository.path().join("file3.txt"), "content3").unwrap();
+    let commit2_id =
+        gitbutler_branch_actions::create_commit(project, branch_id, "commit two", None, false)
+            .unwrap();
+
+    // create commit
+    fs::write(repository.path().join("file4.txt"), "content4").unwrap();
+    let _commit3_id =
+        gitbutler_branch_actions::create_commit(project, branch_id, "commit three", None, false)
+            .unwrap();
+
+    // create default branch
+    // this branch should not be affected by the undo
+    let default_branch_id = gitbutler_branch_actions::create_virtual_branch(
+        project,
+        &BranchCreateRequest {
+            selected_for_changes: Some(true),
+            ..BranchCreateRequest::default()
+        },
+    )
+    .unwrap();
+
+    gitbutler_branch_actions::undo_commit(project, branch_id, commit2_id).unwrap();
+
+    let mut branches = gitbutler_branch_actions::list_virtual_branches(project)
+        .unwrap()
+        .0
+        .into_iter();
+
+    let branch = &branches.find(|b| b.id == branch_id).unwrap();
+    let default_branch = &branches.find(|b| b.id == default_branch_id).unwrap();
+
+    // should be two uncommitted files now (file2.txt and file3.txt)
+    assert_eq!(branch.files.len(), 2);
+    assert_eq!(branch.commits.len(), 2);
+    assert_eq!(branch.commits[0].files.len(), 1);
+    assert_eq!(branch.commits[1].files.len(), 1);
+    assert_eq!(default_branch.files.len(), 0);
+    assert_eq!(default_branch.commits.len(), 0);
+
+    let descriptions = branch
+        .commits
+        .iter()
+        .map(|c| c.description.clone())
+        .collect::<Vec<_>>();
+
+    assert_eq!(descriptions, vec!["commit three", "commit one"]);
+}