Add command for resetting diff hunks (#5736)

2 files changed, 76 insertions, 11 deletions

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index c76e9f2b..ebdfdfde 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -3011,13 +3011,13 @@ fn goto_first_change_impl(cx: &mut Context, reverse: bool) {
     let (view, doc) = current!(editor);
     if let Some(handle) = doc.diff_handle() {
         let hunk = {
-            let hunks = handle.hunks();
+            let diff = handle.load();
             let idx = if reverse {
-                hunks.len().saturating_sub(1)
+                diff.len().saturating_sub(1)
             } else {
                 0
             };
-            hunks.nth_hunk(idx)
+            diff.nth_hunk(idx)
         };
         if hunk != Hunk::NONE {
             let range = hunk_range(hunk, doc.text().slice(..));
@@ -3049,19 +3049,19 @@ fn goto_next_change_impl(cx: &mut Context, direction: Direction) {
         let selection = doc.selection(view.id).clone().transform(|range| {
             let cursor_line = range.cursor_line(doc_text) as u32;
 
-            let hunks = diff_handle.hunks();
+            let diff = diff_handle.load();
             let hunk_idx = match direction {
-                Direction::Forward => hunks
+                Direction::Forward => diff
                     .next_hunk(cursor_line)
-                    .map(|idx| (idx + count).min(hunks.len() - 1)),
-                Direction::Backward => hunks
+                    .map(|idx| (idx + count).min(diff.len() - 1)),
+                Direction::Backward => diff
                     .prev_hunk(cursor_line)
                     .map(|idx| idx.saturating_sub(count)),
             };
             let Some(hunk_idx) = hunk_idx else {
                 return range;
             };
-            let hunk = hunks.nth_hunk(hunk_idx);
+            let hunk = diff.nth_hunk(hunk_idx);
             let new_range = hunk_range(hunk, doc_text);
             if editor.mode == Mode::Select {
                 let head = if new_range.head < range.anchor {
@@ -4721,14 +4721,14 @@ fn select_textobject(cx: &mut Context, objtype: textobject::TextObject) {
 
                 let textobject_change = |range: Range| -> Range {
                     let diff_handle = doc.diff_handle().unwrap();
-                    let hunks = diff_handle.hunks();
+                    let diff = diff_handle.load();
                     let line = range.cursor_line(text);
-                    let hunk_idx = if let Some(hunk_idx) = hunks.hunk_at(line as u32, false) {
+                    let hunk_idx = if let Some(hunk_idx) = diff.hunk_at(line as u32, false) {
                         hunk_idx
                     } else {
                         return range;
                     };
-                    let hunk = hunks.nth_hunk(hunk_idx).after;
+                    let hunk = diff.nth_hunk(hunk_idx).after;
 
                     let start = text.line_to_char(hunk.start as usize);
                     let end = text.line_to_char(hunk.end as usize);
diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index 0ddca6df..77c14321 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -2033,6 +2033,64 @@ fn run_shell_command(
     Ok(())
 }
 
+fn reset_diff_change(
+    cx: &mut compositor::Context,
+    args: &[Cow<str>],
+    event: PromptEvent,
+) -> anyhow::Result<()> {
+    if event != PromptEvent::Validate {
+        return Ok(());
+    }
+    ensure!(args.is_empty(), ":reset-diff-change takes no arguments");
+
+    let editor = &mut cx.editor;
+    let scrolloff = editor.config().scrolloff;
+
+    let (view, doc) = current!(editor);
+    // TODO refactor to use let..else once MSRV is raised to 1.65
+    let handle = match doc.diff_handle() {
+        Some(handle) => handle,
+        None => bail!("Diff is not available in the current buffer"),
+    };
+
+    let diff = handle.load();
+    let doc_text = doc.text().slice(..);
+    let line = doc.selection(view.id).primary().cursor_line(doc_text);
+
+    // TODO refactor to use let..else once MSRV is raised to 1.65
+    let hunk_idx = match diff.hunk_at(line as u32, true) {
+        Some(hunk_idx) => hunk_idx,
+        None => bail!("There is no change at the cursor"),
+    };
+    let hunk = diff.nth_hunk(hunk_idx);
+    let diff_base = diff.diff_base();
+    let before_start = diff_base.line_to_char(hunk.before.start as usize);
+    let before_end = diff_base.line_to_char(hunk.before.end as usize);
+    let text: Tendril = diff
+        .diff_base()
+        .slice(before_start..before_end)
+        .chunks()
+        .collect();
+    let anchor = doc_text.line_to_char(hunk.after.start as usize);
+    let transaction = Transaction::change(
+        doc.text(),
+        [(
+            anchor,
+            doc_text.line_to_char(hunk.after.end as usize),
+            (!text.is_empty()).then_some(text),
+        )]
+        .into_iter(),
+    );
+    drop(diff); // make borrow check happy
+    doc.apply(&transaction, view.id);
+    // select inserted text
+    let text_len = before_end - before_start;
+    doc.set_selection(view.id, Selection::single(anchor, anchor + text_len));
+    doc.append_changes_to_history(view);
+    view.ensure_cursor_in_view(doc, scrolloff);
+    Ok(())
+}
+
 pub const TYPABLE_COMMAND_LIST: &[TypableCommand] = &[
         TypableCommand {
             name: "quit",
@@ -2569,6 +2627,13 @@ pub const TYPABLE_COMMAND_LIST: &[TypableCommand] = &[
             fun: run_shell_command,
             completer: Some(completers::filename),
         },
+       TypableCommand {
+            name: "reset-diff-change",
+            aliases: &["diffget", "diffg"],
+            doc: "Reset the diff change at the cursor position.",
+            fun: reset_diff_change,
+            completer: None,
+        },
     ];
 
 pub static TYPABLE_COMMAND_MAP: Lazy<HashMap<&'static str, &'static TypableCommand>> =