refactor: get hunk locks from HunkDependencies instead of HunkAssignment in absorb

Author: Kiril Videlov

crates/but-api/src/legacy/absorb.rs

@@ -11,7 +11,10 @@ use but_hunk_assignment::{
     AbsorptionReason, AbsorptionTarget, CommitAbsorption, CommitMap, FileAbsorption,
     GroupedChanges, HunkAssignment, convert_assignments_to_diff_specs,
 };
-use but_hunk_dependency::ui::{HunkLock, HunkLockTarget};
+use but_hunk_dependency::ui::{
+    HunkDependencies, HunkLock, HunkLockTarget,
+    hunk_dependencies_for_workspace_changes_by_worktree_dir,
+};
 use but_rebase::graph_rebase::mutate::{InsertSide, RelativeTo};
 use but_workspace::ui::StackDetails;
 use gitbutler_oplog::{
@@ -97,13 +100,14 @@ pub fn absorption_plan(
     ctx: &mut Context,
     target: AbsorptionTarget,
 ) -> anyhow::Result<Vec<CommitAbsorption>> {
-    let assignments = match target {
+    let (assignments, dependencies) = match target {
         AbsorptionTarget::Branch { branch_name } => {
             // Get all worktree changes, assignments, and dependencies
             // TODO: Ideally, there's a simpler way of getting the worktree changes without passing the context to it.
             // At this time, the context is passed pretty deep into the function.
             let worktree_changes = changes_in_worktree(ctx)?;
             let all_assignments = worktree_changes.assignments;
+            let dependencies = worktree_changes.dependencies;
 
             // Get the stack ID for this branch
             let stacks = crate::legacy::workspace::stacks(ctx, None)?;
@@ -131,7 +135,7 @@ pub fn absorption_plan(
                 anyhow::bail!("No uncommitted changes assigned to branch: {branch_name}");
             }
 
-            stack_assignments
+            (stack_assignments, dependencies)
         }
         AbsorptionTarget::TreeChanges {
             changes,
@@ -140,6 +144,7 @@ pub fn absorption_plan(
             // Get all worktree changes, assignments, and dependencies
             let worktree_changes = changes_in_worktree(ctx)?;
             let all_assignments = worktree_changes.assignments;
+            let dependencies = worktree_changes.dependencies;
 
             // Filter assignments to just this stack
             let stack_assignments: Vec<_> = all_assignments
@@ -155,23 +160,34 @@ pub fn absorption_plan(
                 anyhow::bail!("No uncommitted changes assigned to stack: {assigned_stack_id:?}");
             }
 
-            stack_assignments
+            (stack_assignments, dependencies)
+        }
+        AbsorptionTarget::HunkAssignments { assignments } => {
+            // Compute hunk dependencies only for this target since changes_in_worktree isn't called
+            let (_read_guard, repo, ws, _db) = ctx.workspace_and_db()?;
+            let dependencies =
+                hunk_dependencies_for_workspace_changes_by_worktree_dir(&repo, &ws, None).ok();
+            drop((_read_guard, repo, ws, _db));
+            (assignments, dependencies)
         }
-        AbsorptionTarget::HunkAssignments { assignments } => assignments,
         AbsorptionTarget::All => {
             // Get all worktree changes, assignments, and dependencies
             // TODO: Ideally, there's a simpler way of getting the worktree changes without passing the context to it.
             // At this time, the context is passed pretty deep into the function.
             let worktree_changes = changes_in_worktree(ctx)?;
-            worktree_changes.assignments
+            (worktree_changes.assignments, worktree_changes.dependencies)
         }
     };
 
     let mut guard = ctx.exclusive_worktree_access();
 
     // Group all changes by their target commit
-    let changes_by_commit =
-        group_changes_by_target_commit(ctx, &assignments, guard.write_permission())?;
+    let changes_by_commit = group_changes_by_target_commit(
+        ctx,
+        &assignments,
+        dependencies.as_ref(),
+        guard.write_permission(),
+    )?;
 
     // Prepare commit absorptions for display
     let commit_absorptions = prepare_commit_absorptions(ctx, changes_by_commit)?;
@@ -183,17 +199,30 @@ pub fn absorption_plan(
 fn group_changes_by_target_commit(
     ctx: &mut Context,
     assignments: &[HunkAssignment],
+    dependencies: Option<&HunkDependencies>,
     perm: &mut RepoExclusive,
 ) -> anyhow::Result<GroupedChanges> {
     let mut changes_by_commit: GroupedChanges = BTreeMap::new();
 
     let mut stack_details_cache = HashMap::<StackId, StackDetails>::new();
 
+    // Build an index for O(1) lock lookups per assignment
+    let lock_index = dependencies.map(build_lock_index);
+
     // Process each assignment
     for assignment in assignments {
         // Determine the target commit for this assignment
-        let (stack_id, commit_id, reason) =
-            determine_target_commit(ctx, assignment, &mut stack_details_cache, perm)?;
+        let locks = lock_index
+            .as_ref()
+            .map(|idx| locks_for_assignment(idx, assignment))
+            .filter(|l| !l.is_empty());
+        let (stack_id, commit_id, reason) = determine_target_commit(
+            ctx,
+            assignment,
+            locks.as_deref(),
+            &mut stack_details_cache,
+            perm,
+        )?;
 
         let entry = changes_by_commit
             .entry((stack_id, commit_id))
@@ -209,6 +238,71 @@ fn group_changes_by_target_commit(
     Ok(changes_by_commit)
 }
 
+/// Per-file entries of `(DiffHunk, locks)` for range-based lock matching.
+type LockIndex = HashMap<String, Vec<(but_core::unified_diff::DiffHunk, Vec<HunkLock>)>>;
+
+/// Build a lookup index from hunk dependencies, grouped by file path.
+///
+/// Each entry retains the original `DiffHunk` range so that lookups can match
+/// by range overlap rather than exact header equality. This is necessary because
+/// dependency hunks are computed with 0 context lines while assignment hunks use
+/// the user's `context_lines` setting, so their headers differ.
+fn build_lock_index(dependencies: &HunkDependencies) -> LockIndex {
+    let mut index = LockIndex::new();
+    for (path, diff_hunk, locks) in &dependencies.diffs {
+        index
+            .entry(path.clone())
+            .or_default()
+            .push((diff_hunk.clone(), locks.clone()));
+    }
+    index
+}
+
+/// Check whether two line ranges overlap.
+/// Ranges are `[start, start + lines)` (1-based start, length in lines).
+/// A range with 0 lines (pure insertion/deletion) is treated as a point at `start`.
+fn ranges_overlap(start_a: u32, lines_a: u32, start_b: u32, lines_b: u32) -> bool {
+    let end_a = start_a + lines_a.max(1);
+    let end_b = start_b + lines_b.max(1);
+    start_a < end_b && start_b < end_a
+}
+
+/// Look up the dependency locks for an assignment by finding dependency hunks
+/// whose ranges overlap with the assignment's hunk header.
+///
+/// When the assignment has no hunk header (binary/too-large diffs), all locks
+/// for the file are returned as a fallback.
+fn locks_for_assignment(index: &LockIndex, assignment: &HunkAssignment) -> Vec<HunkLock> {
+    let Some(file_entries) = index.get(&assignment.path) else {
+        return Vec::new();
+    };
+
+    match assignment.hunk_header {
+        Some(hunk_header) => {
+            let mut locks = Vec::new();
+            for (dep_hunk, dep_locks) in file_entries {
+                // Match on the new-file side: assignment hunks describe worktree
+                // state (new), and dependency hunks record which committed ranges
+                // they depend on.
+                if ranges_overlap(
+                    dep_hunk.new_start,
+                    dep_hunk.new_lines,
+                    hunk_header.new_start,
+                    hunk_header.new_lines,
+                ) {
+                    locks.extend(dep_locks.iter().cloned());
+                }
+            }
+            locks
+        }
+        // No hunk header (binary/too-large) — return all locks for the file.
+        None => file_entries
+            .iter()
+            .flat_map(|(_, locks)| locks.iter().cloned())
+            .collect(),
+    }
+}
+
 // Find the lock that is highest in the application order (child-most commit)
 fn find_top_most_lock<'a>(
     locks: &'a [HunkLock],
@@ -252,6 +346,7 @@ fn find_top_most_lock<'a>(
 fn determine_target_commit(
     ctx: &mut Context,
     assignment: &HunkAssignment,
+    locks: Option<&[HunkLock]>,
     stack_details_cache: &mut HashMap<StackId, StackDetails>,
     perm: &mut RepoExclusive,
 ) -> anyhow::Result<(
@@ -260,7 +355,7 @@ fn determine_target_commit(
     AbsorptionReason,
 )> {
     // Priority 1: Check if there's a dependency lock for this hunk
-    if let Some(locks) = &assignment.hunk_locks {
+    if let Some(locks) = locks {
         if let Some(lock) = find_top_most_lock(locks, ctx, stack_details_cache) {
             if let HunkLockTarget::Stack(stack_id) = lock.target {
                 return Ok((stack_id, lock.commit_id, AbsorptionReason::HunkDependency));

crates/but/tests/but/command/absorb.rs

@@ -37,7 +37,7 @@ fn uncommitted_file() -> anyhow::Result<()> {
 Found 1 changed file to absorb:
 
 Absorbed to commit: f4ea7f8 a.txt
-  (last commit in the primary lane)
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
     a.txt @6,4 +6,4
 
@@ -116,7 +116,7 @@ k0 a.txt│
 Found 1 changed file to absorb:
 
 Absorbed to commit: f4ea7f8 a.txt
-  (last commit in the primary lane)
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
 
 
@@ -323,9 +323,12 @@ Hint: run `but diff` to see uncommitted changes and `but stage <file>` to stage
         .stdout_eq(snapbox::str![[r#"
 Found 1 changed file to absorb:
 
-Absorbed to commit: a7aa4ef partial change to a.txt 3
-  (last commit in the primary lane)
+Absorbed to commit: 889385c partial change to a.txt 2
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
+
+Absorbed to commit: a7aa4ef partial change to a.txt 3
+  (files locked to commit due to hunk range overlap)
     a.txt @6,4 +6,4
 
 
@@ -343,10 +346,10 @@ Hint: you can run `but undo` to undo these changes
 ┊     no changes
 ┊
 ┊╭┄g0 [A]
-┊●   a309adb partial change to a.txt 3
-┊│     a3:nk M a.txt
-┊●   889385c partial change to a.txt 2
-┊│     88:nk M a.txt
+┊●   4822140 partial change to a.txt 3
+┊│     48:nk M a.txt
+┊●   4593422 partial change to a.txt 2
+┊│     45:nk M a.txt
 ┊●   8dc39e0 partial change to a.txt 1
 ┊│     8d:nk M a.txt
 ┊●   f4ea7f8 a.txt
@@ -415,7 +418,7 @@ fn uncommitted_file_new() -> anyhow::Result<()> {
 Found 1 changed file to absorb:
 
 Created on top of commit: f4ea7f8 a.txt
-  (last commit in the primary lane)
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
     a.txt @6,4 +6,4
 
@@ -533,7 +536,7 @@ k0 a.txt│
 Found 1 changed file to absorb:
 
 Created on top of commit: f4ea7f8 a.txt
-  (last commit in the primary lane)
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
 
 
@@ -636,7 +639,7 @@ fn dry_run_new_shows_plan_without_changes() -> anyhow::Result<()> {
 Found 1 changed file to absorb:
 
 Created on top of commit: f4ea7f8 a.txt
-  (last commit in the primary lane)
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
     a.txt @6,4 +6,4
 
@@ -707,7 +710,7 @@ fn dry_run_shows_plan_without_changes() -> anyhow::Result<()> {
 Found 1 changed file to absorb:
 
 Absorbed to commit: f4ea7f8 a.txt
-  (last commit in the primary lane)
+  (files locked to commit due to hunk range overlap)
     a.txt @1,4 +1,4
     a.txt @6,4 +6,4