vim: Fix count handling to allow pre/post counts

Fixes 2yy, d3d, etc.

For zed-industries/community#970
For zed-industries/community#1496

crates/vim/src/normal/case.rs

@@ -8,7 +8,7 @@ use crate::{normal::ChangeCase, state::Mode, Vim};
 pub fn change_case(_: &mut Workspace, _: &ChangeCase, cx: &mut ViewContext<Workspace>) {
     Vim::update(cx, |vim, cx| {
         vim.record_current_action(cx);
-        let count = vim.pop_number_operator(cx).unwrap_or(1) as u32;
+        let count = vim.take_count().unwrap_or(1) as u32;
         vim.update_active_editor(cx, |editor, cx| {
             let mut ranges = Vec::new();
             let mut cursor_positions = Vec::new();

crates/vim/src/normal/delete.rs

@@ -387,4 +387,40 @@ mod test {
         assert_eq!(cx.active_operator(), None);
         assert_eq!(cx.mode(), Mode::Normal);
     }
+
+    #[gpui::test]
+    async fn test_delete_with_counts(cx: &mut gpui::TestAppContext) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+        cx.set_shared_state(indoc! {"
+                The ˇquick brown
+                fox jumps over
+                the lazy dog"})
+            .await;
+        cx.simulate_shared_keystrokes(["d", "2", "d"]).await;
+        cx.assert_shared_state(indoc! {"
+        the ˇlazy dog"})
+            .await;
+
+        cx.set_shared_state(indoc! {"
+                The ˇquick brown
+                fox jumps over
+                the lazy dog"})
+            .await;
+        cx.simulate_shared_keystrokes(["2", "d", "d"]).await;
+        cx.assert_shared_state(indoc! {"
+        the ˇlazy dog"})
+            .await;
+
+        cx.set_shared_state(indoc! {"
+                The ˇquick brown
+                fox jumps over
+                the moon,
+                a star, and
+                the lazy dog"})
+            .await;
+        cx.simulate_shared_keystrokes(["2", "d", "2", "d"]).await;
+        cx.assert_shared_state(indoc! {"
+        the ˇlazy dog"})
+            .await;
+    }
 }

crates/vim/src/normal/repeat.rs

@@ -34,7 +34,7 @@ pub(crate) fn init(cx: &mut AppContext) {
             let Some(editor) = vim.active_editor.clone() else {
                 return None;
             };
-            let count = vim.pop_number_operator(cx);
+            let count = vim.take_count();
 
             vim.workspace_state.replaying = true;
 
@@ -424,4 +424,42 @@ mod test {
         })
         .await;
     }
+
+    #[gpui::test]
+    async fn test_repeat_motion_counts(
+        deterministic: Arc<Deterministic>,
+        cx: &mut gpui::TestAppContext,
+    ) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+
+        cx.set_shared_state(indoc! {
+            "ˇthe quick brown
+            fox jumps over
+            the lazy dog"
+        })
+        .await;
+        cx.simulate_shared_keystrokes(["3", "d", "3", "l"]).await;
+        cx.assert_shared_state(indoc! {
+            "ˇ brown
+            fox jumps over
+            the lazy dog"
+        })
+        .await;
+        cx.simulate_shared_keystrokes(["j", "."]).await;
+        deterministic.run_until_parked();
+        cx.assert_shared_state(indoc! {
+            " brown
+            ˇ over
+            the lazy dog"
+        })
+        .await;
+        cx.simulate_shared_keystrokes(["j", "2", "."]).await;
+        deterministic.run_until_parked();
+        cx.assert_shared_state(indoc! {
+            " brown
+             over
+            ˇe lazy dog"
+        })
+        .await;
+    }
 }

crates/vim/src/normal/scroll.rs

@@ -48,7 +48,7 @@ pub fn init(cx: &mut AppContext) {
 
 fn scroll(cx: &mut ViewContext<Workspace>, by: fn(c: Option<f32>) -> ScrollAmount) {
     Vim::update(cx, |vim, cx| {
-        let amount = by(vim.pop_number_operator(cx).map(|c| c as f32));
+        let amount = by(vim.take_count().map(|c| c as f32));
         vim.update_active_editor(cx, |editor, cx| scroll_editor(editor, &amount, cx));
     })
 }

crates/vim/src/normal/search.rs

@@ -52,7 +52,7 @@ fn search(workspace: &mut Workspace, action: &Search, cx: &mut ViewContext<Works
         Direction::Next
     };
     Vim::update(cx, |vim, cx| {
-        let count = vim.pop_number_operator(cx).unwrap_or(1);
+        let count = vim.take_count().unwrap_or(1);
         pane.update(cx, |pane, cx| {
             if let Some(search_bar) = pane.toolbar().read(cx).item_of_type::<BufferSearchBar>() {
                 search_bar.update(cx, |search_bar, cx| {
@@ -119,7 +119,7 @@ pub fn move_to_internal(
 ) {
     Vim::update(cx, |vim, cx| {
         let pane = workspace.active_pane().clone();
-        let count = vim.pop_number_operator(cx).unwrap_or(1);
+        let count = vim.take_count().unwrap_or(1);
         pane.update(cx, |pane, cx| {
             if let Some(search_bar) = pane.toolbar().read(cx).item_of_type::<BufferSearchBar>() {
                 let search = search_bar.update(cx, |search_bar, cx| {

crates/vim/src/normal/substitute.rs

@@ -11,7 +11,7 @@ pub(crate) fn init(cx: &mut AppContext) {
     cx.add_action(|_: &mut Workspace, _: &Substitute, cx| {
         Vim::update(cx, |vim, cx| {
             vim.start_recording(cx);
-            let count = vim.pop_number_operator(cx);
+            let count = vim.take_count();
             substitute(vim, count, vim.state().mode == Mode::VisualLine, cx);
         })
     });
@@ -22,7 +22,7 @@ pub(crate) fn init(cx: &mut AppContext) {
             if matches!(vim.state().mode, Mode::VisualBlock | Mode::Visual) {
                 vim.switch_mode(Mode::VisualLine, false, cx)
             }
-            let count = vim.pop_number_operator(cx);
+            let count = vim.take_count();
             substitute(vim, count, true, cx)
         })
     });