vim: Add U to undo last line (#33571)

Closes #14760

Still TODO:

* Vim actually undoes *many* changes if they're all on the same line.

Release Notes:

- vim: Add `U` to return to the last changed line and undo

assets/keymaps/vim.json

@@ -364,6 +364,7 @@
       "p": "vim::Paste",
       "shift-p": ["vim::Paste", { "before": true }],
       "u": "vim::Undo",
+      "shift-u": "vim::UndoLastLine",
       "r": "vim::PushReplace",
       "s": "vim::Substitute",
       "shift-s": "vim::SubstituteLine",

crates/editor/src/editor.rs

@@ -865,9 +865,19 @@ pub trait Addon: 'static {
     }
 }
 
+struct ChangeLocation {
+    current: Option<Vec<Anchor>>,
+    original: Vec<Anchor>,
+}
+impl ChangeLocation {
+    fn locations(&self) -> &[Anchor] {
+        self.current.as_ref().unwrap_or(&self.original)
+    }
+}
+
 /// A set of caret positions, registered when the editor was edited.
 pub struct ChangeList {
-    changes: Vec<Vec<Anchor>>,
+    changes: Vec<ChangeLocation>,
     /// Currently "selected" change.
     position: Option<usize>,
 }
@@ -894,20 +904,38 @@ impl ChangeList {
             (prev + count).min(self.changes.len() - 1)
         };
         self.position = Some(next);
-        self.changes.get(next).map(|anchors| anchors.as_slice())
+        self.changes.get(next).map(|change| change.locations())
     }
 
     /// Adds a new change to the list, resetting the change list position.
-    pub fn push_to_change_list(&mut self, pop_state: bool, new_positions: Vec<Anchor>) {
+    pub fn push_to_change_list(&mut self, group: bool, new_positions: Vec<Anchor>) {
         self.position.take();
-        if pop_state {
-            self.changes.pop();
+        if let Some(last) = self.changes.last_mut()
+            && group
+        {
+            last.current = Some(new_positions)
+        } else {
+            self.changes.push(ChangeLocation {
+                original: new_positions,
+                current: None,
+            });
         }
-        self.changes.push(new_positions.clone());
     }
 
     pub fn last(&self) -> Option<&[Anchor]> {
-        self.changes.last().map(|anchors| anchors.as_slice())
+        self.changes.last().map(|change| change.locations())
+    }
+
+    pub fn last_before_grouping(&self) -> Option<&[Anchor]> {
+        self.changes.last().map(|change| change.original.as_slice())
+    }
+
+    pub fn invert_last_group(&mut self) {
+        if let Some(last) = self.changes.last_mut() {
+            if let Some(current) = last.current.as_mut() {
+                mem::swap(&mut last.original, current);
+            }
+        }
     }
 }
 

crates/vim/src/normal.rs

@@ -24,9 +24,9 @@ use crate::{
 };
 use collections::BTreeSet;
 use convert::ConvertTarget;
-use editor::Bias;
 use editor::Editor;
 use editor::{Anchor, SelectionEffects};
+use editor::{Bias, ToPoint};
 use editor::{display_map::ToDisplayPoint, movement};
 use gpui::{Context, Window, actions};
 use language::{Point, SelectionGoal};
@@ -90,6 +90,8 @@ actions!(
         Undo,
         /// Redoes the last undone change.
         Redo,
+        /// Undoes all changes to the most recently changed line.
+        UndoLastLine,
     ]
 );
 
@@ -194,6 +196,120 @@ pub(crate) fn register(editor: &mut Editor, cx: &mut Context<Vim>) {
             }
         });
     });
+    Vim::action(editor, cx, |vim, _: &UndoLastLine, window, cx| {
+        Vim::take_forced_motion(cx);
+        vim.update_editor(window, cx, |vim, editor, window, cx| {
+            let snapshot = editor.buffer().read(cx).snapshot(cx);
+            let Some(last_change) = editor.change_list.last_before_grouping() else {
+                return;
+            };
+
+            let anchors = last_change.iter().cloned().collect::<Vec<_>>();
+            let mut last_row = None;
+            let ranges: Vec<_> = anchors
+                .iter()
+                .filter_map(|anchor| {
+                    let point = anchor.to_point(&snapshot);
+                    if last_row == Some(point.row) {
+                        return None;
+                    }
+                    last_row = Some(point.row);
+                    let line_range = Point::new(point.row, 0)
+                        ..Point::new(point.row, snapshot.line_len(MultiBufferRow(point.row)));
+                    Some((
+                        snapshot.anchor_before(line_range.start)
+                            ..snapshot.anchor_after(line_range.end),
+                        line_range,
+                    ))
+                })
+                .collect();
+
+            let edits = editor.buffer().update(cx, |buffer, cx| {
+                let current_content = ranges
+                    .iter()
+                    .map(|(anchors, _)| {
+                        buffer
+                            .snapshot(cx)
+                            .text_for_range(anchors.clone())
+                            .collect::<String>()
+                    })
+                    .collect::<Vec<_>>();
+                let mut content_before_undo = current_content.clone();
+                let mut undo_count = 0;
+
+                loop {
+                    let undone_tx = buffer.undo(cx);
+                    undo_count += 1;
+                    let mut content_after_undo = Vec::new();
+
+                    let mut line_changed = false;
+                    for ((anchors, _), text_before_undo) in
+                        ranges.iter().zip(content_before_undo.iter())
+                    {
+                        let snapshot = buffer.snapshot(cx);
+                        let text_after_undo =
+                            snapshot.text_for_range(anchors.clone()).collect::<String>();
+
+                        if &text_after_undo != text_before_undo {
+                            line_changed = true;
+                        }
+                        content_after_undo.push(text_after_undo);
+                    }
+
+                    content_before_undo = content_after_undo;
+                    if !line_changed {
+                        break;
+                    }
+                    if undone_tx == vim.undo_last_line_tx {
+                        break;
+                    }
+                }
+
+                let edits = ranges
+                    .into_iter()
+                    .zip(content_before_undo.into_iter().zip(current_content))
+                    .filter_map(|((_, mut points), (mut old_text, new_text))| {
+                        if new_text == old_text {
+                            return None;
+                        }
+                        let common_suffix_starts_at = old_text
+                            .char_indices()
+                            .rev()
+                            .zip(new_text.chars().rev())
+                            .find_map(
+                                |((i, a), b)| {
+                                    if a != b { Some(i + a.len_utf8()) } else { None }
+                                },
+                            )
+                            .unwrap_or(old_text.len());
+                        points.end.column -= (old_text.len() - common_suffix_starts_at) as u32;
+                        old_text = old_text.split_at(common_suffix_starts_at).0.to_string();
+                        let common_prefix_len = old_text
+                            .char_indices()
+                            .zip(new_text.chars())
+                            .find_map(|((i, a), b)| if a != b { Some(i) } else { None })
+                            .unwrap_or(0);
+                        points.start.column = common_prefix_len as u32;
+                        old_text = old_text.split_at(common_prefix_len).1.to_string();
+
+                        Some((points, old_text))
+                    })
+                    .collect::<Vec<_>>();
+
+                for _ in 0..undo_count {
+                    buffer.redo(cx);
+                }
+                edits
+            });
+            vim.undo_last_line_tx = editor.transact(window, cx, |editor, window, cx| {
+                editor.change_list.invert_last_group();
+                editor.edit(edits, cx);
+                editor.change_selections(SelectionEffects::default(), window, cx, |s| {
+                    s.select_anchor_ranges(anchors.into_iter().map(|a| a..a));
+                })
+            });
+        });
+    });
 
     repeat::register(editor, cx);
     scroll::register(editor, cx);
@@ -1876,4 +1992,102 @@ mod test {
         cx.simulate_shared_keystrokes("ctrl-o").await;
         cx.shared_state().await.assert_matches();
     }
+
+    #[gpui::test]
+    async fn test_undo_last_line(cx: &mut gpui::TestAppContext) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state(indoc! {"
+            ˇfn a() { }
+            fn a() { }
+            fn a() { }
+        "})
+            .await;
+        // do a jump to reset vim's undo grouping
+        cx.simulate_shared_keystrokes("shift-g").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("r a").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("g g shift-u").await;
+        cx.shared_state().await.assert_matches();
+    }
+
+    #[gpui::test]
+    async fn test_undo_last_line_newline(cx: &mut gpui::TestAppContext) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state(indoc! {"
+            ˇfn a() { }
+            fn a() { }
+            fn a() { }
+        "})
+            .await;
+        // do a jump to reset vim's undo grouping
+        cx.simulate_shared_keystrokes("shift-g k").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("o h e l l o escape").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+    }
+
+    #[gpui::test]
+    async fn test_undo_last_line_newline_many_changes(cx: &mut gpui::TestAppContext) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state(indoc! {"
+            ˇfn a() { }
+            fn a() { }
+            fn a() { }
+        "})
+            .await;
+        // do a jump to reset vim's undo grouping
+        cx.simulate_shared_keystrokes("x shift-g k").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("x f a x f { x").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+        cx.simulate_shared_keystrokes("shift-u").await;
+        cx.shared_state().await.assert_matches();
+    }
+
+    #[gpui::test]
+    async fn test_undo_last_line_multicursor(cx: &mut gpui::TestAppContext) {
+        let mut cx = VimTestContext::new(cx, true).await;
+
+        cx.set_state(
+            indoc! {"
+            ˇone two ˇone
+            two ˇone two
+        "},
+            Mode::Normal,
+        );
+        cx.simulate_keystrokes("3 r a");
+        cx.assert_state(
+            indoc! {"
+            aaˇa two aaˇa
+            two aaˇa two
+        "},
+            Mode::Normal,
+        );
+        cx.simulate_keystrokes("escape escape");
+        cx.simulate_keystrokes("shift-u");
+        cx.set_state(
+            indoc! {"
+            onˇe two onˇe
+            two onˇe two
+        "},
+            Mode::Normal,
+        );
+    }
 }

crates/vim/src/vim.rs

@@ -375,6 +375,7 @@ pub(crate) struct Vim {
     pub(crate) current_tx: Option<TransactionId>,
     pub(crate) current_anchor: Option<Selection<Anchor>>,
     pub(crate) undo_modes: HashMap<TransactionId, Mode>,
+    pub(crate) undo_last_line_tx: Option<TransactionId>,
 
     selected_register: Option<char>,
     pub search: SearchState,
@@ -422,6 +423,7 @@ impl Vim {
 
             stored_visual_mode: None,
             current_tx: None,
+            undo_last_line_tx: None,
             current_anchor: None,
             undo_modes: HashMap::default(),
 

crates/vim/test_data/test_undo_last_line.json

@@ -0,0 +1,14 @@
+{"Put":{"state":"ˇfn a() { }\nfn a() { }\nfn a() { }\n"}}
+{"Key":"shift-g"}
+{"Get":{"state":"fn a() { }\nfn a() { }\nfn a() { }\nˇ","mode":"Normal"}}
+{"Key":"r"}
+{"Key":"a"}
+{"Get":{"state":"fn a() { }\nfn a() { }\nfn a() { }\nˇ","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"ˇ\nfn a() { }\nfn a() { }\n","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"ˇfn a() { }\nfn a() { }\nfn a() { }\n","mode":"Normal"}}
+{"Key":"g"}
+{"Key":"g"}
+{"Key":"shift-u"}
+{"Get":{"state":"ˇ\nfn a() { }\nfn a() { }\n","mode":"Normal"}}

crates/vim/test_data/test_undo_last_line_newline.json

@@ -0,0 +1,15 @@
+{"Put":{"state":"ˇfn a() { }\nfn a() { }\nfn a() { }\n"}}
+{"Key":"shift-g"}
+{"Key":"k"}
+{"Get":{"state":"fn a() { }\nfn a() { }\nˇfn a() { }\n","mode":"Normal"}}
+{"Key":"o"}
+{"Key":"h"}
+{"Key":"e"}
+{"Key":"l"}
+{"Key":"l"}
+{"Key":"o"}
+{"Key":"escape"}
+{"Get":{"state":"fn a() { }\nfn a() { }\nfn a() { }\nhellˇo\n","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"fn a() { }\nfn a() { }\nfn a() { }\nˇ\n","mode":"Normal"}}
+{"Key":"shift-u"}

crates/vim/test_data/test_undo_last_line_newline_many_changes.json

@@ -0,0 +1,21 @@
+{"Put":{"state":"ˇfn a() { }\nfn a() { }\nfn a() { }\n"}}
+{"Key":"x"}
+{"Key":"shift-g"}
+{"Key":"k"}
+{"Get":{"state":"n a() { }\nfn a() { }\nˇfn a() { }\n","mode":"Normal"}}
+{"Key":"x"}
+{"Key":"f"}
+{"Key":"a"}
+{"Key":"x"}
+{"Key":"f"}
+{"Key":"{"}
+{"Key":"x"}
+{"Get":{"state":"n a() { }\nfn a() { }\nn () ˇ }\n","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"n a() { }\nfn a() { }\nˇfn a() { }\n","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"n a() { }\nfn a() { }\nn () ˇ }\n","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"n a() { }\nfn a() { }\nˇfn a() { }\n","mode":"Normal"}}
+{"Key":"shift-u"}
+{"Get":{"state":"n a() { }\nfn a() { }\nn () ˇ }\n","mode":"Normal"}}