From e0c2d82127f5ade1651e519532d91bea50bcf97f Mon Sep 17 00:00:00 2001
From: Mattias Granlund <mtsgrd@gmail.com>
Date: Thu, 25 Apr 2024 16:08:38 +0200
Subject: [PATCH] Reset files into the same branch they were committted to

- regressed when switching to blame for hunk locking
---
 .../src/virtual_branches/virtual.rs           | 43 +++++++++++-
 .../suite/virtual_branches/create_commit.rs   | 70 +++++++++++++++++++
 2 files changed, 112 insertions(+), 1 deletion(-)

diff --git a/crates/gitbutler-core/src/virtual_branches/virtual.rs b/crates/gitbutler-core/src/virtual_branches/virtual.rs
index d9a563b9b..0caff1d5f 100644
--- a/crates/gitbutler-core/src/virtual_branches/virtual.rs
+++ b/crates/gitbutler-core/src/virtual_branches/virtual.rs
@@ -16,13 +16,14 @@ use git2_hooks::HookResult;
 use regex::Regex;
 use serde::Serialize;
 
+use super::integration::get_workspace_head;
 use super::{
     branch::{
         self, Branch, BranchCreateRequest, BranchId, BranchOwnershipClaims, Hunk, OwnershipClaim,
     },
     branch_to_remote_branch, errors, target, RemoteBranch, VirtualBranchesHandle,
 };
-use crate::git::diff::{diff_files_into_hunks, DiffByPathMap};
+use crate::git::diff::{diff_files_into_hunks, trees, DiffByPathMap};
 use crate::virtual_branches::branch::HunkHash;
 use crate::{
     askpass::AskpassBroker,
@@ -2042,11 +2043,51 @@ pub fn reset_branch(
         ));
     }
 
+    // Compute the old workspace before resetting so we can can figure out
+    // what hunks were released by this reset, and assign them to this branch.
+    let old_head = get_workspace_head(&vb_state, project_repository)?;
+
     branch.head = target_commit_oid;
     vb_state
         .set_branch(branch.clone())
         .context("failed to write branch")?;
 
+    let updated_head = get_workspace_head(&vb_state, project_repository)?;
+    let repo = &project_repository.git_repository;
+    let diff = trees(
+        repo,
+        &repo
+            .find_commit(updated_head)
+            .map_err(anyhow::Error::from)?
+            .tree()
+            .map_err(anyhow::Error::from)?,
+        &repo
+            .find_commit(old_head)
+            .map_err(anyhow::Error::from)?
+            .tree()
+            .map_err(anyhow::Error::from)?,
+    )?;
+
+    // Assign the new hunks to the branch we're working on.
+    for (path, filediff) in diff {
+        for hunk in filediff.hunks {
+            let hash = Hunk::hash_diff(hunk.diff_lines.as_ref());
+            branch.ownership.put(
+                &format!(
+                    "{}:{}-{}-{:?}",
+                    path.display(),
+                    hunk.new_start,
+                    hunk.new_start + hunk.new_lines,
+                    &hash
+                )
+                .parse()?,
+            );
+        }
+    }
+    vb_state
+        .set_branch(branch)
+        .context("failed to write branch")?;
+
     super::integration::update_gitbutler_integration(&vb_state, project_repository)
         .context("failed to update gitbutler integration")?;
 
diff --git a/crates/gitbutler-core/tests/suite/virtual_branches/create_commit.rs b/crates/gitbutler-core/tests/suite/virtual_branches/create_commit.rs
index e3f71f85d..2bdc46b70 100644
--- a/crates/gitbutler-core/tests/suite/virtual_branches/create_commit.rs
+++ b/crates/gitbutler-core/tests/suite/virtual_branches/create_commit.rs
@@ -159,6 +159,76 @@ async fn should_not_lock_disjointed_hunks() {
     }
 }
 
+#[tokio::test]
+async fn should_reset_into_same_branch() {
+    let Test {
+        project_id,
+        controller,
+        repository,
+        ..
+    } = &Test::default();
+
+    let mut lines = gen_file(repository, "file.txt", 7);
+    commit_and_push_initial(repository);
+
+    let base_branch = controller
+        .set_base_branch(project_id, &"refs/remotes/origin/master".parse().unwrap())
+        .await
+        .unwrap();
+
+    controller
+        .create_virtual_branch(project_id, &branch::BranchCreateRequest::default())
+        .await
+        .unwrap();
+
+    let branch_2_id = controller
+        .create_virtual_branch(
+            project_id,
+            &branch::BranchCreateRequest {
+                selected_for_changes: Some(true),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+
+    lines[0] = "change 1".to_string();
+    write_file(repository, "file.txt", &lines);
+
+    controller
+        .create_commit(project_id, &branch_2_id, "commit to branch 2", None, false)
+        .await
+        .unwrap();
+
+    let files = get_virtual_branch(controller, project_id, branch_2_id)
+        .await
+        .files;
+    assert_eq!(files.len(), 0);
+
+    // Set target to branch 1 and verify the file resets into branch 2.
+    controller
+        .update_virtual_branch(
+            project_id,
+            branch::BranchUpdateRequest {
+                id: branch_2_id,
+                selected_for_changes: Some(true),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+
+    controller
+        .reset_virtual_branch(project_id, &branch_2_id, base_branch.base_sha)
+        .await
+        .unwrap();
+
+    let files = get_virtual_branch(controller, project_id, branch_2_id)
+        .await
+        .files;
+    assert_eq!(files.len(), 1);
+}
+
 #[tokio::test]
 async fn should_double_lock() {
     let Test {