Improve vim undo (#9317)

The important change here is to ensure that undo never lands you in
visual mode; but we also take care to restore the selection the same way
vim does (visual line goes to beginning of line, visual block to the top
left, etc.).

To help make this behaviour feel right we also group any deletions that
started insert mode with the first text inserted.

Fixes: #7521

Release Notes:

- vim: Improved undo. It will now restore you to normal mode in the same
position as vim, and group deletions caused by `c` or `s` with the
concomitant insert.
([#7521](https://github.com/zed-industries/zed/issues/7521)).

assets/keymaps/vim.json

@@ -515,7 +515,9 @@
       "ctrl-w": "editor::DeleteToPreviousWordStart",
       "ctrl-u": "editor::DeleteToBeginningOfLine",
       "ctrl-t": "vim::Indent",
-      "ctrl-d": "vim::Outdent"
+      "ctrl-d": "vim::Outdent",
+      "ctrl-r \"": "editor::Paste",
+      "ctrl-r +": "editor::Paste"
     }
   },
   {

crates/editor/src/editor.rs

@@ -5697,6 +5697,9 @@ impl Editor {
             self.unmark_text(cx);
             self.refresh_copilot_suggestions(true, cx);
             cx.emit(EditorEvent::Edited);
+            cx.emit(EditorEvent::TransactionUndone {
+                transaction_id: tx_id,
+            });
         }
     }
 
@@ -5724,6 +5727,11 @@ impl Editor {
             .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
     }
 
+    pub fn group_until_transaction(&mut self, tx_id: TransactionId, cx: &mut ViewContext<Self>) {
+        self.buffer
+            .update(cx, |buffer, cx| buffer.group_until_transaction(tx_id, cx));
+    }
+
     pub fn move_left(&mut self, _: &MoveLeft, cx: &mut ViewContext<Self>) {
         self.change_selections(Some(Autoscroll::fit()), cx, |s| {
             let line_mode = s.line_mode;
@@ -8551,6 +8559,9 @@ impl Editor {
         {
             self.selection_history
                 .insert_transaction(tx_id, self.selections.disjoint_anchors());
+            cx.emit(EditorEvent::TransactionBegun {
+                transaction_id: tx_id,
+            })
         }
     }
 
@@ -10165,6 +10176,12 @@ pub enum EditorEvent {
         autoscroll: bool,
     },
     Closed,
+    TransactionUndone {
+        transaction_id: clock::Lamport,
+    },
+    TransactionBegun {
+        transaction_id: clock::Lamport,
+    },
 }
 
 impl EventEmitter<EditorEvent> for Editor {}

crates/vim/src/editor_events.rs

@@ -29,7 +29,6 @@ pub fn init(cx: &mut AppContext) {
     })
     .detach();
 }
-
 fn focused(editor: View<Editor>, cx: &mut WindowContext) {
     Vim::update(cx, |vim, cx| {
         if !vim.enabled {

crates/vim/src/state.rs

@@ -1,8 +1,9 @@
 use std::{fmt::Display, ops::Range, sync::Arc};
 
 use collections::HashMap;
+use editor::Anchor;
 use gpui::{Action, KeyContext};
-use language::CursorShape;
+use language::{CursorShape, Selection, TransactionId};
 use serde::{Deserialize, Serialize};
 use workspace::searchable::Direction;
 
@@ -66,6 +67,10 @@ pub struct EditorState {
     pub post_count: Option<usize>,
 
     pub operator_stack: Vec<Operator>,
+
+    pub current_tx: Option<TransactionId>,
+    pub current_anchor: Option<Selection<Anchor>>,
+    pub undo_modes: HashMap<TransactionId, Mode>,
 }
 
 #[derive(Default, Clone, Debug)]

crates/vim/src/test.rs

@@ -981,3 +981,55 @@ async fn test_remap(cx: &mut gpui::TestAppContext) {
     cx.simulate_keystrokes(["g", "t"]);
     cx.assert_state("12ˇ 34", Mode::Normal);
 }
+
+#[gpui::test]
+async fn test_undo(cx: &mut gpui::TestAppContext) {
+    let mut cx = NeovimBackedTestContext::new(cx).await;
+
+    cx.set_shared_state("hello quˇoel world").await;
+    cx.simulate_shared_keystrokes(["v", "i", "w", "s", "c", "o", "escape", "u"])
+        .await;
+    cx.assert_shared_state("hello ˇquoel world").await;
+    cx.simulate_shared_keystrokes(["ctrl-r"]).await;
+    cx.assert_shared_state("hello ˇco world").await;
+    cx.simulate_shared_keystrokes(["a", "o", "right", "l", "escape"])
+        .await;
+    cx.assert_shared_state("hello cooˇl world").await;
+    cx.simulate_shared_keystrokes(["u"]).await;
+    cx.assert_shared_state("hello cooˇ world").await;
+    cx.simulate_shared_keystrokes(["u"]).await;
+    cx.assert_shared_state("hello cˇo world").await;
+    cx.simulate_shared_keystrokes(["u"]).await;
+    cx.assert_shared_state("hello ˇquoel world").await;
+
+    cx.set_shared_state("hello quˇoel world").await;
+    cx.simulate_shared_keystrokes(["v", "i", "w", "~", "u"])
+        .await;
+    cx.assert_shared_state("hello ˇquoel world").await;
+
+    cx.set_shared_state("\nhello quˇoel world\n").await;
+    cx.simulate_shared_keystrokes(["shift-v", "s", "c", "escape", "u"])
+        .await;
+    cx.assert_shared_state("\nˇhello quoel world\n").await;
+
+    cx.set_shared_state(indoc! {"
+        ˇ1
+        2
+        3"})
+        .await;
+
+    cx.simulate_shared_keystrokes(["ctrl-v", "shift-g", "ctrl-a"])
+        .await;
+    cx.assert_shared_state(indoc! {"
+        ˇ2
+        3
+        4"})
+        .await;
+
+    cx.simulate_shared_keystrokes(["u"]).await;
+    cx.assert_shared_state(indoc! {"
+        ˇ1
+        2
+        3"})
+        .await;
+}

crates/vim/src/vim.rs

@@ -25,7 +25,7 @@ use gpui::{
     actions, impl_actions, Action, AppContext, EntityId, Global, KeystrokeEvent, Subscription,
     View, ViewContext, WeakView, WindowContext,
 };
-use language::{CursorShape, Point, Selection, SelectionGoal};
+use language::{CursorShape, Point, Selection, SelectionGoal, TransactionId};
 pub use mode_indicator::ModeIndicator;
 use motion::Motion;
 use normal::normal_replace;
@@ -209,9 +209,11 @@ impl Vim {
             EditorEvent::SelectionsChanged { local: true } => {
                 let editor = editor.read(cx);
                 if editor.leader_peer_id().is_none() {
-                    let newest = editor.selections.newest::<usize>(cx);
+                    let newest = editor.selections.newest_anchor().clone();
                     let is_multicursor = editor.selections.count() > 1;
-                    local_selections_changed(newest, is_multicursor, cx);
+                    Vim::update(cx, |vim, cx| {
+                        vim.local_selections_changed(newest, is_multicursor, cx);
+                    })
                 }
             }
             EditorEvent::InputIgnored { text } => {
@@ -222,6 +224,12 @@ impl Vim {
                 text,
                 utf16_range_to_replace: range_to_replace,
             } => Vim::record_insertion(text, range_to_replace.clone(), cx),
+            EditorEvent::TransactionBegun { transaction_id } => Vim::update(cx, |vim, cx| {
+                vim.transaction_begun(*transaction_id, cx);
+            }),
+            EditorEvent::TransactionUndone { transaction_id } => Vim::update(cx, |vim, cx| {
+                vim.transaction_undone(transaction_id, cx);
+            }),
             _ => {}
         }));
 
@@ -350,10 +358,13 @@ impl Vim {
         let state = self.state();
         let last_mode = state.mode;
         let prior_mode = state.last_mode;
+        let prior_tx = state.current_tx;
         self.update_state(|state| {
             state.last_mode = last_mode;
             state.mode = mode;
             state.operator_stack.clear();
+            state.current_tx.take();
+            state.current_anchor.take();
         });
         if mode != Mode::Insert {
             self.take_count(cx);
@@ -372,6 +383,11 @@ impl Vim {
             {
                 visual_block_motion(true, editor, cx, |_, point, goal| Some((point, goal)))
             }
+            if last_mode == Mode::Insert {
+                if let Some(prior_tx) = prior_tx {
+                    editor.group_until_transaction(prior_tx, cx)
+                }
+            }
 
             editor.change_selections(None, cx, |s| {
                 // we cheat with visual block mode and use multiple cursors.
@@ -472,6 +488,83 @@ impl Vim {
         self.state().operator_stack.last().copied()
     }
 
+    fn transaction_begun(&mut self, transaction_id: TransactionId, _: &mut WindowContext) {
+        self.update_state(|state| {
+            let mode = if (state.mode == Mode::Insert || state.mode == Mode::Normal)
+                && state.current_tx.is_none()
+            {
+                state.current_tx = Some(transaction_id);
+                state.last_mode
+            } else {
+                state.mode
+            };
+            if mode == Mode::VisualLine || mode == Mode::VisualBlock {
+                state.undo_modes.insert(transaction_id, mode);
+            }
+        });
+    }
+
+    fn transaction_undone(&mut self, transaction_id: &TransactionId, cx: &mut WindowContext) {
+        if !self.state().mode.is_visual() {
+            return;
+        };
+        self.update_active_editor(cx, |vim, editor, cx| {
+            let original_mode = vim.state().undo_modes.get(transaction_id);
+            editor.change_selections(None, cx, |s| match original_mode {
+                Some(Mode::VisualLine) => {
+                    s.move_with(|map, selection| {
+                        selection.collapse_to(
+                            map.prev_line_boundary(selection.start.to_point(map)).1,
+                            SelectionGoal::None,
+                        )
+                    });
+                }
+                Some(Mode::VisualBlock) => {
+                    let mut first = s.first_anchor();
+                    first.collapse_to(first.start, first.goal);
+                    s.select_anchors(vec![first]);
+                }
+                _ => {
+                    s.move_with(|_, selection| {
+                        selection.collapse_to(selection.start, selection.goal);
+                    });
+                }
+            });
+        });
+        self.switch_mode(Mode::Normal, true, cx)
+    }
+
+    fn local_selections_changed(
+        &mut self,
+        newest: Selection<editor::Anchor>,
+        is_multicursor: bool,
+        cx: &mut WindowContext,
+    ) {
+        let state = self.state();
+        if state.mode == Mode::Insert && state.current_tx.is_some() {
+            if state.current_anchor.is_none() {
+                self.update_state(|state| state.current_anchor = Some(newest));
+            } else if state.current_anchor.as_ref().unwrap() != &newest {
+                if let Some(tx_id) = self.update_state(|state| state.current_tx.take()) {
+                    self.update_active_editor(cx, |_, editor, cx| {
+                        editor.group_until_transaction(tx_id, cx)
+                    });
+                }
+            }
+        } else if state.mode == Mode::Normal && newest.start != newest.end {
+            if matches!(newest.goal, SelectionGoal::HorizontalRange { .. }) {
+                self.switch_mode(Mode::VisualBlock, false, cx);
+            } else {
+                self.switch_mode(Mode::Visual, false, cx)
+            }
+        } else if newest.start == newest.end
+            && !is_multicursor
+            && [Mode::Visual, Mode::VisualLine, Mode::VisualBlock].contains(&state.mode)
+        {
+            self.switch_mode(Mode::Normal, true, cx)
+        }
+    }
+
     fn active_editor_input_ignored(text: Arc<str>, cx: &mut WindowContext) {
         if text.is_empty() {
             return;
@@ -667,24 +760,3 @@ impl Settings for VimSettings {
         Self::load_via_json_merge(default_value, user_values)
     }
 }
-
-fn local_selections_changed(
-    newest: Selection<usize>,
-    is_multicursor: bool,
-    cx: &mut WindowContext,
-) {
-    Vim::update(cx, |vim, cx| {
-        if vim.state().mode == Mode::Normal && !newest.is_empty() {
-            if matches!(newest.goal, SelectionGoal::HorizontalRange { .. }) {
-                vim.switch_mode(Mode::VisualBlock, false, cx);
-            } else {
-                vim.switch_mode(Mode::Visual, false, cx)
-            }
-        } else if newest.is_empty()
-            && !is_multicursor
-            && [Mode::Visual, Mode::VisualLine, Mode::VisualBlock].contains(&vim.state().mode)
-        {
-            vim.switch_mode(Mode::Normal, true, cx)
-        }
-    })
-}

crates/vim/test_data/test_undo.json

@@ -0,0 +1,45 @@
+{"Put":{"state":"hello quˇoel world"}}
+{"Key":"v"}
+{"Key":"i"}
+{"Key":"w"}
+{"Key":"s"}
+{"Key":"c"}
+{"Key":"o"}
+{"Key":"escape"}
+{"Key":"u"}
+{"Get":{"state":"hello ˇquoel world","mode":"Normal"}}
+{"Key":"ctrl-r"}
+{"Get":{"state":"hello ˇco world","mode":"Normal"}}
+{"Key":"a"}
+{"Key":"o"}
+{"Key":"right"}
+{"Key":"l"}
+{"Key":"escape"}
+{"Get":{"state":"hello cooˇl world","mode":"Normal"}}
+{"Key":"u"}
+{"Get":{"state":"hello cooˇ world","mode":"Normal"}}
+{"Key":"u"}
+{"Get":{"state":"hello cˇo world","mode":"Normal"}}
+{"Key":"u"}
+{"Get":{"state":"hello ˇquoel world","mode":"Normal"}}
+{"Put":{"state":"hello quˇoel world"}}
+{"Key":"v"}
+{"Key":"i"}
+{"Key":"w"}
+{"Key":"~"}
+{"Key":"u"}
+{"Get":{"state":"hello ˇquoel world","mode":"Normal"}}
+{"Put":{"state":"\nhello quˇoel world\n"}}
+{"Key":"shift-v"}
+{"Key":"s"}
+{"Key":"c"}
+{"Key":"escape"}
+{"Key":"u"}
+{"Get":{"state":"\nˇhello quoel world\n","mode":"Normal"}}
+{"Put":{"state":"ˇ1\n2\n3"}}
+{"Key":"ctrl-v"}
+{"Key":"shift-g"}
+{"Key":"ctrl-a"}
+{"Get":{"state":"ˇ2\n3\n4","mode":"Normal"}}
+{"Key":"u"}
+{"Get":{"state":"ˇ1\n2\n3","mode":"Normal"}}