Fix crash when staging a hunk that overlaps multiple unstaged hunks (#27545)

João Marcos and Max Brunsfeld created

Release Notes:

- Git: Fix crash when staging a hunk that overlaps multiple unstaged
hunks.

---------

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>

Change summary

crates/buffer_diff/src/buffer_diff.rs | 164 +++++++++++++++++++++-------
1 file changed, 123 insertions(+), 41 deletions(-)

Detailed changes

crates/buffer_diff/src/buffer_diff.rs 🔗

@@ -3,9 +3,7 @@ use git2::{DiffLineType as GitDiffLineType, DiffOptions as GitOptions, Patch as
 use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Task};
 use language::{Language, LanguageRegistry};
 use rope::Rope;
-use std::cmp::Ordering;
-use std::mem;
-use std::{future::Future, iter, ops::Range, sync::Arc};
+use std::{cmp::Ordering, future::Future, iter, mem, ops::Range, sync::Arc};
 use sum_tree::SumTree;
 use text::{Anchor, Bias, BufferId, OffsetRangeExt, Point, ToOffset as _};
 use util::ResultExt;
@@ -189,7 +187,7 @@ impl BufferDiffSnapshot {
 impl BufferDiffInner {
     /// Returns the new index text and new pending hunks.
     fn stage_or_unstage_hunks_impl(
-        &mut self,
+        &self,
         unstaged_diff: &Self,
         stage: bool,
         hunks: &[DiffHunk],
@@ -261,7 +259,6 @@ impl BufferDiffInner {
                 old_pending_hunks.next(buffer);
             }
 
-            // merge into pending hunks
             if (stage && secondary_status == DiffHunkSecondaryStatus::NoSecondaryHunk)
                 || (!stage && secondary_status == DiffHunkSecondaryStatus::HasSecondaryHunk)
             {
@@ -288,56 +285,71 @@ impl BufferDiffInner {
         let mut unstaged_hunk_cursor = unstaged_diff.hunks.cursor::<DiffHunkSummary>(buffer);
         unstaged_hunk_cursor.next(buffer);
 
-        let mut prev_unstaged_hunk_buffer_offset = 0;
-        let mut prev_unstaged_hunk_base_text_offset = 0;
-        let mut edits = Vec::<(Range<usize>, String)>::new();
-
         // then, iterate over all pending hunks (both new ones and the existing ones) and compute the edits
-        for PendingHunk {
+        let mut prev_unstaged_hunk_buffer_end = 0;
+        let mut prev_unstaged_hunk_base_text_end = 0;
+        let mut edits = Vec::<(Range<usize>, String)>::new();
+        let mut pending_hunks_iter = pending_hunks.iter().cloned().peekable();
+        while let Some(PendingHunk {
             buffer_range,
             diff_base_byte_range,
             ..
-        } in pending_hunks.iter().cloned()
+        }) = pending_hunks_iter.next()
         {
-            let skipped_hunks = unstaged_hunk_cursor.slice(&buffer_range.start, Bias::Left, buffer);
+            // Advance unstaged_hunk_cursor to skip unstaged hunks before current hunk
+            let skipped_unstaged =
+                unstaged_hunk_cursor.slice(&buffer_range.start, Bias::Left, buffer);
 
-            if let Some(secondary_hunk) = skipped_hunks.last() {
-                prev_unstaged_hunk_base_text_offset = secondary_hunk.diff_base_byte_range.end;
-                prev_unstaged_hunk_buffer_offset =
-                    secondary_hunk.buffer_range.end.to_offset(buffer);
+            if let Some(unstaged_hunk) = skipped_unstaged.last() {
+                prev_unstaged_hunk_base_text_end = unstaged_hunk.diff_base_byte_range.end;
+                prev_unstaged_hunk_buffer_end = unstaged_hunk.buffer_range.end.to_offset(buffer);
             }
 
+            // Find where this hunk is in the index if it doesn't overlap
             let mut buffer_offset_range = buffer_range.to_offset(buffer);
-            let start_overshoot = buffer_offset_range.start - prev_unstaged_hunk_buffer_offset;
-            let mut index_start = prev_unstaged_hunk_base_text_offset + start_overshoot;
-
-            while let Some(unstaged_hunk) = unstaged_hunk_cursor.item().filter(|item| {
-                item.buffer_range
-                    .start
-                    .cmp(&buffer_range.end, buffer)
-                    .is_le()
-            }) {
-                let unstaged_hunk_offset_range = unstaged_hunk.buffer_range.to_offset(buffer);
-                prev_unstaged_hunk_base_text_offset = unstaged_hunk.diff_base_byte_range.end;
-                prev_unstaged_hunk_buffer_offset = unstaged_hunk_offset_range.end;
+            let start_overshoot = buffer_offset_range.start - prev_unstaged_hunk_buffer_end;
+            let mut index_start = prev_unstaged_hunk_base_text_end + start_overshoot;
+
+            loop {
+                // Merge this hunk with any overlapping unstaged hunks.
+                if let Some(unstaged_hunk) = unstaged_hunk_cursor.item() {
+                    let unstaged_hunk_offset_range = unstaged_hunk.buffer_range.to_offset(buffer);
+                    if unstaged_hunk_offset_range.start <= buffer_offset_range.end {
+                        prev_unstaged_hunk_base_text_end = unstaged_hunk.diff_base_byte_range.end;
+                        prev_unstaged_hunk_buffer_end = unstaged_hunk_offset_range.end;
+
+                        index_start = index_start.min(unstaged_hunk.diff_base_byte_range.start);
+                        buffer_offset_range.start = buffer_offset_range
+                            .start
+                            .min(unstaged_hunk_offset_range.start);
+                        buffer_offset_range.end =
+                            buffer_offset_range.end.max(unstaged_hunk_offset_range.end);
+
+                        unstaged_hunk_cursor.next(buffer);
+                        continue;
+                    }
+                }
 
-                index_start = index_start.min(unstaged_hunk.diff_base_byte_range.start);
-                buffer_offset_range.start = buffer_offset_range
-                    .start
-                    .min(unstaged_hunk_offset_range.start);
+                // If any unstaged hunks were merged, then subsequent pending hunks may
+                // now overlap this hunk. Merge them.
+                if let Some(next_pending_hunk) = pending_hunks_iter.peek() {
+                    let next_pending_hunk_offset_range =
+                        next_pending_hunk.buffer_range.to_offset(buffer);
+                    if next_pending_hunk_offset_range.start <= buffer_offset_range.end {
+                        buffer_offset_range.end = next_pending_hunk_offset_range.end;
+                        pending_hunks_iter.next();
+                        continue;
+                    }
+                }
 
-                unstaged_hunk_cursor.next(buffer);
+                break;
             }
 
             let end_overshoot = buffer_offset_range
                 .end
-                .saturating_sub(prev_unstaged_hunk_buffer_offset);
-            let index_end = prev_unstaged_hunk_base_text_offset + end_overshoot;
-
-            let index_range = index_start..index_end;
-            buffer_offset_range.end = buffer_offset_range
-                .end
-                .max(prev_unstaged_hunk_buffer_offset);
+                .saturating_sub(prev_unstaged_hunk_buffer_end);
+            let index_end = prev_unstaged_hunk_base_text_end + end_overshoot;
+            let index_byte_range = index_start..index_end;
 
             let replacement_text = if stage {
                 log::debug!("stage hunk {:?}", buffer_offset_range);
@@ -351,8 +363,9 @@ impl BufferDiffInner {
                     .collect::<String>()
             };
 
-            edits.push((index_range, replacement_text));
+            edits.push((index_byte_range, replacement_text));
         }
+        drop(pending_hunks_iter);
 
         #[cfg(debug_assertions)] // invariants: non-overlapping and sorted
         {
@@ -1649,6 +1662,75 @@ mod tests {
                 "
                 .unindent(),
             },
+            Example {
+                name: "one unstaged hunk that contains two uncommitted hunks",
+                head_text: "
+                    one
+                    two
+
+                    three
+                    four
+                "
+                .unindent(),
+                index_text: "
+                    one
+                    two
+                    three
+                    four
+                "
+                .unindent(),
+                buffer_marked_text: "
+                    «one
+
+                    three // modified
+                    four»
+                "
+                .unindent(),
+                final_index_text: "
+                    one
+
+                    three // modified
+                    four
+                "
+                .unindent(),
+            },
+            Example {
+                name: "one uncommitted hunk that contains two unstaged hunks",
+                head_text: "
+                    one
+                    two
+                    three
+                    four
+                    five
+                "
+                .unindent(),
+                index_text: "
+                    ZERO
+                    one
+                    TWO
+                    THREE
+                    FOUR
+                    five
+                "
+                .unindent(),
+                buffer_marked_text: "
+                    «one
+                    TWO_HUNDRED
+                    THREE
+                    FOUR_HUNDRED
+                    five»
+                "
+                .unindent(),
+                final_index_text: "
+                    ZERO
+                    one
+                    TWO_HUNDRED
+                    THREE
+                    FOUR_HUNDRED
+                    five
+                "
+                .unindent(),
+            },
         ];
 
         for example in table {