vim: Allow count and repeat for "r" and "shift-r" action (#13287)

Fixing the "r" action just involved adapting `normal_replace` to replace
multiple characters.

Fixing the "shift-r" command was less straightforward. The bindings for
`vim::BeforeNormal` in replace mode were being overwritten and several
other steps required for action repetition were not performed. Finally,
the cursor adjustment after re-entering normal mode was duplicated
(`vim::BeforeNormal` was now triggered correctly) so I removed the
special case for replace mode.

Release Notes:

- Fixed vim "r" action to accept a count argument
- Fixed vim "shift-r" action to accept a count argument and allow
repetition

---------

Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>

assets/keymaps/vim.json

@@ -645,11 +645,13 @@
       "escape": "vim::NormalBefore",
       "ctrl-c": "vim::NormalBefore",
       "ctrl-[": "vim::NormalBefore",
+      "tab": "vim::Tab",
+      "enter": "vim::Enter",
       "backspace": "vim::UndoReplace"
     }
   },
   {
-    "context": "Editor && VimWaiting",
+    "context": "Editor && vim_mode != replace && VimWaiting",
     "bindings": {
       "tab": "vim::Tab",
       "enter": "vim::Enter",

crates/vim/src/normal.rs

@@ -484,6 +484,7 @@ fn restore_selection_cursors(
 
 pub(crate) fn normal_replace(text: Arc<str>, cx: &mut WindowContext) {
     Vim::update(cx, |vim, cx| {
+        let count = vim.take_count(cx).unwrap_or(1);
         vim.stop_recording();
         vim.update_active_editor(cx, |_, editor, cx| {
             editor.transact(cx, |editor, cx| {
@@ -506,13 +507,13 @@ pub(crate) fn normal_replace(text: Arc<str>, cx: &mut WindowContext) {
                     .into_iter()
                     .map(|selection| {
                         let mut range = selection.range();
-                        *range.end.column_mut() += 1;
+                        range.end = right(&map, range.end, count);
-                        range.end = map.clip_point(range.end, Bias::Right);
+                        let repeated_text = text.repeat(count);
 
                         (
                             range.start.to_offset(&map, Bias::Left)
                                 ..range.end.to_offset(&map, Bias::Left),
-                            text.clone(),
+                            repeated_text,
                         )
                     })
                     .collect::<Vec<_>>();
@@ -523,6 +524,11 @@ pub(crate) fn normal_replace(text: Arc<str>, cx: &mut WindowContext) {
                 editor.set_clip_at_line_ends(true, cx);
                 editor.change_selections(None, cx, |s| {
                     s.select_anchor_ranges(stable_anchors);
+                    if count > 1 {
+                        s.move_cursors_with(|map, point, _| {
+                            (right(map, point, count - 1), SelectionGoal::None)
+                        });
+                    }
                 });
             });
         });
@@ -1415,4 +1421,25 @@ mod test {
             indoc! {"asserˇt_binding"},
         );
     }
+
+    #[gpui::test]
+    async fn test_r(cx: &mut gpui::TestAppContext) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state("ˇhello\n").await;
+        cx.simulate_shared_keystrokes("r -").await;
+        cx.shared_state().await.assert_eq("ˇ-ello\n");
+
+        cx.set_shared_state("ˇhello\n").await;
+        cx.simulate_shared_keystrokes("3 r -").await;
+        cx.shared_state().await.assert_eq("--ˇ-lo\n");
+
+        cx.set_shared_state("ˇhello\n").await;
+        cx.simulate_shared_keystrokes("r - 2 l .").await;
+        cx.shared_state().await.assert_eq("-eˇ-lo\n");
+
+        cx.set_shared_state("ˇhello world\n").await;
+        cx.simulate_shared_keystrokes("2 r - f w .").await;
+        cx.shared_state().await.assert_eq("--llo -ˇ-rld\n");
+    }
 }

crates/vim/src/normal/repeat.rs

@@ -31,6 +31,8 @@ fn repeatable_insert(action: &ReplayableAction) -> Option<Box<dyn Action>> {
                 || super::InsertLineBelow.partial_eq(&**action)
             {
                 Some(super::InsertLineBelow.boxed_clone())
+            } else if crate::replace::ToggleReplace.partial_eq(&**action) {
+                Some(crate::replace::ToggleReplace.boxed_clone())
             } else {
                 None
             }

crates/vim/src/replace.rs

@@ -16,6 +16,7 @@ pub fn register(workspace: &mut Workspace, _: &mut ViewContext<Workspace>) {
     workspace.register_action(|_, _: &ToggleReplace, cx: &mut ViewContext<Workspace>| {
         Vim::update(cx, |vim, cx| {
             vim.update_state(|state| state.replacements = vec![]);
+            vim.start_recording(cx);
             vim.switch_mode(Mode::Replace, false, cx);
         });
     });
@@ -237,6 +238,30 @@ mod test {
         );
     }
 
+    #[gpui::test]
+    async fn test_replace_mode_with_counts(cx: &mut gpui::TestAppContext) {
+        let mut cx: NeovimBackedTestContext = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state("ˇhello\n").await;
+        cx.simulate_shared_keystrokes("3 shift-r - escape").await;
+        cx.shared_state().await.assert_eq("--ˇ-lo\n");
+
+        cx.set_shared_state("ˇhello\n").await;
+        cx.simulate_shared_keystrokes("3 shift-r a b c escape")
+            .await;
+        cx.shared_state().await.assert_eq("abcabcabˇc\n");
+    }
+
+    #[gpui::test]
+    async fn test_replace_mode_repeat(cx: &mut gpui::TestAppContext) {
+        let mut cx: NeovimBackedTestContext = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state("ˇhello world\n").await;
+        cx.simulate_shared_keystrokes("shift-r - - - escape 4 l .")
+            .await;
+        cx.shared_state().await.assert_eq("---lo --ˇ-ld\n");
+    }
+
     #[gpui::test]
     async fn test_replace_mode_undo(cx: &mut gpui::TestAppContext) {
         let mut cx: NeovimBackedTestContext = NeovimBackedTestContext::new(cx).await;

crates/vim/src/vim.rs

@@ -421,7 +421,7 @@ impl Vim {
             state.current_tx.take();
             state.current_anchor.take();
         });
-        if mode != Mode::Insert {
+        if mode != Mode::Insert && mode != Mode::Replace {
             self.take_count(cx);
         }
 
@@ -488,11 +488,6 @@ impl Vim {
                         if selection.is_empty() {
                             selection.end = movement::right(map, selection.start);
                         }
-                    } else if last_mode == Mode::Replace {
-                        if selection.head().column() != 0 {
-                            let point = movement::left(map, selection.head());
-                            selection.collapse_to(point, selection.goal)
-                        }
                     }
                 });
             })

crates/vim/test_data/test_r.json

@@ -0,0 +1,24 @@
+{"Put":{"state":"ˇhello\n"}}
+{"Key":"r"}
+{"Key":"-"}
+{"Get":{"state":"ˇ-ello\n","mode":"Normal"}}
+{"Put":{"state":"ˇhello\n"}}
+{"Key":"3"}
+{"Key":"r"}
+{"Key":"-"}
+{"Get":{"state":"--ˇ-lo\n","mode":"Normal"}}
+{"Put":{"state":"ˇhello\n"}}
+{"Key":"r"}
+{"Key":"-"}
+{"Key":"2"}
+{"Key":"l"}
+{"Key":"."}
+{"Get":{"state":"-eˇ-lo\n","mode":"Normal"}}
+{"Put":{"state":"ˇhello world\n"}}
+{"Key":"2"}
+{"Key":"r"}
+{"Key":"-"}
+{"Key":"f"}
+{"Key":"w"}
+{"Key":"."}
+{"Get":{"state":"--llo -ˇ-rld\n","mode":"Normal"}}

crates/vim/test_data/test_replace_mode_repeat.json

@@ -0,0 +1,10 @@
+{"Put":{"state":"ˇhello world\n"}}
+{"Key":"shift-r"}
+{"Key":"-"}
+{"Key":"-"}
+{"Key":"-"}
+{"Key":"escape"}
+{"Key":"4"}
+{"Key":"l"}
+{"Key":"."}
+{"Get":{"state":"---lo --ˇ-ld\n","mode":"Normal"}}

crates/vim/test_data/test_replace_mode_with_counts.json

@@ -0,0 +1,14 @@
+{"Put":{"state":"ˇhello\n"}}
+{"Key":"3"}
+{"Key":"shift-r"}
+{"Key":"-"}
+{"Key":"escape"}
+{"Get":{"state":"--ˇ-lo\n","mode":"Normal"}}
+{"Put":{"state":"ˇhello\n"}}
+{"Key":"3"}
+{"Key":"shift-r"}
+{"Key":"a"}
+{"Key":"b"}
+{"Key":"c"}
+{"Key":"escape"}
+{"Get":{"state":"abcabcabˇc\n","mode":"Normal"}}