From 43472a6f02c8f14e91b53e781a9ceaa8ba9af9cb Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Mon, 25 May 2026 18:18:55 -0300 Subject: [PATCH] feat(tui): add vim command counts --- codex-rs/tui/src/bottom_pane/textarea.rs | 310 +++++++++++++----- codex-rs/tui/src/bottom_pane/textarea/vim.rs | 1 + .../tui/src/bottom_pane/textarea/vim/count.rs | 165 ++++++++++ .../tui/src/bottom_pane/textarea/vim/find.rs | 38 ++- 4 files changed, 422 insertions(+), 92 deletions(-) create mode 100644 codex-rs/tui/src/bottom_pane/textarea/vim/count.rs diff --git a/codex-rs/tui/src/bottom_pane/textarea.rs b/codex-rs/tui/src/bottom_pane/textarea.rs index ff7207da2e..66f8e3f583 100644 --- a/codex-rs/tui/src/bottom_pane/textarea.rs +++ b/codex-rs/tui/src/bottom_pane/textarea.rs @@ -111,6 +111,8 @@ pub(crate) struct TextArea { vim_enabled: bool, vim_mode: VimMode, vim_pending: VimPending, + vim_count: Option, + vim_operator_count: Option, vim_last_find: Option, editor_keymap: EditorKeymap, vim_normal_keymap: VimNormalKeymap, @@ -153,6 +155,8 @@ impl TextArea { vim_enabled: false, vim_mode: VimMode::Insert, vim_pending: VimPending::None, + vim_count: None, + vim_operator_count: None, vim_last_find: None, editor_keymap: defaults.editor, vim_normal_keymap: defaults.vim_normal, @@ -234,6 +238,7 @@ impl TextArea { pub(crate) fn set_vim_enabled(&mut self, enabled: bool) { self.vim_enabled = enabled; self.vim_pending = VimPending::None; + self.clear_vim_counts(); self.vim_mode = if enabled { VimMode::Normal } else { @@ -281,6 +286,7 @@ impl TextArea { if self.vim_enabled { self.vim_mode = VimMode::Insert; self.vim_pending = VimPending::None; + self.clear_vim_counts(); } } @@ -294,6 +300,7 @@ impl TextArea { if self.vim_enabled { self.vim_mode = VimMode::Normal; self.vim_pending = VimPending::None; + self.clear_vim_counts(); self.preferred_col = None; } } @@ -645,6 +652,11 @@ impl TextArea { } fn handle_vim_normal(&mut self, event: KeyEvent) { + if matches!(self.vim_pending, VimPending::None | VimPending::Operator(_)) + && self.capture_vim_count_digit(event) + { + return; + } let pending = std::mem::replace(&mut self.vim_pending, VimPending::None); match pending { VimPending::None => {} @@ -663,26 +675,31 @@ impl TextArea { } if self.vim_normal_keymap.enter_insert.is_pressed(event) { + self.clear_vim_counts(); self.vim_mode = VimMode::Insert; return; } if self.vim_normal_keymap.append_after_cursor.is_pressed(event) { + self.clear_vim_counts(); let next = self.next_atomic_boundary(self.cursor_pos); self.set_cursor(next); self.vim_mode = VimMode::Insert; return; } if self.vim_normal_keymap.append_line_end.is_pressed(event) { + self.clear_vim_counts(); self.set_cursor(self.end_of_current_line()); self.vim_mode = VimMode::Insert; return; } if self.vim_normal_keymap.insert_line_start.is_pressed(event) { + self.clear_vim_counts(); self.set_cursor(self.first_non_blank_of_current_line()); self.vim_mode = VimMode::Insert; return; } if self.vim_normal_keymap.open_line_below.is_pressed(event) { + self.clear_vim_counts(); let eol = self.end_of_current_line(); let insert_at = if eol < self.text.len() { eol + 1 } else { eol }; self.insert_str_at(insert_at, "\n"); @@ -696,6 +713,7 @@ impl TextArea { return; } if self.vim_normal_keymap.open_line_above.is_pressed(event) { + self.clear_vim_counts(); let bol = self.beginning_of_current_line(); self.insert_str_at(bol, "\n"); self.set_cursor(bol); @@ -703,39 +721,48 @@ impl TextArea { return; } if self.vim_normal_keymap.move_left.is_pressed(event) { - self.move_cursor_left(); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::Left, count); return; } if self.vim_normal_keymap.move_right.is_pressed(event) { - self.move_cursor_right(); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::Right, count); return; } if self.vim_normal_keymap.move_down.is_pressed(event) { - self.move_cursor_down(); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::Down, count); return; } if self.vim_normal_keymap.move_up.is_pressed(event) { - self.move_cursor_up(); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::Up, count); return; } if self.vim_normal_keymap.move_word_forward.is_pressed(event) { - self.set_cursor(self.beginning_of_next_word()); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::WordForward, count); return; } if self.vim_normal_keymap.move_word_backward.is_pressed(event) { - self.set_cursor(self.beginning_of_previous_word()); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::WordBackward, count); return; } if self.vim_normal_keymap.move_word_end.is_pressed(event) { - self.set_cursor(self.vim_word_end_cursor()); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::WordEnd, count); return; } if self.vim_normal_keymap.move_line_start.is_pressed(event) { - self.set_cursor(self.beginning_of_current_line()); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::LineStart, count); return; } if self.vim_normal_keymap.move_line_end.is_pressed(event) { - self.set_cursor(self.vim_line_end_cursor()); + let count = self.take_vim_count(); + self.apply_counted_vim_normal_motion(VimMotion::LineEnd, count); return; } if let Some(kind) = self.vim_normal_find_kind_for_event(event) { @@ -746,40 +773,49 @@ impl TextArea { return; } if self.vim_normal_keymap.repeat_find.is_pressed(event) { - self.repeat_vim_find(/*operator*/ None, /*reverse*/ false); + let count = self.take_vim_count(); + self.repeat_vim_find(/*operator*/ None, /*reverse*/ false, count); return; } if self.vim_normal_keymap.repeat_find_reverse.is_pressed(event) { - self.repeat_vim_find(/*operator*/ None, /*reverse*/ true); + let count = self.take_vim_count(); + self.repeat_vim_find(/*operator*/ None, /*reverse*/ true, count); return; } if self.vim_normal_keymap.delete_char.is_pressed(event) { - self.delete_forward_kill(/*n*/ 1); + let count = self.take_vim_count(); + self.delete_forward_kill(count); return; } if self.vim_normal_keymap.delete_to_line_end.is_pressed(event) { - self.vim_kill_to_end_of_line(); + let count = self.take_vim_count(); + self.apply_vim_operator(VimOperator::Delete, VimMotion::LineEnd, count); return; } if self.vim_normal_keymap.change_to_line_end.is_pressed(event) { - self.apply_vim_operator(VimOperator::Change, VimMotion::LineEnd); + let count = self.take_vim_count(); + self.apply_vim_operator(VimOperator::Change, VimMotion::LineEnd, count); return; } if self.vim_normal_keymap.substitute_char.is_pressed(event) { - self.delete_forward_kill(/*n*/ 1); + let count = self.take_vim_count(); + self.delete_forward_kill(count); self.vim_mode = VimMode::Insert; return; } if self.vim_normal_keymap.substitute_line.is_pressed(event) { - self.change_current_line(); + let count = self.take_vim_count(); + self.vim_change_current_lines(count); return; } if self.vim_normal_keymap.yank_line.is_pressed(event) { - self.yank_current_line(); + let count = self.take_vim_count(); + self.vim_yank_current_lines(count); return; } if self.vim_normal_keymap.paste_after.is_pressed(event) { - self.paste_after_cursor(); + let count = self.take_vim_count(); + self.paste_after_cursor_counted(count); return; } if self @@ -787,11 +823,11 @@ impl TextArea { .start_delete_operator .is_pressed(event) { - self.vim_pending = VimPending::Operator(VimOperator::Delete); + self.begin_vim_operator(VimOperator::Delete); return; } if self.vim_normal_keymap.start_yank_operator.is_pressed(event) { - self.vim_pending = VimPending::Operator(VimOperator::Yank); + self.begin_vim_operator(VimOperator::Yank); return; } if self @@ -799,28 +835,35 @@ impl TextArea { .start_change_operator .is_pressed(event) { - self.vim_pending = VimPending::Operator(VimOperator::Change); + self.begin_vim_operator(VimOperator::Change); return; } if self.vim_normal_keymap.cancel_operator.is_pressed(event) { self.vim_pending = VimPending::None; + self.clear_vim_counts(); + return; } + self.clear_vim_counts(); } fn handle_vim_operator(&mut self, op: VimOperator, event: KeyEvent) -> bool { if op == VimOperator::Delete && self.vim_operator_keymap.delete_line.is_pressed(event) { - self.kill_current_line(); + let count = self.take_vim_operator_count(); + self.vim_kill_current_lines(count); return true; } if op == VimOperator::Yank && self.vim_operator_keymap.yank_line.is_pressed(event) { - self.yank_current_line(); + let count = self.take_vim_operator_count(); + self.vim_yank_current_lines(count); return true; } if op == VimOperator::Change && self.vim_operator_keymap.change_line.is_pressed(event) { - self.change_current_line(); + let count = self.take_vim_operator_count(); + self.vim_change_current_lines(count); return true; } if self.vim_operator_keymap.cancel.is_pressed(event) { + self.clear_vim_counts(); return true; } if let Some(scope) = self.vim_text_object_scope_for_event(event) { @@ -838,7 +881,8 @@ impl TextArea { return true; } if self.vim_operator_keymap.repeat_find.is_pressed(event) { - self.repeat_vim_find(Some(op), /*reverse*/ false); + let count = self.take_vim_operator_count(); + self.repeat_vim_find(Some(op), /*reverse*/ false, count); return true; } if self @@ -846,14 +890,17 @@ impl TextArea { .repeat_find_reverse .is_pressed(event) { - self.repeat_vim_find(Some(op), /*reverse*/ true); + let count = self.take_vim_operator_count(); + self.repeat_vim_find(Some(op), /*reverse*/ true, count); return true; } if let Some(motion) = self.vim_motion_for_event(event) { - self.apply_vim_operator(op, motion); + let count = self.take_vim_operator_count(); + self.apply_vim_operator(op, motion, count); return true; } + self.clear_vim_counts(); false } @@ -864,12 +911,15 @@ impl TextArea { event: KeyEvent, ) -> bool { if self.vim_text_object_keymap.cancel.is_pressed(event) { + self.clear_vim_counts(); return true; } let Some(object) = self.vim_text_object_for_event(event) else { + self.clear_vim_counts(); return false; }; - if let Some(range) = self.text_object_range(object, scope) { + let count = self.take_vim_operator_count(); + if let Some(range) = self.counted_text_object_range(object, scope, count) { self.apply_vim_operator_to_range(op, range); } true @@ -914,20 +964,20 @@ impl TextArea { None } - fn apply_vim_operator(&mut self, op: VimOperator, motion: VimMotion) { + fn apply_vim_operator(&mut self, op: VimOperator, motion: VimMotion, count: usize) { match op { VimOperator::Delete => { - if let Some(range) = self.range_for_motion(motion) { + if let Some(range) = self.range_for_motion(motion, count) { self.kill_range(range); } } VimOperator::Yank => { - if let Some(range) = self.range_for_motion(motion) { + if let Some(range) = self.range_for_motion(motion, count) { self.yank_range(range); } } VimOperator::Change => { - if let Some(range) = self.range_for_change_motion(motion) { + if let Some(range) = self.range_for_change_motion(motion, count) { self.kill_range(range); self.vim_mode = VimMode::Insert; } else if motion == VimMotion::LineEnd { @@ -937,8 +987,9 @@ impl TextArea { } } - fn range_for_change_motion(&mut self, motion: VimMotion) -> Option> { - if motion == VimMotion::WordForward + fn range_for_change_motion(&mut self, motion: VimMotion, count: usize) -> Option> { + if count == 1 + && motion == VimMotion::WordForward && self.text[self.cursor_pos..] .chars() .next() @@ -947,7 +998,7 @@ impl TextArea { let end = self.end_of_next_word(); return (self.cursor_pos < end).then_some(self.cursor_pos..end); } - self.range_for_motion(motion) + self.range_for_motion(motion, count) } fn apply_vim_operator_to_range(&mut self, op: VimOperator, range: Range) { @@ -961,12 +1012,12 @@ impl TextArea { } } - fn range_for_motion(&mut self, motion: VimMotion) -> Option> { + fn range_for_motion(&mut self, motion: VimMotion, count: usize) -> Option> { if matches!(motion, VimMotion::Up | VimMotion::Down) { - return self.linewise_range_for_vertical_motion(motion); + return self.linewise_range_for_vertical_motion(motion, count); } let start = self.cursor_pos; - let target = self.target_for_motion(motion); + let target = self.target_for_motion(motion, count); if start == target { return None; } @@ -978,28 +1029,36 @@ impl TextArea { Some(range_start..range_end) } - fn linewise_range_for_vertical_motion(&self, motion: VimMotion) -> Option> { + fn linewise_range_for_vertical_motion( + &self, + motion: VimMotion, + count: usize, + ) -> Option> { let current = self.current_line_range_with_newline(); let range = match motion { VimMotion::Up => { - let start = if current.start == 0 { - current.start - } else { - self.beginning_of_line(current.start.saturating_sub(1)) - }; + let mut start = current.start; + for _ in 0..count { + if start == 0 { + break; + } + start = self.beginning_of_line(start.saturating_sub(1)); + } start..current.end } VimMotion::Down => { - let end = if current.end >= self.text.len() { - current.end - } else { - let next_eol = self.end_of_line(current.end); - if next_eol < self.text.len() { + let mut end = current.end; + for _ in 0..count { + if end >= self.text.len() { + break; + } + let next_eol = self.end_of_line(end); + end = if next_eol < self.text.len() { next_eol + 1 } else { next_eol - } - }; + }; + } current.start..end } VimMotion::Left @@ -1013,19 +1072,28 @@ impl TextArea { (range.start < range.end).then_some(range) } - fn target_for_motion(&mut self, motion: VimMotion) -> usize { + fn target_for_motion(&mut self, motion: VimMotion, count: usize) -> usize { let original_cursor = self.cursor_pos; let original_preferred = self.preferred_col; - match motion { - VimMotion::Left => self.move_cursor_left(), - VimMotion::Right => self.move_cursor_right(), - VimMotion::Up => self.move_cursor_up(), - VimMotion::Down => self.move_cursor_down(), - VimMotion::WordForward => self.set_cursor(self.beginning_of_next_word()), - VimMotion::WordBackward => self.set_cursor(self.beginning_of_previous_word()), - VimMotion::WordEnd => self.set_cursor(self.vim_word_end_exclusive()), - VimMotion::LineStart => self.set_cursor(self.beginning_of_current_line()), - VimMotion::LineEnd => self.set_cursor(self.end_of_current_line()), + if motion == VimMotion::LineEnd { + for _ in 1..count { + self.move_cursor_down(); + } + self.set_cursor(self.end_of_current_line()); + } else { + for _ in 0..count { + match motion { + VimMotion::Left => self.move_cursor_left(), + VimMotion::Right => self.move_cursor_right(), + VimMotion::Up => self.move_cursor_up(), + VimMotion::Down => self.move_cursor_down(), + VimMotion::WordForward => self.set_cursor(self.beginning_of_next_word()), + VimMotion::WordBackward => self.set_cursor(self.beginning_of_previous_word()), + VimMotion::WordEnd => self.set_cursor(self.vim_word_end_exclusive()), + VimMotion::LineStart => self.set_cursor(self.beginning_of_current_line()), + VimMotion::LineEnd => unreachable!("handled before count loop"), + } + } } let target = self.cursor_pos; self.cursor_pos = original_cursor; @@ -1116,13 +1184,6 @@ impl TextArea { } } - fn vim_kill_to_end_of_line(&mut self) { - let eol = self.end_of_current_line(); - if self.cursor_pos < eol { - self.kill_range(self.cursor_pos..eol); - } - } - pub fn kill_to_beginning_of_line(&mut self) { let bol = self.beginning_of_current_line(); let range = if self.cursor_pos == bol { @@ -1232,22 +1293,11 @@ impl TextArea { self.set_cursor(cursor.min(self.text.len())); } - fn yank_current_line(&mut self) { - let range = self.current_line_range_with_newline(); - self.yank_line_range(range); - } - fn kill_current_line(&mut self) { let range = self.current_line_range_with_newline(); self.kill_line_range(range); } - fn change_current_line(&mut self) { - let range = self.beginning_of_current_line()..self.end_of_current_line(); - self.kill_line_range(range); - self.vim_mode = VimMode::Insert; - } - fn current_line_range_with_newline(&self) -> Range { let bol = self.beginning_of_current_line(); let eol = self.end_of_current_line(); @@ -2598,6 +2648,106 @@ mod tests { assert_eq!(t.vim_mode_label(), Some("Insert")); } + #[test] + fn vim_counts_apply_to_normal_motions_and_multiply_operator_counts() { + let mut t = ta_with("one two three four"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + + t.input(KeyEvent::new(KeyCode::Char('3'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('w'), KeyModifiers::NONE)); + assert_eq!(t.cursor(), "one two three ".len()); + + let mut t = ta_with("one two three four five six seven"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('d'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('3'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('w'), KeyModifiers::NONE)); + + assert_eq!(t.text(), "seven"); + assert_eq!(t.kill_buffer, "one two three four five six "); + } + + #[test] + fn vim_count_zero_is_a_digit_only_after_a_prefix_begins() { + let mut t = ta_with("abc\ndef"); + t.set_cursor(/*pos*/ "abc\nde".len()); + t.set_vim_enabled(/*enabled*/ true); + + t.input(KeyEvent::new(KeyCode::Char('0'), KeyModifiers::NONE)); + assert_eq!(t.cursor(), "abc\n".len()); + + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('0'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('l'), KeyModifiers::NONE)); + assert_eq!(t.cursor(), t.text().len()); + + let mut t = ta_with("one\ntwo\nthree"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('$'), KeyModifiers::NONE)); + assert_eq!(t.cursor(), "one\ntw".len()); + } + + #[test] + fn vim_counts_cover_text_objects_and_find_targets() { + let mut t = ta_with("one two three"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('d'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('a'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('w'), KeyModifiers::NONE)); + + assert_eq!(t.text(), "three"); + assert_eq!(t.kill_buffer, "one two "); + + let mut t = ta_with("a,b,c,d"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('d'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('f'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char(','), KeyModifiers::NONE)); + + assert_eq!(t.text(), "c,d"); + assert_eq!(t.kill_buffer, "a,b,"); + } + + #[test] + fn vim_counts_cover_paste_substitution_and_linewise_change() { + let mut t = ta_with("a"); + t.kill_buffer = "!".to_string(); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('3'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('p'), KeyModifiers::NONE)); + assert_eq!(t.text(), "a!!!"); + + let mut t = ta_with("👍@file rest"); + t.add_element_range("👍".len().."👍@file".len()) + .expect("valid element"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('s'), KeyModifiers::NONE)); + assert_eq!(t.text(), " rest"); + assert_eq!(t.kill_buffer, "👍@file"); + assert_eq!(t.vim_mode_label(), Some("Insert")); + + let mut t = ta_with("one\ntwo\nthree"); + t.set_cursor(/*pos*/ 0); + t.set_vim_enabled(/*enabled*/ true); + t.input(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + t.input(KeyEvent::new(KeyCode::Char('S'), KeyModifiers::NONE)); + assert_eq!(t.text(), "\nthree"); + assert_eq!(t.kill_buffer, "one\ntwo"); + assert_eq!(t.vim_mode_label(), Some("Insert")); + } + #[test] fn vim_find_till_motions_repeat_and_stay_on_current_line() { let mut t = ta_with("ab,cd,ef\naz"); diff --git a/codex-rs/tui/src/bottom_pane/textarea/vim.rs b/codex-rs/tui/src/bottom_pane/textarea/vim.rs index 91276a140b..adbcc4e4fe 100644 --- a/codex-rs/tui/src/bottom_pane/textarea/vim.rs +++ b/codex-rs/tui/src/bottom_pane/textarea/vim.rs @@ -4,6 +4,7 @@ use crate::key_hint::KeyBindingListExt; use crossterm::event::KeyEvent; use std::ops::Range; +mod count; mod find; mod prose; mod tag; diff --git a/codex-rs/tui/src/bottom_pane/textarea/vim/count.rs b/codex-rs/tui/src/bottom_pane/textarea/vim/count.rs new file mode 100644 index 0000000000..93f11e8f87 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/textarea/vim/count.rs @@ -0,0 +1,165 @@ +use super::super::TextArea; +use super::VimMotion; +use super::VimOperator; +use super::VimTextObject; +use super::VimTextObjectScope; +use crossterm::event::KeyCode; +use crossterm::event::KeyEvent; +use crossterm::event::KeyModifiers; +use std::ops::Range; + +impl TextArea { + pub(in super::super) fn capture_vim_count_digit(&mut self, event: KeyEvent) -> bool { + let KeyEvent { + code: KeyCode::Char(digit @ '0'..='9'), + modifiers: KeyModifiers::NONE, + .. + } = event + else { + return false; + }; + if digit == '0' && self.vim_count.is_none() { + return false; + } + let digit = digit.to_digit(10).unwrap_or_default() as usize; + self.vim_count = Some( + self.vim_count + .unwrap_or_default() + .saturating_mul(10) + .saturating_add(digit), + ); + true + } + + pub(in super::super) fn begin_vim_operator(&mut self, operator: VimOperator) { + self.vim_operator_count = self.vim_count.take(); + self.vim_pending = super::VimPending::Operator(operator); + } + + pub(in super::super) fn take_vim_count(&mut self) -> usize { + self.vim_count.take().unwrap_or(1).max(1) + } + + pub(in super::super) fn take_vim_operator_count(&mut self) -> usize { + let operator = self.vim_operator_count.take().unwrap_or(1); + operator.saturating_mul(self.take_vim_count()).max(1) + } + + pub(in super::super) fn clear_vim_counts(&mut self) { + self.vim_count = None; + self.vim_operator_count = None; + } + + pub(in super::super) fn apply_counted_vim_normal_motion( + &mut self, + motion: VimMotion, + count: usize, + ) { + if motion == VimMotion::LineEnd { + for _ in 1..count { + self.move_cursor_down(); + } + self.set_cursor(self.vim_line_end_cursor()); + return; + } + for _ in 0..count { + match motion { + VimMotion::Left => self.move_cursor_left(), + VimMotion::Right => self.move_cursor_right(), + VimMotion::Up => self.move_cursor_up(), + VimMotion::Down => self.move_cursor_down(), + VimMotion::WordForward => self.set_cursor(self.beginning_of_next_word()), + VimMotion::WordBackward => self.set_cursor(self.beginning_of_previous_word()), + VimMotion::WordEnd => self.set_cursor(self.vim_word_end_cursor()), + VimMotion::LineStart => self.set_cursor(self.beginning_of_current_line()), + VimMotion::LineEnd => unreachable!("handled before count loop"), + } + } + } + + pub(in super::super) fn counted_text_object_range( + &mut self, + object: VimTextObject, + scope: VimTextObjectScope, + count: usize, + ) -> Option> { + let original_cursor = self.cursor_pos; + let mut aggregate = self.text_object_range(object, scope)?; + for _ in 1..count { + let previous_end = aggregate.end; + let mut probe = previous_end; + let mut next_range = None; + while probe <= self.text.len() { + self.cursor_pos = probe; + if let Some(candidate) = self.text_object_range(object, scope) + && candidate.end > previous_end + { + next_range = Some(candidate); + break; + } + if probe == self.text.len() { + break; + } + let next = self.next_atomic_boundary(probe); + if next <= probe { + break; + } + probe = next; + } + let Some(next_range) = next_range else { + break; + }; + aggregate.start = aggregate.start.min(next_range.start); + aggregate.end = aggregate.end.max(next_range.end); + } + self.cursor_pos = original_cursor; + Some(aggregate) + } + + pub(in super::super) fn vim_yank_current_lines(&mut self, count: usize) { + let range = self.current_line_range_with_count(count); + self.yank_line_range(range); + } + + pub(in super::super) fn vim_kill_current_lines(&mut self, count: usize) { + let range = self.current_line_range_with_count(count); + self.kill_line_range(range); + } + + pub(in super::super) fn vim_change_current_lines(&mut self, count: usize) { + let range = self.current_lines_change_range(count); + self.kill_line_range(range); + self.vim_mode = super::VimMode::Insert; + } + + pub(in super::super) fn paste_after_cursor_counted(&mut self, count: usize) { + for _ in 0..count { + self.paste_after_cursor(); + } + } + + fn current_line_range_with_count(&self, count: usize) -> Range { + let start = self.beginning_of_current_line(); + let mut end = self.current_line_range_with_newline().end; + for _ in 1..count { + if end >= self.text.len() { + break; + } + let eol = self.end_of_line(end); + end = if eol < self.text.len() { eol + 1 } else { eol }; + } + start..end + } + + fn current_lines_change_range(&self, count: usize) -> Range { + let start = self.beginning_of_current_line(); + let mut end = self.end_of_current_line(); + for _ in 1..count { + if end >= self.text.len() { + break; + } + end = self.end_of_line(end + 1); + } + start..end + } +} diff --git a/codex-rs/tui/src/bottom_pane/textarea/vim/find.rs b/codex-rs/tui/src/bottom_pane/textarea/vim/find.rs index da5ee12881..3d33075a05 100644 --- a/codex-rs/tui/src/bottom_pane/textarea/vim/find.rs +++ b/codex-rs/tui/src/bottom_pane/textarea/vim/find.rs @@ -28,14 +28,21 @@ impl TextArea { if self.vim_normal_keymap.cancel_operator.is_pressed(event) || self.vim_operator_keymap.cancel.is_pressed(event) { + self.clear_vim_counts(); return true; } let KeyCode::Char(target) = event.code else { + self.clear_vim_counts(); return true; }; let find = VimFind { kind, target }; self.vim_last_find = Some(find); - self.execute_vim_find(operator, find); + let count = if operator.is_some() { + self.take_vim_operator_count() + } else { + self.take_vim_count() + }; + self.execute_vim_find(operator, find, count); true } @@ -81,6 +88,7 @@ impl TextArea { &mut self, operator: Option, reverse: bool, + count: usize, ) { let Some(mut find) = self.vim_last_find else { return; @@ -88,21 +96,21 @@ impl TextArea { if reverse { find.kind = find.kind.reverse(); } - self.execute_vim_find(operator, find); + self.execute_vim_find(operator, find, count); } - fn execute_vim_find(&mut self, operator: Option, find: VimFind) { + fn execute_vim_find(&mut self, operator: Option, find: VimFind, count: usize) { if let Some(operator) = operator { - if let Some(range) = self.range_for_find(find) { + if let Some(range) = self.range_for_find(find, count) { self.apply_vim_operator_to_range(operator, range); } - } else if let Some(target) = self.target_for_find(find) { + } else if let Some(target) = self.target_for_find(find, count) { self.set_cursor(target); } } - pub(in super::super) fn target_for_find(&self, find: VimFind) -> Option { - let matched = self.find_match(find)?; + pub(in super::super) fn target_for_find(&self, find: VimFind, count: usize) -> Option { + let matched = self.find_match(find, count)?; Some(match find.kind { VimFindKind::FindForward | VimFindKind::FindBackward => matched, VimFindKind::TillForward => self.prev_atomic_boundary(matched), @@ -110,8 +118,12 @@ impl TextArea { }) } - pub(in super::super) fn range_for_find(&self, find: VimFind) -> Option> { - let matched = self.find_match(find)?; + pub(in super::super) fn range_for_find( + &self, + find: VimFind, + count: usize, + ) -> Option> { + let matched = self.find_match(find, count)?; let range = match find.kind { VimFindKind::FindForward => self.cursor_pos..self.next_atomic_boundary(matched), VimFindKind::TillForward => self.cursor_pos..matched, @@ -121,7 +133,7 @@ impl TextArea { (range.start < range.end).then_some(range) } - fn find_match(&self, find: VimFind) -> Option { + fn find_match(&self, find: VimFind, count: usize) -> Option { let line_start = self.beginning_of_current_line(); let line_end = self.end_of_current_line(); match find.kind { @@ -130,14 +142,16 @@ impl TextArea { self.text[start..line_end] .char_indices() .map(|(offset, _)| start + offset) - .find(|&idx| self.matches_find_target(idx, find.target)) + .filter(|&idx| self.matches_find_target(idx, find.target)) + .nth(count.saturating_sub(1)) } VimFindKind::FindBackward | VimFindKind::TillBackward => self.text [line_start..self.cursor_pos] .char_indices() .map(|(offset, _)| line_start + offset) .rev() - .find(|&idx| self.matches_find_target(idx, find.target)), + .filter(|&idx| self.matches_find_target(idx, find.target)) + .nth(count.saturating_sub(1)), } }