vim: Fix count handling to allow pre/post counts

Fixes 2yy, d3d, etc.

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

assets/keymaps/vim.json

@@ -328,7 +328,7 @@
     }
   },
   {
-    "context": "Editor && vim_mode == normal && (vim_operator == none || vim_operator == n) && !VimWaiting",
+    "context": "Editor && vim_mode == normal && vim_operator == none && !VimWaiting",
     "bindings": {
       ".": "vim::Repeat",
       "c": [
@@ -391,7 +391,7 @@
     }
   },
   {
-    "context": "Editor && vim_operator == n",
+    "context": "Editor && VimCount",
     "bindings": {
       "0": [
         "vim::Number",

crates/vim/src/motion.rs

@@ -229,11 +229,11 @@ pub(crate) fn motion(motion: Motion, cx: &mut WindowContext) {
         Vim::update(cx, |vim, cx| vim.pop_operator(cx));
     }
 
-    let times = Vim::update(cx, |vim, cx| vim.pop_number_operator(cx));
+    let count = Vim::update(cx, |vim, _| vim.take_count());
     let operator = Vim::read(cx).active_operator();
     match Vim::read(cx).state().mode {
-        Mode::Normal => normal_motion(motion, operator, times, cx),
-        Mode::Visual | Mode::VisualLine | Mode::VisualBlock => visual_motion(motion, times, cx),
+        Mode::Normal => normal_motion(motion, operator, count, cx),
+        Mode::Visual | Mode::VisualLine | Mode::VisualBlock => visual_motion(motion, count, cx),
         Mode::Insert => {
             // Shouldn't execute a motion in insert mode. Ignoring
         }
@@ -412,7 +412,7 @@ impl Motion {
                 map.clip_at_line_end(movement::end_of_paragraph(map, point, times)),
                 SelectionGoal::None,
             ),
-            CurrentLine => (next_line_end(map, point, 1), SelectionGoal::None),
+            CurrentLine => (next_line_end(map, point, times), SelectionGoal::None),
             StartOfDocument => (start_of_document(map, point, times), SelectionGoal::None),
             EndOfDocument => (
                 end_of_document(map, point, maybe_times),

crates/vim/src/normal.rs

@@ -68,21 +68,21 @@ pub fn init(cx: &mut AppContext) {
     cx.add_action(|_: &mut Workspace, _: &DeleteLeft, cx| {
         Vim::update(cx, |vim, cx| {
             vim.record_current_action(cx);
-            let times = vim.pop_number_operator(cx);
+            let times = vim.take_count();
             delete_motion(vim, Motion::Left, times, cx);
         })
     });
     cx.add_action(|_: &mut Workspace, _: &DeleteRight, cx| {
         Vim::update(cx, |vim, cx| {
             vim.record_current_action(cx);
-            let times = vim.pop_number_operator(cx);
+            let times = vim.take_count();
             delete_motion(vim, Motion::Right, times, cx);
         })
     });
     cx.add_action(|_: &mut Workspace, _: &ChangeToEndOfLine, cx| {
         Vim::update(cx, |vim, cx| {
             vim.start_recording(cx);
-            let times = vim.pop_number_operator(cx);
+            let times = vim.take_count();
             change_motion(
                 vim,
                 Motion::EndOfLine {
@@ -96,7 +96,7 @@ pub fn init(cx: &mut AppContext) {
     cx.add_action(|_: &mut Workspace, _: &DeleteToEndOfLine, cx| {
         Vim::update(cx, |vim, cx| {
             vim.record_current_action(cx);
-            let times = vim.pop_number_operator(cx);
+            let times = vim.take_count();
             delete_motion(
                 vim,
                 Motion::EndOfLine {
@@ -110,7 +110,7 @@ pub fn init(cx: &mut AppContext) {
     cx.add_action(|_: &mut Workspace, _: &JoinLines, cx| {
         Vim::update(cx, |vim, cx| {
             vim.record_current_action(cx);
-            let mut times = vim.pop_number_operator(cx).unwrap_or(1);
+            let mut times = vim.take_count().unwrap_or(1);
             if vim.state().mode.is_visual() {
                 times = 1;
             } else if times > 1 {

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)
         })
     });

crates/vim/src/state.rs

@@ -33,7 +33,6 @@ impl Default for Mode {
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Deserialize)]
 pub enum Operator {
-    Number(usize),
     Change,
     Delete,
     Yank,
@@ -47,6 +46,12 @@ pub enum Operator {
 pub struct EditorState {
     pub mode: Mode,
     pub last_mode: Mode,
+
+    /// pre_count is the number before an operator is specified (3 in 3d2d)
+    pub pre_count: Option<usize>,
+    /// post_count is the number after an operator is specified (2 in 3d2d)
+    pub post_count: Option<usize>,
+
     pub operator_stack: Vec<Operator>,
 }
 
@@ -158,6 +163,10 @@ impl EditorState {
         }
     }
 
+    pub fn active_operator(&self) -> Option<Operator> {
+        self.operator_stack.last().copied()
+    }
+
     pub fn keymap_context_layer(&self) -> KeymapContext {
         let mut context = KeymapContext::default();
         context.add_identifier("VimEnabled");
@@ -174,7 +183,13 @@ impl EditorState {
             context.add_identifier("VimControl");
         }
 
-        let active_operator = self.operator_stack.last();
+        if self.active_operator().is_none() && self.pre_count.is_some()
+            || self.active_operator().is_some() && self.post_count.is_some()
+        {
+            context.add_identifier("VimCount");
+        }
+
+        let active_operator = self.active_operator();
 
         if let Some(active_operator) = active_operator {
             for context_flag in active_operator.context_flags().into_iter() {
@@ -194,7 +209,6 @@ impl EditorState {
 impl Operator {
     pub fn id(&self) -> &'static str {
         match self {
-            Operator::Number(_) => "n",
             Operator::Object { around: false } => "i",
             Operator::Object { around: true } => "a",
             Operator::Change => "c",

crates/vim/src/vim.rs

@@ -40,9 +40,12 @@ pub struct SwitchMode(pub Mode);
 pub struct PushOperator(pub Operator);
 
 #[derive(Clone, Deserialize, PartialEq)]
-struct Number(u8);
+struct Number(usize);
 
-actions!(vim, [Tab, Enter]);
+actions!(
+    vim,
+    [Tab, Enter, Object, InnerObject, FindForward, FindBackward]
+);
 impl_actions!(vim, [Number, SwitchMode, PushOperator]);
 
 #[derive(Copy, Clone, Debug)]
@@ -70,7 +73,7 @@ pub fn init(cx: &mut AppContext) {
         },
     );
     cx.add_action(|_: &mut Workspace, n: &Number, cx: _| {
-        Vim::update(cx, |vim, cx| vim.push_number(n, cx));
+        Vim::update(cx, |vim, _| vim.push_count_digit(n.0));
     });
 
     cx.add_action(|_: &mut Workspace, _: &Tab, cx| {
@@ -236,12 +239,7 @@ impl Vim {
         if !self.workspace_state.replaying {
             self.workspace_state.recording = true;
             self.workspace_state.recorded_actions = Default::default();
-            self.workspace_state.recorded_count =
-                if let Some(Operator::Number(number)) = self.active_operator() {
-                    Some(number)
-                } else {
-                    None
-                };
+            self.workspace_state.recorded_count = None;
 
             let selections = self
                 .active_editor
@@ -352,6 +350,36 @@ impl Vim {
         });
     }
 
+    fn push_count_digit(&mut self, number: usize) {
+        if self.active_operator().is_some() {
+            self.update_state(|state| {
+                state.post_count = Some(state.post_count.unwrap_or(0) * 10 + number)
+            })
+        } else {
+            self.update_state(|state| {
+                state.pre_count = Some(state.pre_count.unwrap_or(0) * 10 + number)
+            })
+        }
+    }
+
+    fn take_count(&mut self) -> Option<usize> {
+        if self.workspace_state.replaying {
+            return self.workspace_state.recorded_count;
+        }
+
+        let count = if self.state().post_count == None && self.state().pre_count == None {
+            return None;
+        } else {
+            Some(self.update_state(|state| {
+                state.post_count.take().unwrap_or(1) * state.pre_count.take().unwrap_or(1)
+            }))
+        };
+        if self.workspace_state.recording {
+            self.workspace_state.recorded_count = count;
+        }
+        count
+    }
+
     fn push_operator(&mut self, operator: Operator, cx: &mut WindowContext) {
         if matches!(
             operator,
@@ -363,15 +391,6 @@ impl Vim {
         self.sync_vim_settings(cx);
     }
 
-    fn push_number(&mut self, Number(number): &Number, cx: &mut WindowContext) {
-        if let Some(Operator::Number(current_number)) = self.active_operator() {
-            self.pop_operator(cx);
-            self.push_operator(Operator::Number(current_number * 10 + *number as usize), cx);
-        } else {
-            self.push_operator(Operator::Number(*number as usize), cx);
-        }
-    }
-
     fn maybe_pop_operator(&mut self) -> Option<Operator> {
         self.update_state(|state| state.operator_stack.pop())
     }
@@ -382,21 +401,6 @@ impl Vim {
         self.sync_vim_settings(cx);
         popped_operator
     }
-
-    fn pop_number_operator(&mut self, cx: &mut WindowContext) -> Option<usize> {
-        if self.workspace_state.replaying {
-            if let Some(number) = self.workspace_state.recorded_count {
-                return Some(number);
-            }
-        }
-
-        if let Some(Operator::Number(number)) = self.active_operator() {
-            self.pop_operator(cx);
-            return Some(number);
-        }
-        None
-    }
-
     fn clear_operator(&mut self, cx: &mut WindowContext) {
         self.update_state(|state| state.operator_stack.clear());
         self.sync_vim_settings(cx);

crates/vim/test_data/test_delete_with_counts.json

@@ -0,0 +1,16 @@
+{"Put":{"state":"The ˇquick brown\nfox jumps over\nthe lazy dog"}}
+{"Key":"d"}
+{"Key":"2"}
+{"Key":"d"}
+{"Get":{"state":"the ˇlazy dog","mode":"Normal"}}
+{"Put":{"state":"The ˇquick brown\nfox jumps over\nthe lazy dog"}}
+{"Key":"2"}
+{"Key":"d"}
+{"Key":"d"}
+{"Get":{"state":"the ˇlazy dog","mode":"Normal"}}
+{"Put":{"state":"The ˇquick brown\nfox jumps over\nthe moon,\na star, and\nthe lazy dog"}}
+{"Key":"2"}
+{"Key":"d"}
+{"Key":"2"}
+{"Key":"d"}
+{"Get":{"state":"the ˇlazy dog","mode":"Normal"}}

crates/vim/test_data/test_repeat_motion_counts.json

@@ -0,0 +1,13 @@
+{"Put":{"state":"ˇthe quick brown\nfox jumps over\nthe lazy dog"}}
+{"Key":"3"}
+{"Key":"d"}
+{"Key":"3"}
+{"Key":"l"}
+{"Get":{"state":"ˇ brown\nfox jumps over\nthe lazy dog","mode":"Normal"}}
+{"Key":"j"}
+{"Key":"."}
+{"Get":{"state":" brown\nˇ over\nthe lazy dog","mode":"Normal"}}
+{"Key":"j"}
+{"Key":"2"}
+{"Key":"."}
+{"Get":{"state":" brown\n over\nˇe lazy dog","mode":"Normal"}}