Implement staging and unstaging hunks (#24606)

- [x] Staging hunks
- [x] Unstaging hunks
- [x] Write a randomized test
- [x] Get test passing
- [x] Fix existing bug in diff_base_byte_range computation
- [x] Remote project support
- [ ] ~~Improve performance of
buffer_range_to_unchanged_diff_base_range~~
- [ ] ~~Bug: project diff editor scrolls to top when staging/unstaging
hunk~~ existing issue
- [ ] ~~UI~~ deferred
- [x] Tricky cases
  - [x] Correctly handle acting on multiple hunks for a single file
- [x] Remove path from index when unstaging the last staged hunk, if
it's absent from HEAD, or staging the only hunk, if it's deleted in the
working copy

Release Notes:

- Add `ToggleStagedSelectedDiffHunks` action for staging and unstaging
individual diff hunks

Cargo.lock

@@ -2029,11 +2029,15 @@ name = "buffer_diff"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "ctor",
+ "env_logger 0.11.6",
  "futures 0.3.31",
  "git2",
  "gpui",
  "language",
+ "log",
  "pretty_assertions",
+ "rand 0.8.5",
  "rope",
  "serde_json",
  "sum_tree",

crates/buffer_diff/Cargo.toml

@@ -20,14 +20,18 @@ futures.workspace = true
 git2.workspace = true
 gpui.workspace = true
 language.workspace = true
+log.workspace = true
 rope.workspace = true
 sum_tree.workspace = true
 text.workspace = true
 util.workspace = true
 
 [dev-dependencies]
+ctor.workspace = true
+env_logger.workspace = true
+gpui = { workspace = true, features = ["test-support"] }
 pretty_assertions.workspace = true
+rand.workspace = true
 serde_json.workspace = true
 text = { workspace = true, features = ["test-support"] }
-gpui = { workspace = true, features = ["test-support"] }
 unindent.workspace = true

crates/buffer_diff/src/buffer_diff.rs

@@ -5,6 +5,7 @@ use language::{Language, LanguageRegistry};
 use rope::Rope;
 use std::{cmp, future::Future, iter, ops::Range, sync::Arc};
 use sum_tree::SumTree;
+use text::ToOffset as _;
 use text::{Anchor, Bias, BufferId, OffsetRangeExt, Point};
 use util::ResultExt;
 
@@ -14,10 +15,11 @@ pub struct BufferDiff {
     secondary_diff: Option<Entity<BufferDiff>>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct BufferDiffSnapshot {
     inner: BufferDiffInner,
     secondary_diff: Option<Box<BufferDiffSnapshot>>,
+    pub is_single_insertion: bool,
 }
 
 #[derive(Clone)]
@@ -40,21 +42,6 @@ pub enum DiffHunkSecondaryStatus {
     None,
 }
 
-// to stage a hunk:
-// - assume hunk starts out as not staged
-// - hunk exists with the same buffer range in the unstaged diff and the uncommitted diff
-// - we want to construct a "version" of the file that
-//   - starts from the index base text
-//   - has the single hunk applied to it
-//     - the hunk is the one from the UNSTAGED diff, so that the diff base offset range is correct to apply to that diff base
-// - write that new version of the file into the index
-
-// to unstage a hunk
-// - no hunk in the unstaged diff intersects this hunk from the uncommitted diff
-// - we want to compute the hunk that
-//   - we can apply to the index text
-//   - at the end of applying it,
-
 /// A diff hunk resolved to rows in the buffer.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct DiffHunk {
@@ -65,6 +52,7 @@ pub struct DiffHunk {
     /// The range in the buffer's diff base text to which this hunk corresponds.
     pub diff_base_byte_range: Range<usize>,
     pub secondary_status: DiffHunkSecondaryStatus,
+    pub secondary_diff_base_byte_range: Option<Range<usize>>,
 }
 
 /// We store [`InternalDiffHunk`]s internally so we don't need to store the additional row range.
@@ -166,6 +154,99 @@ impl BufferDiffSnapshot {
             }
         }
     }
+
+    fn buffer_range_to_unchanged_diff_base_range(
+        &self,
+        buffer_range: Range<Anchor>,
+        buffer: &text::BufferSnapshot,
+    ) -> Option<Range<usize>> {
+        let mut hunks = self.inner.hunks.iter();
+        let mut start = 0;
+        let mut pos = buffer.anchor_before(0);
+        while let Some(hunk) = hunks.next() {
+            assert!(buffer_range.start.cmp(&pos, buffer).is_ge());
+            assert!(hunk.buffer_range.start.cmp(&pos, buffer).is_ge());
+            if hunk
+                .buffer_range
+                .start
+                .cmp(&buffer_range.end, buffer)
+                .is_ge()
+            {
+                // target buffer range is contained in the unchanged stretch leading up to this next hunk,
+                // so do a final adjustment based on that
+                break;
+            }
+
+            // if the target buffer range intersects this hunk at all, no dice
+            if buffer_range
+                .start
+                .cmp(&hunk.buffer_range.end, buffer)
+                .is_lt()
+            {
+                return None;
+            }
+
+            start += hunk.buffer_range.start.to_offset(buffer) - pos.to_offset(buffer);
+            start += hunk.diff_base_byte_range.end - hunk.diff_base_byte_range.start;
+            pos = hunk.buffer_range.end;
+        }
+        start += buffer_range.start.to_offset(buffer) - pos.to_offset(buffer);
+        let end = start + buffer_range.end.to_offset(buffer) - buffer_range.start.to_offset(buffer);
+        Some(start..end)
+    }
+
+    pub fn secondary_edits_for_stage_or_unstage(
+        &self,
+        stage: bool,
+        hunks: impl Iterator<Item = (Range<usize>, Option<Range<usize>>, Range<Anchor>)>,
+        buffer: &text::BufferSnapshot,
+    ) -> Vec<(Range<usize>, String)> {
+        let Some(secondary_diff) = self.secondary_diff() else {
+            log::debug!("no secondary diff");
+            return Vec::new();
+        };
+        let index_base = secondary_diff.base_text().map_or_else(
+            || Rope::from(""),
+            |snapshot| snapshot.text.as_rope().clone(),
+        );
+        let head_base = self.base_text().map_or_else(
+            || Rope::from(""),
+            |snapshot| snapshot.text.as_rope().clone(),
+        );
+        log::debug!("original: {:?}", index_base.to_string());
+        let mut edits = Vec::new();
+        for (diff_base_byte_range, secondary_diff_base_byte_range, buffer_range) in hunks {
+            let (index_byte_range, replacement_text) = if stage {
+                log::debug!("staging");
+                let mut replacement_text = String::new();
+                let Some(index_byte_range) = secondary_diff_base_byte_range.clone() else {
+                    log::debug!("not a stageable hunk");
+                    continue;
+                };
+                log::debug!("using {:?}", index_byte_range);
+                for chunk in buffer.text_for_range(buffer_range.clone()) {
+                    replacement_text.push_str(chunk);
+                }
+                (index_byte_range, replacement_text)
+            } else {
+                log::debug!("unstaging");
+                let mut replacement_text = String::new();
+                let Some(index_byte_range) = secondary_diff
+                    .buffer_range_to_unchanged_diff_base_range(buffer_range.clone(), &buffer)
+                else {
+                    log::debug!("not an unstageable hunk");
+                    continue;
+                };
+                for chunk in head_base.chunks_in_range(diff_base_byte_range.clone()) {
+                    replacement_text.push_str(chunk);
+                }
+                (index_byte_range, replacement_text)
+            };
+            edits.push((index_byte_range, replacement_text));
+        }
+        log::debug!("edits: {edits:?}");
+        edits
+    }
 }
 
 impl BufferDiffInner {
@@ -225,6 +306,7 @@ impl BufferDiffInner {
             }
 
             let mut secondary_status = DiffHunkSecondaryStatus::None;
+            let mut secondary_diff_base_byte_range = None;
             if let Some(secondary_cursor) = secondary_cursor.as_mut() {
                 if start_anchor
                     .cmp(&secondary_cursor.start().buffer_range.start, buffer)
@@ -234,9 +316,15 @@ impl BufferDiffInner {
                 }
 
                 if let Some(secondary_hunk) = secondary_cursor.item() {
-                    let secondary_range = secondary_hunk.buffer_range.to_point(buffer);
+                    let mut secondary_range = secondary_hunk.buffer_range.to_point(buffer);
+                    if secondary_range.end.column > 0 {
+                        secondary_range.end.row += 1;
+                        secondary_range.end.column = 0;
+                    }
                     if secondary_range == (start_point..end_point) {
                         secondary_status = DiffHunkSecondaryStatus::HasSecondaryHunk;
+                        secondary_diff_base_byte_range =
+                            Some(secondary_hunk.diff_base_byte_range.clone());
                     } else if secondary_range.start <= end_point {
                         secondary_status = DiffHunkSecondaryStatus::OverlapsWithSecondaryHunk;
                     }
@@ -248,6 +336,7 @@ impl BufferDiffInner {
                 diff_base_byte_range: start_base..end_base,
                 buffer_range: start_anchor..end_anchor,
                 secondary_status,
+                secondary_diff_base_byte_range,
             });
         })
     }
@@ -282,6 +371,7 @@ impl BufferDiffInner {
                 buffer_range: hunk.buffer_range.clone(),
                 // The secondary status is not used by callers of this method.
                 secondary_status: DiffHunkSecondaryStatus::None,
+                secondary_diff_base_byte_range: None,
             })
         })
     }
@@ -351,12 +441,12 @@ impl BufferDiffInner {
 }
 
 fn compute_hunks(
-    diff_base: Option<Arc<String>>,
+    diff_base: Option<(Arc<String>, Rope)>,
     buffer: text::BufferSnapshot,
 ) -> SumTree<InternalDiffHunk> {
     let mut tree = SumTree::new(&buffer);
 
-    if let Some(diff_base) = diff_base {
+    if let Some((diff_base, diff_base_rope)) = diff_base {
         let buffer_text = buffer.as_rope().to_string();
 
         let mut options = GitOptions::default();
@@ -387,7 +477,13 @@ fn compute_hunks(
         if let Some(patch) = patch {
             let mut divergence = 0;
             for hunk_index in 0..patch.num_hunks() {
-                let hunk = process_patch_hunk(&patch, hunk_index, &buffer, &mut divergence);
+                let hunk = process_patch_hunk(
+                    &patch,
+                    hunk_index,
+                    &diff_base_rope,
+                    &buffer,
+                    &mut divergence,
+                );
                 tree.push(hunk, &buffer);
             }
         }
@@ -399,6 +495,7 @@ fn compute_hunks(
 fn process_patch_hunk(
     patch: &GitPatch<'_>,
     hunk_index: usize,
+    diff_base: &Rope,
     buffer: &text::BufferSnapshot,
     buffer_row_divergence: &mut i64,
 ) -> InternalDiffHunk {
@@ -408,50 +505,59 @@ fn process_patch_hunk(
     let mut first_deletion_buffer_row: Option<u32> = None;
     let mut buffer_row_range: Option<Range<u32>> = None;
     let mut diff_base_byte_range: Option<Range<usize>> = None;
+    let mut first_addition_old_row: Option<u32> = None;
 
     for line_index in 0..line_item_count {
         let line = patch.line_in_hunk(hunk_index, line_index).unwrap();
         let kind = line.origin_value();
         let content_offset = line.content_offset() as isize;
         let content_len = line.content().len() as isize;
+        match kind {
+            GitDiffLineType::Addition => {
+                if first_addition_old_row.is_none() {
+                    first_addition_old_row = Some(
+                        (line.new_lineno().unwrap() as i64 - *buffer_row_divergence - 1) as u32,
+                    );
+                }
+                *buffer_row_divergence += 1;
+                let row = line.new_lineno().unwrap().saturating_sub(1);
 
-        if kind == GitDiffLineType::Addition {
+                match &mut buffer_row_range {
-            *buffer_row_divergence += 1;
+                    Some(Range { end, .. }) => *end = row + 1,
-            let row = line.new_lineno().unwrap().saturating_sub(1);
+                    None => buffer_row_range = Some(row..row + 1),
-
+                }
-            match &mut buffer_row_range {
-                Some(buffer_row_range) => buffer_row_range.end = row + 1,
-                None => buffer_row_range = Some(row..row + 1),
             }
-        }
+            GitDiffLineType::Deletion => {
+                let end = content_offset + content_len;
 
-        if kind == GitDiffLineType::Deletion {
+                match &mut diff_base_byte_range {
-            let end = content_offset + content_len;
+                    Some(head_byte_range) => head_byte_range.end = end as usize,
+                    None => diff_base_byte_range = Some(content_offset as usize..end as usize),
+                }
 
-            match &mut diff_base_byte_range {
+                if first_deletion_buffer_row.is_none() {
-                Some(head_byte_range) => head_byte_range.end = end as usize,
+                    let old_row = line.old_lineno().unwrap().saturating_sub(1);
-                None => diff_base_byte_range = Some(content_offset as usize..end as usize),
+                    let row = old_row as i64 + *buffer_row_divergence;
+                    first_deletion_buffer_row = Some(row as u32);
+                }
+
+                *buffer_row_divergence -= 1;
             }
-
+            _ => {}
-            if first_deletion_buffer_row.is_none() {
-                let old_row = line.old_lineno().unwrap().saturating_sub(1);
-                let row = old_row as i64 + *buffer_row_divergence;
-                first_deletion_buffer_row = Some(row as u32);
-            }
-
-            *buffer_row_divergence -= 1;
         }
     }
 
-    //unwrap_or deletion without addition
     let buffer_row_range = buffer_row_range.unwrap_or_else(|| {
-        //we cannot have an addition-less hunk without deletion(s) or else there would be no hunk
+        // Pure deletion hunk without addition.
         let row = first_deletion_buffer_row.unwrap();
         row..row
     });
-
+    let diff_base_byte_range = diff_base_byte_range.unwrap_or_else(|| {
-    //unwrap_or addition without deletion
+        // Pure addition hunk without deletion.
-    let diff_base_byte_range = diff_base_byte_range.unwrap_or(0..0);
+        let row = first_addition_old_row.unwrap();
+        let offset = diff_base.point_to_offset(Point::new(row, 0));
+        offset..offset
+    });
 
     let start = Point::new(buffer_row_range.start, 0);
     let end = Point::new(buffer_row_range.end, 0);
@@ -499,9 +605,11 @@ impl BufferDiff {
         language_registry: Option<Arc<LanguageRegistry>>,
         cx: &mut App,
     ) -> impl Future<Output = BufferDiffInner> {
-        let base_text_snapshot = diff_base.as_ref().map(|base_text| {
+        let diff_base =
+            diff_base.map(|diff_base| (diff_base.clone(), Rope::from(diff_base.as_str())));
+        let base_text_snapshot = diff_base.as_ref().map(|(_, diff_base)| {
             language::Buffer::build_snapshot(
-                Rope::from(base_text.as_str()),
+                diff_base.clone(),
                 language.clone(),
                 language_registry.clone(),
                 cx,
@@ -528,6 +636,11 @@ impl BufferDiff {
         diff_base_buffer: Option<language::BufferSnapshot>,
         cx: &App,
     ) -> impl Future<Output = BufferDiffInner> {
+        let diff_base = diff_base.clone().zip(
+            diff_base_buffer
+                .clone()
+                .map(|buffer| buffer.as_rope().clone()),
+        );
         cx.background_executor().spawn(async move {
             BufferDiffInner {
                 hunks: compute_hunks(diff_base, buffer),
@@ -545,6 +658,7 @@ impl BufferDiff {
 
     pub fn build_with_single_insertion(
         insertion_present_in_secondary_diff: bool,
+        buffer: language::BufferSnapshot,
         cx: &mut App,
     ) -> BufferDiffSnapshot {
         let base_text = language::Buffer::build_empty_snapshot(cx);
@@ -560,17 +674,23 @@ impl BufferDiff {
                 hunks: hunks.clone(),
                 base_text: Some(base_text.clone()),
             },
-            secondary_diff: if insertion_present_in_secondary_diff {
+            secondary_diff: Some(Box::new(BufferDiffSnapshot {
-                Some(Box::new(BufferDiffSnapshot {
+                inner: BufferDiffInner {
-                    inner: BufferDiffInner {
+                    hunks: if insertion_present_in_secondary_diff {
-                        hunks,
+                        hunks
-                        base_text: Some(base_text),
+                    } else {
+                        SumTree::new(&buffer.text)
                     },
-                    secondary_diff: None,
+                    base_text: Some(if insertion_present_in_secondary_diff {
-                }))
+                        base_text
-            } else {
+                    } else {
-                None
+                        buffer
-            },
+                    }),
+                },
+                secondary_diff: None,
+                is_single_insertion: true,
+            })),
+            is_single_insertion: true,
         }
     }
 
@@ -675,6 +795,7 @@ impl BufferDiff {
                 .secondary_diff
                 .as_ref()
                 .map(|diff| Box::new(diff.read(cx).snapshot(cx))),
+            is_single_insertion: false,
         }
     }
 
@@ -875,13 +996,21 @@ pub fn assert_hunks<Iter>(
 
 #[cfg(test)]
 mod tests {
-    use std::assert_eq;
+    use std::fmt::Write as _;
 
     use super::*;
-    use gpui::TestAppContext;
+    use gpui::{AppContext as _, TestAppContext};
-    use text::{Buffer, BufferId};
+    use rand::{rngs::StdRng, Rng as _};
+    use text::{Buffer, BufferId, Rope};
     use unindent::Unindent as _;
 
+    #[ctor::ctor]
+    fn init_logger() {
+        if std::env::var("RUST_LOG").is_ok() {
+            env_logger::init();
+        }
+    }
+
     #[gpui::test]
     async fn test_buffer_diff_simple(cx: &mut gpui::TestAppContext) {
         let diff_base = "
@@ -1200,4 +1329,192 @@ mod tests {
         let range = diff_6.compare(&diff_5, &buffer).unwrap();
         assert_eq!(range.to_point(&buffer), Point::new(7, 0)..Point::new(8, 0));
     }
+
+    #[gpui::test(iterations = 100)]
+    async fn test_secondary_edits_for_stage_unstage(cx: &mut TestAppContext, mut rng: StdRng) {
+        fn gen_line(rng: &mut StdRng) -> String {
+            if rng.gen_bool(0.2) {
+                "\n".to_owned()
+            } else {
+                let c = rng.gen_range('A'..='Z');
+                format!("{c}{c}{c}\n")
+            }
+        }
+
+        fn gen_working_copy(rng: &mut StdRng, head: &str) -> String {
+            let mut old_lines = {
+                let mut old_lines = Vec::new();
+                let mut old_lines_iter = head.lines();
+                while let Some(line) = old_lines_iter.next() {
+                    assert!(!line.ends_with("\n"));
+                    old_lines.push(line.to_owned());
+                }
+                if old_lines.last().is_some_and(|line| line.is_empty()) {
+                    old_lines.pop();
+                }
+                old_lines.into_iter()
+            };
+            let mut result = String::new();
+            let unchanged_count = rng.gen_range(0..=old_lines.len());
+            result +=
+                &old_lines
+                    .by_ref()
+                    .take(unchanged_count)
+                    .fold(String::new(), |mut s, line| {
+                        writeln!(&mut s, "{line}").unwrap();
+                        s
+                    });
+            while old_lines.len() > 0 {
+                let deleted_count = rng.gen_range(0..=old_lines.len());
+                let _advance = old_lines
+                    .by_ref()
+                    .take(deleted_count)
+                    .map(|line| line.len() + 1)
+                    .sum::<usize>();
+                let minimum_added = if deleted_count == 0 { 1 } else { 0 };
+                let added_count = rng.gen_range(minimum_added..=5);
+                let addition = (0..added_count).map(|_| gen_line(rng)).collect::<String>();
+                result += &addition;
+
+                if old_lines.len() > 0 {
+                    let blank_lines = old_lines.clone().take_while(|line| line.is_empty()).count();
+                    if blank_lines == old_lines.len() {
+                        break;
+                    };
+                    let unchanged_count = rng.gen_range((blank_lines + 1).max(1)..=old_lines.len());
+                    result += &old_lines.by_ref().take(unchanged_count).fold(
+                        String::new(),
+                        |mut s, line| {
+                            writeln!(&mut s, "{line}").unwrap();
+                            s
+                        },
+                    );
+                }
+            }
+            result
+        }
+
+        fn uncommitted_diff(
+            working_copy: &language::BufferSnapshot,
+            index_text: &Entity<language::Buffer>,
+            head_text: String,
+            cx: &mut TestAppContext,
+        ) -> BufferDiff {
+            let inner = BufferDiff::build_sync(working_copy.text.clone(), head_text, cx);
+            let secondary = BufferDiff {
+                buffer_id: working_copy.remote_id(),
+                inner: BufferDiff::build_sync(
+                    working_copy.text.clone(),
+                    index_text.read_with(cx, |index_text, _| index_text.text()),
+                    cx,
+                ),
+                secondary_diff: None,
+            };
+            let secondary = cx.new(|_| secondary);
+            BufferDiff {
+                buffer_id: working_copy.remote_id(),
+                inner,
+                secondary_diff: Some(secondary),
+            }
+        }
+
+        let operations = std::env::var("OPERATIONS")
+            .map(|i| i.parse().expect("invalid `OPERATIONS` variable"))
+            .unwrap_or(10);
+
+        let rng = &mut rng;
+        let head_text = ('a'..='z').fold(String::new(), |mut s, c| {
+            writeln!(&mut s, "{c}{c}{c}").unwrap();
+            s
+        });
+        let working_copy = gen_working_copy(rng, &head_text);
+        let working_copy = cx.new(|cx| {
+            language::Buffer::local_normalized(
+                Rope::from(working_copy.as_str()),
+                text::LineEnding::default(),
+                cx,
+            )
+        });
+        let working_copy = working_copy.read_with(cx, |working_copy, _| working_copy.snapshot());
+        let index_text = cx.new(|cx| {
+            language::Buffer::local_normalized(
+                if rng.gen() {
+                    Rope::from(head_text.as_str())
+                } else {
+                    working_copy.as_rope().clone()
+                },
+                text::LineEnding::default(),
+                cx,
+            )
+        });
+
+        let mut diff = uncommitted_diff(&working_copy, &index_text, head_text.clone(), cx);
+        let mut hunks = cx.update(|cx| {
+            diff.hunks_intersecting_range(Anchor::MIN..Anchor::MAX, &working_copy, cx)
+                .collect::<Vec<_>>()
+        });
+        if hunks.len() == 0 {
+            return;
+        }
+
+        for _ in 0..operations {
+            let i = rng.gen_range(0..hunks.len());
+            let hunk = &mut hunks[i];
+            let hunk_fields = (
+                hunk.diff_base_byte_range.clone(),
+                hunk.secondary_diff_base_byte_range.clone(),
+                hunk.buffer_range.clone(),
+            );
+            let stage = match (
+                hunk.secondary_status,
+                hunk.secondary_diff_base_byte_range.clone(),
+            ) {
+                (DiffHunkSecondaryStatus::HasSecondaryHunk, Some(_)) => {
+                    hunk.secondary_status = DiffHunkSecondaryStatus::None;
+                    hunk.secondary_diff_base_byte_range = None;
+                    true
+                }
+                (DiffHunkSecondaryStatus::None, None) => {
+                    hunk.secondary_status = DiffHunkSecondaryStatus::HasSecondaryHunk;
+                    // We don't look at this, just notice whether it's Some or not.
+                    hunk.secondary_diff_base_byte_range = Some(17..17);
+                    false
+                }
+                _ => unreachable!(),
+            };
+
+            let snapshot = cx.update(|cx| diff.snapshot(cx));
+            let edits = snapshot.secondary_edits_for_stage_or_unstage(
+                stage,
+                [hunk_fields].into_iter(),
+                &working_copy,
+            );
+            index_text.update(cx, |index_text, cx| {
+                index_text.edit(edits, None, cx);
+            });
+
+            diff = uncommitted_diff(&working_copy, &index_text, head_text.clone(), cx);
+            let found_hunks = cx.update(|cx| {
+                diff.hunks_intersecting_range(Anchor::MIN..Anchor::MAX, &working_copy, cx)
+                    .collect::<Vec<_>>()
+            });
+            assert_eq!(hunks.len(), found_hunks.len());
+            for (expected_hunk, found_hunk) in hunks.iter().zip(&found_hunks) {
+                assert_eq!(
+                    expected_hunk.buffer_range.to_point(&working_copy),
+                    found_hunk.buffer_range.to_point(&working_copy)
+                );
+                assert_eq!(
+                    expected_hunk.diff_base_byte_range,
+                    found_hunk.diff_base_byte_range
+                );
+                assert_eq!(expected_hunk.secondary_status, found_hunk.secondary_status);
+                assert_eq!(
+                    expected_hunk.secondary_diff_base_byte_range.is_some(),
+                    found_hunk.secondary_diff_base_byte_range.is_some()
+                )
+            }
+            hunks = found_hunks;
+        }
+    }
 }

crates/collab/src/rpc.rs

@@ -395,6 +395,7 @@ impl Server {
             .add_request_handler(forward_mutating_project_request::<proto::Stage>)
             .add_request_handler(forward_mutating_project_request::<proto::Unstage>)
             .add_request_handler(forward_mutating_project_request::<proto::Commit>)
+            .add_request_handler(forward_mutating_project_request::<proto::SetIndexText>)
             .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>)
             .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>)
             .add_message_handler(update_context)

crates/editor/src/actions.rs

@@ -402,6 +402,7 @@ gpui::actions!(
         ToggleInlayHints,
         ToggleEditPrediction,
         ToggleLineNumbers,
+        ToggleStagedSelectedDiffHunks,
         SwapSelectionEnds,
         SetMark,
         ToggleRelativeLineNumbers,

crates/editor/src/editor.rs

@@ -52,6 +52,7 @@ pub use actions::{AcceptEditPrediction, OpenExcerpts, OpenExcerptsSplit};
 use aho_corasick::AhoCorasick;
 use anyhow::{anyhow, Context as _, Result};
 use blink_manager::BlinkManager;
+use buffer_diff::DiffHunkSecondaryStatus;
 use client::{Collaborator, ParticipantIndex};
 use clock::ReplicaId;
 use collections::{BTreeMap, HashMap, HashSet, VecDeque};
@@ -95,7 +96,7 @@ use itertools::Itertools;
 use language::{
     language_settings::{self, all_language_settings, language_settings, InlayHintSettings},
     markdown, point_from_lsp, AutoindentMode, BracketPair, Buffer, Capability, CharKind, CodeLabel,
-    CompletionDocumentation, CursorShape, Diagnostic, EditPredictionsMode, EditPreview,
+    CompletionDocumentation, CursorShape, Diagnostic, DiskState, EditPredictionsMode, EditPreview,
     HighlightedText, IndentKind, IndentSize, Language, OffsetRangeExt, Point, Selection,
     SelectionGoal, TextObject, TransactionId, TreeSitterOptions,
 };
@@ -12431,6 +12432,121 @@ impl Editor {
         self.toggle_diff_hunks_in_ranges(ranges, cx);
     }
 
+    fn diff_hunks_in_ranges<'a>(
+        &'a self,
+        ranges: &'a [Range<Anchor>],
+        buffer: &'a MultiBufferSnapshot,
+    ) -> impl 'a + Iterator<Item = MultiBufferDiffHunk> {
+        ranges.iter().flat_map(move |range| {
+            let end_excerpt_id = range.end.excerpt_id;
+            let range = range.to_point(buffer);
+            let mut peek_end = range.end;
+            if range.end.row < buffer.max_row().0 {
+                peek_end = Point::new(range.end.row + 1, 0);
+            }
+            buffer
+                .diff_hunks_in_range(range.start..peek_end)
+                .filter(move |hunk| hunk.excerpt_id.cmp(&end_excerpt_id, buffer).is_le())
+        })
+    }
+
+    pub fn has_stageable_diff_hunks_in_ranges(
+        &self,
+        ranges: &[Range<Anchor>],
+        snapshot: &MultiBufferSnapshot,
+    ) -> bool {
+        let mut hunks = self.diff_hunks_in_ranges(ranges, &snapshot);
+        hunks.any(|hunk| {
+            log::debug!("considering {hunk:?}");
+            hunk.secondary_status == DiffHunkSecondaryStatus::HasSecondaryHunk
+        })
+    }
+
+    pub fn toggle_staged_selected_diff_hunks(
+        &mut self,
+        _: &ToggleStagedSelectedDiffHunks,
+        _window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let ranges: Vec<_> = self.selections.disjoint.iter().map(|s| s.range()).collect();
+        self.stage_or_unstage_diff_hunks(&ranges, cx);
+    }
+
+    pub fn stage_or_unstage_diff_hunks(
+        &mut self,
+        ranges: &[Range<Anchor>],
+        cx: &mut Context<Self>,
+    ) {
+        let Some(project) = &self.project else {
+            return;
+        };
+        let snapshot = self.buffer.read(cx).snapshot(cx);
+        let stage = self.has_stageable_diff_hunks_in_ranges(ranges, &snapshot);
+
+        let chunk_by = self
+            .diff_hunks_in_ranges(&ranges, &snapshot)
+            .chunk_by(|hunk| hunk.buffer_id);
+        for (buffer_id, hunks) in &chunk_by {
+            let Some(buffer) = project.read(cx).buffer_for_id(buffer_id, cx) else {
+                log::debug!("no buffer for id");
+                continue;
+            };
+            let buffer = buffer.read(cx).snapshot();
+            let Some((repo, path)) = project
+                .read(cx)
+                .repository_and_path_for_buffer_id(buffer_id, cx)
+            else {
+                log::debug!("no git repo for buffer id");
+                continue;
+            };
+            let Some(diff) = snapshot.diff_for_buffer_id(buffer_id) else {
+                log::debug!("no diff for buffer id");
+                continue;
+            };
+            let Some(secondary_diff) = diff.secondary_diff() else {
+                log::debug!("no secondary diff for buffer id");
+                continue;
+            };
+
+            let edits = diff.secondary_edits_for_stage_or_unstage(
+                stage,
+                hunks.map(|hunk| {
+                    (
+                        hunk.diff_base_byte_range.clone(),
+                        hunk.secondary_diff_base_byte_range.clone(),
+                        hunk.buffer_range.clone(),
+                    )
+                }),
+                &buffer,
+            );
+
+            let index_base = secondary_diff.base_text().map_or_else(
+                || Rope::from(""),
+                |snapshot| snapshot.text.as_rope().clone(),
+            );
+            let index_buffer = cx.new(|cx| {
+                Buffer::local_normalized(index_base.clone(), text::LineEnding::default(), cx)
+            });
+            let new_index_text = index_buffer.update(cx, |index_buffer, cx| {
+                index_buffer.edit(edits, None, cx);
+                index_buffer.snapshot().as_rope().to_string()
+            });
+            let new_index_text = if new_index_text.is_empty()
+                && (diff.is_single_insertion
+                    || buffer
+                        .file()
+                        .map_or(false, |file| file.disk_state() == DiskState::New))
+            {
+                log::debug!("removing from index");
+                None
+            } else {
+                Some(new_index_text)
+            };
+
+            let _ = repo.read(cx).set_index_text(&path, new_index_text);
+        }
+    }
+
     pub fn expand_selected_diff_hunks(&mut self, cx: &mut Context<Self>) {
         let ranges: Vec<_> = self.selections.disjoint.iter().map(|s| s.range()).collect();
         self.buffer

crates/editor/src/editor_tests.rs

@@ -14047,6 +14047,59 @@ async fn test_edit_after_expanded_modification_hunk(
     );
 }
 
+#[gpui::test]
+async fn test_stage_and_unstage_added_file_hunk(
+    executor: BackgroundExecutor,
+    cx: &mut gpui::TestAppContext,
+) {
+    init_test(cx, |_| {});
+
+    let mut cx = EditorTestContext::new(cx).await;
+    cx.update_editor(|editor, _, cx| {
+        editor.set_expand_all_diff_hunks(cx);
+    });
+
+    let working_copy = r#"
+            ˇfn main() {
+                println!("hello, world!");
+            }
+        "#
+    .unindent();
+
+    cx.set_state(&working_copy);
+    executor.run_until_parked();
+
+    cx.assert_state_with_diff(
+        r#"
+            + ˇfn main() {
+            +     println!("hello, world!");
+            + }
+        "#
+        .unindent(),
+    );
+    cx.assert_index_text(None);
+
+    cx.update_editor(|editor, window, cx| {
+        editor.toggle_staged_selected_diff_hunks(&ToggleStagedSelectedDiffHunks, window, cx);
+    });
+    executor.run_until_parked();
+    cx.assert_index_text(Some(&working_copy.replace("ˇ", "")));
+    cx.assert_state_with_diff(
+        r#"
+            + ˇfn main() {
+            +     println!("hello, world!");
+            + }
+        "#
+        .unindent(),
+    );
+
+    cx.update_editor(|editor, window, cx| {
+        editor.toggle_staged_selected_diff_hunks(&ToggleStagedSelectedDiffHunks, window, cx);
+    });
+    executor.run_until_parked();
+    cx.assert_index_text(None);
+}
+
 async fn setup_indent_guides_editor(
     text: &str,
     cx: &mut gpui::TestAppContext,

crates/editor/src/element.rs

@@ -417,7 +417,9 @@ impl EditorElement {
         register_action(editor, window, Editor::toggle_git_blame);
         register_action(editor, window, Editor::toggle_git_blame_inline);
         register_action(editor, window, Editor::toggle_selected_diff_hunks);
+        register_action(editor, window, Editor::toggle_staged_selected_diff_hunks);
         register_action(editor, window, Editor::expand_all_diff_hunks);
+
         register_action(editor, window, |editor, action, window, cx| {
             if let Some(task) = editor.format(action, window, cx) {
                 task.detach_and_notify_err(window, cx);

crates/editor/src/test/editor_test_context.rs

@@ -298,6 +298,18 @@ impl EditorTestContext {
         self.cx.run_until_parked();
     }
 
+    pub fn assert_index_text(&mut self, expected: Option<&str>) {
+        let fs = self.update_editor(|editor, _, cx| {
+            editor.project.as_ref().unwrap().read(cx).fs().as_fake()
+        });
+        let path = self.update_buffer(|buffer, _| buffer.file().unwrap().path().clone());
+        let mut found = None;
+        fs.with_git_state(&Self::root_path().join(".git"), false, |git_state| {
+            found = git_state.index_contents.get(path.as_ref()).cloned();
+        });
+        assert_eq!(expected, found.as_deref());
+    }
+
     /// Change the editor's text and selections using a string containing
     /// embedded range markers that represent the ranges and directions of
     /// each selection.

crates/git/src/repository.rs

@@ -8,6 +8,8 @@ use gpui::SharedString;
 use parking_lot::Mutex;
 use rope::Rope;
 use std::borrow::Borrow;
+use std::io::Write as _;
+use std::process::Stdio;
 use std::sync::LazyLock;
 use std::{
     cmp::Ordering,
@@ -39,6 +41,8 @@ pub trait GitRepository: Send + Sync {
     /// Note that for symlink entries, this will return the contents of the symlink, not the target.
     fn load_committed_text(&self, path: &RepoPath) -> Option<String>;
 
+    fn set_index_text(&self, path: &RepoPath, content: Option<String>) -> anyhow::Result<()>;
+
     /// Returns the URL of the remote with the given name.
     fn remote_url(&self, name: &str) -> Option<String>;
     fn branch_name(&self) -> Option<String>;
@@ -161,6 +165,50 @@ impl GitRepository for RealGitRepository {
         Some(content)
     }
 
+    fn set_index_text(&self, path: &RepoPath, content: Option<String>) -> anyhow::Result<()> {
+        let working_directory = self
+            .repository
+            .lock()
+            .workdir()
+            .context("failed to read git work directory")?
+            .to_path_buf();
+        if let Some(content) = content {
+            let mut child = new_std_command(&self.git_binary_path)
+                .current_dir(&working_directory)
+                .args(["hash-object", "-w", "--stdin"])
+                .stdin(Stdio::piped())
+                .stdout(Stdio::piped())
+                .spawn()?;
+            child.stdin.take().unwrap().write_all(content.as_bytes())?;
+            let output = child.wait_with_output()?.stdout;
+            let sha = String::from_utf8(output)?;
+
+            log::debug!("indexing SHA: {sha}, path {path:?}");
+
+            let status = new_std_command(&self.git_binary_path)
+                .current_dir(&working_directory)
+                .args(["update-index", "--add", "--cacheinfo", "100644", &sha])
+                .arg(path.as_ref())
+                .status()?;
+
+            if !status.success() {
+                return Err(anyhow!("Failed to add to index: {status:?}"));
+            }
+        } else {
+            let status = new_std_command(&self.git_binary_path)
+                .current_dir(&working_directory)
+                .args(["update-index", "--force-remove"])
+                .arg(path.as_ref())
+                .status()?;
+
+            if !status.success() {
+                return Err(anyhow!("Failed to remove from index: {status:?}"));
+            }
+        }
+
+        Ok(())
+    }
+
     fn remote_url(&self, name: &str) -> Option<String> {
         let repo = self.repository.lock();
         let remote = repo.find_remote(name).ok()?;
@@ -412,6 +460,20 @@ impl GitRepository for FakeGitRepository {
         state.head_contents.get(path.as_ref()).cloned()
     }
 
+    fn set_index_text(&self, path: &RepoPath, content: Option<String>) -> anyhow::Result<()> {
+        let mut state = self.state.lock();
+        if let Some(content) = content {
+            state.index_contents.insert(path.clone(), content);
+        } else {
+            state.index_contents.remove(path);
+        }
+        state
+            .event_emitter
+            .try_send(state.path.clone())
+            .expect("Dropped repo change event");
+        Ok(())
+    }
+
     fn remote_url(&self, _name: &str) -> Option<String> {
         None
     }

crates/multi_buffer/src/multi_buffer.rs

@@ -133,6 +133,7 @@ pub struct MultiBufferDiffHunk {
     pub diff_base_byte_range: Range<usize>,
     /// Whether or not this hunk also appears in the 'secondary diff'.
     pub secondary_status: DiffHunkSecondaryStatus,
+    pub secondary_diff_base_byte_range: Option<Range<usize>>,
 }
 
 impl MultiBufferDiffHunk {
@@ -2191,7 +2192,11 @@ impl MultiBuffer {
             let secondary_diff_insertion = new_diff
                 .secondary_diff()
                 .map_or(true, |secondary_diff| secondary_diff.base_text().is_none());
-            new_diff = BufferDiff::build_with_single_insertion(secondary_diff_insertion, cx);
+            new_diff = BufferDiff::build_with_single_insertion(
+                secondary_diff_insertion,
+                buffer.snapshot(),
+                cx,
+            );
         }
 
         let mut snapshot = self.snapshot.borrow_mut();
@@ -3477,6 +3482,7 @@ impl MultiBufferSnapshot {
                 buffer_range: hunk.buffer_range.clone(),
                 diff_base_byte_range: hunk.diff_base_byte_range.clone(),
                 secondary_status: hunk.secondary_status,
+                secondary_diff_base_byte_range: hunk.secondary_diff_base_byte_range,
             })
         })
     }
@@ -3846,6 +3852,7 @@ impl MultiBufferSnapshot {
                         buffer_range: hunk.buffer_range.clone(),
                         diff_base_byte_range: hunk.diff_base_byte_range.clone(),
                         secondary_status: hunk.secondary_status,
+                        secondary_diff_base_byte_range: hunk.secondary_diff_base_byte_range,
                     });
                 }
             }
@@ -5937,6 +5944,10 @@ impl MultiBufferSnapshot {
     pub fn show_headers(&self) -> bool {
         self.show_headers
     }
+
+    pub fn diff_for_buffer_id(&self, buffer_id: BufferId) -> Option<&BufferDiffSnapshot> {
+        self.diffs.get(&buffer_id)
+    }
 }
 
 #[cfg(any(test, feature = "test-support"))]

crates/project/src/buffer_store.rs

@@ -189,6 +189,7 @@ impl BufferDiffState {
         buffer: text::BufferSnapshot,
         cx: &mut Context<Self>,
     ) -> oneshot::Receiver<()> {
+        log::debug!("recalculate diffs");
         let (tx, rx) = oneshot::channel();
         self.diff_updated_futures.push(tx);
 

crates/project/src/git.rs

@@ -23,11 +23,11 @@ use util::{maybe, ResultExt};
 use worktree::{ProjectEntryId, RepositoryEntry, StatusEntry};
 
 pub struct GitState {
-    project_id: Option<ProjectId>,
+    pub(super) project_id: Option<ProjectId>,
-    client: Option<AnyProtoClient>,
+    pub(super) client: Option<AnyProtoClient>,
+    pub update_sender: mpsc::UnboundedSender<(Message, oneshot::Sender<anyhow::Result<()>>)>,
     repositories: Vec<Entity<Repository>>,
     active_index: Option<usize>,
-    update_sender: mpsc::UnboundedSender<(Message, oneshot::Sender<anyhow::Result<()>>)>,
     _subscription: Subscription,
 }
 
@@ -51,7 +51,7 @@ pub enum GitRepo {
     },
 }
 
-enum Message {
+pub enum Message {
     Commit {
         git_repo: GitRepo,
         message: SharedString,
@@ -59,6 +59,7 @@ enum Message {
     },
     Stage(GitRepo, Vec<RepoPath>),
     Unstage(GitRepo, Vec<RepoPath>),
+    SetIndexText(GitRepo, RepoPath, Option<String>),
 }
 
 pub enum GitEvent {
@@ -291,11 +292,32 @@ impl GitState {
                 }
                 Ok(())
             }
+            Message::SetIndexText(git_repo, path, text) => match git_repo {
+                GitRepo::Local(repo) => repo.set_index_text(&path, text),
+                GitRepo::Remote {
+                    project_id,
+                    client,
+                    worktree_id,
+                    work_directory_id,
+                } => client.send(proto::SetIndexText {
+                    project_id: project_id.0,
+                    worktree_id: worktree_id.to_proto(),
+                    work_directory_id: work_directory_id.to_proto(),
+                    path: path.as_ref().to_proto(),
+                    text,
+                }),
+            },
         }
     }
 }
 
+impl GitRepo {}
+
 impl Repository {
+    pub fn git_state(&self) -> Option<Entity<GitState>> {
+        self.git_state.upgrade()
+    }
+
     fn id(&self) -> (WorktreeId, ProjectEntryId) {
         (self.worktree_id, self.repository_entry.work_directory_id())
     }
@@ -522,4 +544,19 @@ impl Repository {
             .ok();
         result_rx
     }
+
+    pub fn set_index_text(
+        &self,
+        path: &RepoPath,
+        content: Option<String>,
+    ) -> oneshot::Receiver<anyhow::Result<()>> {
+        let (result_tx, result_rx) = futures::channel::oneshot::channel();
+        self.update_sender
+            .unbounded_send((
+                Message::SetIndexText(self.git_repo.clone(), path.clone(), content),
+                result_tx,
+            ))
+            .ok();
+        result_rx
+    }
 }

crates/project/src/project.rs

@@ -610,6 +610,7 @@ impl Project {
         client.add_entity_request_handler(Self::handle_stage);
         client.add_entity_request_handler(Self::handle_unstage);
         client.add_entity_request_handler(Self::handle_commit);
+        client.add_entity_request_handler(Self::handle_set_index_text);
         client.add_entity_request_handler(Self::handle_open_commit_message_buffer);
 
         WorktreeStore::init(&client);
@@ -4092,6 +4093,27 @@ impl Project {
         Ok(proto::Ack {})
     }
 
+    async fn handle_set_index_text(
+        this: Entity<Self>,
+        envelope: TypedEnvelope<proto::SetIndexText>,
+        mut cx: AsyncApp,
+    ) -> Result<proto::Ack> {
+        let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
+        let work_directory_id = ProjectEntryId::from_proto(envelope.payload.work_directory_id);
+        let repository_handle =
+            Self::repository_for_request(&this, worktree_id, work_directory_id, &mut cx)?;
+
+        repository_handle
+            .update(&mut cx, |repository_handle, _| {
+                repository_handle.set_index_text(
+                    &RepoPath::from_str(&envelope.payload.path),
+                    envelope.payload.text,
+                )
+            })?
+            .await??;
+        Ok(proto::Ack {})
+    }
+
     async fn handle_open_commit_message_buffer(
         this: Entity<Self>,
         envelope: TypedEnvelope<proto::OpenCommitMessageBuffer>,
@@ -4336,6 +4358,27 @@ impl Project {
     pub fn all_repositories(&self, cx: &App) -> Vec<Entity<Repository>> {
         self.git_state.read(cx).all_repositories()
     }
+
+    pub fn repository_and_path_for_buffer_id(
+        &self,
+        buffer_id: BufferId,
+        cx: &App,
+    ) -> Option<(Entity<Repository>, RepoPath)> {
+        let path = self
+            .buffer_for_id(buffer_id, cx)?
+            .read(cx)
+            .project_path(cx)?;
+        self.git_state
+            .read(cx)
+            .all_repositories()
+            .into_iter()
+            .find_map(|repo| {
+                Some((
+                    repo.clone(),
+                    repo.read(cx).repository_entry.relativize(&path.path).ok()?,
+                ))
+            })
+    }
 }
 
 fn deserialize_code_actions(code_actions: &HashMap<String, bool>) -> Vec<lsp::CodeActionKind> {

crates/proto/proto/zed.proto

@@ -315,7 +315,9 @@ message Envelope {
         OpenCommitMessageBuffer open_commit_message_buffer = 296;
 
         OpenUncommittedDiff open_uncommitted_diff = 297;
-        OpenUncommittedDiffResponse open_uncommitted_diff_response = 298; // current max
+        OpenUncommittedDiffResponse open_uncommitted_diff_response = 298;
+
+        SetIndexText set_index_text = 299; // current max
     }
 
     reserved 87 to 88;
@@ -2087,6 +2089,14 @@ message OpenUncommittedDiffResponse {
     Mode mode = 3;
 }
 
+message SetIndexText {
+    uint64 project_id = 1;
+    uint64 worktree_id = 2;
+    uint64 work_directory_id = 3;
+    string path = 4;
+    optional string text = 5;
+}
+
 message GetNotifications {
     optional uint64 before_id = 1;
 }

crates/proto/src/proto.rs

@@ -440,6 +440,7 @@ messages!(
     (SyncExtensionsResponse, Background),
     (InstallExtension, Background),
     (RegisterBufferWithLanguageServers, Background),
+    (SetIndexText, Background),
 );
 
 request_messages!(
@@ -573,6 +574,7 @@ request_messages!(
     (SyncExtensions, SyncExtensionsResponse),
     (InstallExtension, Ack),
     (RegisterBufferWithLanguageServers, Ack),
+    (SetIndexText, Ack),
 );
 
 entity_messages!(
@@ -665,6 +667,7 @@ entity_messages!(
     GetPathMetadata,
     CancelLanguageServerWork,
     RegisterBufferWithLanguageServers,
+    SetIndexText,
 );
 
 entity_messages!(

crates/remote_server/src/headless_project.rs

@@ -200,6 +200,7 @@ impl HeadlessProject {
         client.add_entity_request_handler(Self::handle_stage);
         client.add_entity_request_handler(Self::handle_unstage);
         client.add_entity_request_handler(Self::handle_commit);
+        client.add_entity_request_handler(Self::handle_set_index_text);
         client.add_entity_request_handler(Self::handle_open_commit_message_buffer);
 
         client.add_request_handler(
@@ -691,6 +692,26 @@ impl HeadlessProject {
         Ok(proto::Ack {})
     }
 
+    async fn handle_set_index_text(
+        this: Entity<Self>,
+        envelope: TypedEnvelope<proto::SetIndexText>,
+        mut cx: AsyncApp,
+    ) -> Result<proto::Ack> {
+        let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
+        let work_directory_id = ProjectEntryId::from_proto(envelope.payload.work_directory_id);
+        let repository =
+            Self::repository_for_request(&this, worktree_id, work_directory_id, &mut cx)?;
+        repository
+            .update(&mut cx, |repository, _| {
+                repository.set_index_text(
+                    &RepoPath::from(envelope.payload.path.as_str()),
+                    envelope.payload.text,
+                )
+            })?
+            .await??;
+        Ok(proto::Ack {})
+    }
+
     async fn handle_open_commit_message_buffer(
         this: Entity<Self>,
         envelope: TypedEnvelope<proto::OpenCommitMessageBuffer>,

crates/sum_tree/src/cursor.rs

@@ -447,7 +447,7 @@ where
         summary.0
     }
 
-    /// Returns whether we found the item you where seeking for
+    /// Returns whether we found the item you were seeking for
     #[track_caller]
     fn seek_internal(
         &mut self,