From 9723ca95e3e959d44d347b449fdbbfd4690d273f Mon Sep 17 00:00:00 2001
From: Wind <WindSoilder@outlook.com>
Date: Fri, 1 Mar 2024 21:00:31 +0800
Subject: [PATCH] Vim mode: make `motion::EndOfLine` works with times. (#8591)

Release Notes:

- Fixed `$` in Vim mode not taking a numeric argument (i.e. `2$` or
`4$`) ([#8007](https://github.com/zed-industries/zed/issues/8007)).

---------

Co-authored-by: Thorsten Ball <mrnugget@gmail.com>
---
 crates/vim/src/motion.rs                      | 15 ++++++++----
 crates/vim/src/normal.rs                      | 23 +++++++++++++++++--
 crates/vim/src/test.rs                        | 18 +++++++++++++++
 .../test_end_of_line_with_neovim.json         |  9 ++++++++
 4 files changed, 58 insertions(+), 7 deletions(-)
 create mode 100644 crates/vim/test_data/test_end_of_line_with_neovim.json

diff --git a/crates/vim/src/motion.rs b/crates/vim/src/motion.rs
index 1747daf813..177fbe4294 100644
--- a/crates/vim/src/motion.rs
+++ b/crates/vim/src/motion.rs
@@ -492,9 +492,10 @@ impl Motion {
                 start_of_line(map, *display_lines, point),
                 SelectionGoal::None,
             ),
-            EndOfLine { display_lines } => {
-                (end_of_line(map, *display_lines, point), SelectionGoal::None)
-            }
+            EndOfLine { display_lines } => (
+                end_of_line(map, *display_lines, point, times),
+                SelectionGoal::None,
+            ),
             StartOfParagraph => (
                 movement::start_of_paragraph(map, point, times),
                 SelectionGoal::None,
@@ -949,8 +950,12 @@ pub(crate) fn start_of_line(
 pub(crate) fn end_of_line(
     map: &DisplaySnapshot,
     display_lines: bool,
-    point: DisplayPoint,
+    mut point: DisplayPoint,
+    times: usize,
 ) -> DisplayPoint {
+    if times > 1 {
+        point = start_of_relative_buffer_row(map, point, times as isize - 1);
+    }
     if display_lines {
         map.clip_point(
             DisplayPoint::new(point.row(), map.line_len(point.row())),
@@ -1119,7 +1124,7 @@ pub(crate) fn next_line_end(
     if times > 1 {
         point = start_of_relative_buffer_row(map, point, times as isize - 1);
     }
-    end_of_line(map, false, point)
+    end_of_line(map, false, point, 1)
 }
 
 fn window_top(
diff --git a/crates/vim/src/normal.rs b/crates/vim/src/normal.rs
index b7d6a0c828..380cbcb393 100644
--- a/crates/vim/src/normal.rs
+++ b/crates/vim/src/normal.rs
@@ -301,7 +301,7 @@ fn insert_line_above(_: &mut Workspace, _: &InsertLineAbove, cx: &mut ViewContex
                 editor.change_selections(Some(Autoscroll::fit()), cx, |s| {
                     s.move_cursors_with(|map, cursor, _| {
                         let previous_line = motion::start_of_relative_buffer_row(map, cursor, -1);
-                        let insert_point = motion::end_of_line(map, false, previous_line);
+                        let insert_point = motion::end_of_line(map, false, previous_line, 1);
                         (insert_point, SelectionGoal::None)
                     });
                 });
@@ -326,7 +326,8 @@ fn insert_line_below(_: &mut Workspace, _: &InsertLineBelow, cx: &mut ViewContex
                 let edits = selection_end_rows.into_iter().map(|row| {
                     let (indent, _) = map.line_indent(row);
                     let end_of_line =
-                        motion::end_of_line(&map, false, DisplayPoint::new(row, 0)).to_point(&map);
+                        motion::end_of_line(&map, false, DisplayPoint::new(row, 0), 1)
+                            .to_point(&map);
 
                     let mut new_text = "\n".to_string();
                     new_text.push_str(&" ".repeat(indent as usize));
@@ -1026,4 +1027,22 @@ mod test {
             .await;
         cx.assert_all("let result = curried_funˇ(ˇ)ˇ(ˇ)ˇ;").await;
     }
+
+    #[gpui::test]
+    async fn test_end_of_line_with_neovim(cx: &mut gpui::TestAppContext) {
+        let mut cx = NeovimBackedTestContext::new(cx).await;
+
+        // goes to current line end
+        cx.set_shared_state(indoc! {"ˇaa\nbb\ncc"}).await;
+        cx.simulate_shared_keystrokes(["$"]).await;
+        cx.assert_shared_state(indoc! {"aˇa\nbb\ncc"}).await;
+
+        // goes to next line end
+        cx.simulate_shared_keystrokes(["2", "$"]).await;
+        cx.assert_shared_state("aa\nbˇb\ncc").await;
+
+        // try to exceed the final line.
+        cx.simulate_shared_keystrokes(["4", "$"]).await;
+        cx.assert_shared_state("aa\nbb\ncˇc").await;
+    }
 }
diff --git a/crates/vim/src/test.rs b/crates/vim/src/test.rs
index b1ed6b0389..8628be7298 100644
--- a/crates/vim/src/test.rs
+++ b/crates/vim/src/test.rs
@@ -152,6 +152,24 @@ async fn test_end_of_document_710(cx: &mut gpui::TestAppContext) {
     cx.assert_editor_state("aˇa\nbb\ncc");
 }
 
+#[gpui::test]
+async fn test_end_of_line_with_times(cx: &mut gpui::TestAppContext) {
+    let mut cx = VimTestContext::new(cx, true).await;
+
+    // goes to current line end
+    cx.set_state(indoc! {"ˇaa\nbb\ncc"}, Mode::Normal);
+    cx.simulate_keystrokes(["$"]);
+    cx.assert_editor_state("aˇa\nbb\ncc");
+
+    // goes to next line end
+    cx.simulate_keystrokes(["2", "$"]);
+    cx.assert_editor_state("aa\nbˇb\ncc");
+
+    // try to exceed the final line.
+    cx.simulate_keystrokes(["4", "$"]);
+    cx.assert_editor_state("aa\nbb\ncˇc");
+}
+
 #[gpui::test]
 async fn test_indent_outdent(cx: &mut gpui::TestAppContext) {
     let mut cx = VimTestContext::new(cx, true).await;
diff --git a/crates/vim/test_data/test_end_of_line_with_neovim.json b/crates/vim/test_data/test_end_of_line_with_neovim.json
new file mode 100644
index 0000000000..58ac05134a
--- /dev/null
+++ b/crates/vim/test_data/test_end_of_line_with_neovim.json
@@ -0,0 +1,9 @@
+{"Put":{"state":"ˇaa\nbb\ncc"}}
+{"Key":"$"}
+{"Get":{"state":"aˇa\nbb\ncc","mode":"Normal"}}
+{"Key":"2"}
+{"Key":"$"}
+{"Get":{"state":"aa\nbˇb\ncc","mode":"Normal"}}
+{"Key":"4"}
+{"Key":"$"}
+{"Get":{"state":"aa\nbb\ncˇc","mode":"Normal"}}