Reset selection goal after helix motion (#33184)

fantacell created

Closes #33060

Motions like `NextWordStart` don't reset the selection goal in vim mode
`helix_normal` unlike in `normal` which can lead to the cursor jumping
back to the previous horizontal position after going up or down.

Release Notes:

- N/A

Change summary

crates/vim/src/helix.rs | 88 ++++++++++++++++++++++--------------------
1 file changed, 46 insertions(+), 42 deletions(-)

Detailed changes

crates/vim/src/helix.rs 🔗

@@ -2,6 +2,7 @@ use editor::{DisplayPoint, Editor, movement, scroll::Autoscroll};
 use gpui::{Action, actions};
 use gpui::{Context, Window};
 use language::{CharClassifier, CharKind};
+use text::SelectionGoal;
 
 use crate::{Vim, motion::Motion, state::Mode};
 
@@ -49,43 +50,43 @@ impl Vim {
             editor.change_selections(Some(Autoscroll::fit()), window, cx, |s| {
                 s.move_with(|map, selection| {
                     let times = times.unwrap_or(1);
+                    let new_goal = SelectionGoal::None;
+                    let mut head = selection.head();
+                    let mut tail = selection.tail();
 
-                    if selection.head() == map.max_point() {
+                    if head == map.max_point() {
                         return;
                     }
 
                     // collapse to block cursor
-                    if selection.tail() < selection.head() {
-                        selection.set_tail(movement::left(map, selection.head()), selection.goal);
+                    if tail < head {
+                        tail = movement::left(map, head);
                     } else {
-                        selection.set_tail(selection.head(), selection.goal);
-                        selection.set_head(movement::right(map, selection.head()), selection.goal);
+                        tail = head;
+                        head = movement::right(map, head);
                     }
 
                     // create a classifier
-                    let classifier = map
-                        .buffer_snapshot
-                        .char_classifier_at(selection.head().to_point(map));
+                    let classifier = map.buffer_snapshot.char_classifier_at(head.to_point(map));
 
-                    let mut last_selection = selection.clone();
                     for _ in 0..times {
-                        let (new_tail, new_head) =
-                            movement::find_boundary_trail(map, selection.head(), |left, right| {
+                        let (maybe_next_tail, next_head) =
+                            movement::find_boundary_trail(map, head, |left, right| {
                                 is_boundary(left, right, &classifier)
                             });
 
-                        selection.set_head(new_head, selection.goal);
-                        if let Some(new_tail) = new_tail {
-                            selection.set_tail(new_tail, selection.goal);
+                        if next_head == head && maybe_next_tail.unwrap_or(next_head) == tail {
+                            break;
                         }
 
-                        if selection.head() == last_selection.head()
-                            && selection.tail() == last_selection.tail()
-                        {
-                            break;
+                        head = next_head;
+                        if let Some(next_tail) = maybe_next_tail {
+                            tail = next_tail;
                         }
-                        last_selection = selection.clone();
                     }
+
+                    selection.set_tail(tail, new_goal);
+                    selection.set_head(head, new_goal);
                 });
             });
         });
@@ -102,47 +103,50 @@ impl Vim {
             editor.change_selections(Some(Autoscroll::fit()), window, cx, |s| {
                 s.move_with(|map, selection| {
                     let times = times.unwrap_or(1);
+                    let new_goal = SelectionGoal::None;
+                    let mut head = selection.head();
+                    let mut tail = selection.tail();
 
-                    if selection.head() == DisplayPoint::zero() {
+                    if head == DisplayPoint::zero() {
                         return;
                     }
 
                     // collapse to block cursor
-                    if selection.tail() < selection.head() {
-                        selection.set_tail(movement::left(map, selection.head()), selection.goal);
+                    if tail < head {
+                        tail = movement::left(map, head);
                     } else {
-                        selection.set_tail(selection.head(), selection.goal);
-                        selection.set_head(movement::right(map, selection.head()), selection.goal);
+                        tail = head;
+                        head = movement::right(map, head);
                     }
 
+                    selection.set_head(head, new_goal);
+                    selection.set_tail(tail, new_goal);
                     // flip the selection
                     selection.swap_head_tail();
+                    head = selection.head();
+                    tail = selection.tail();
 
                     // create a classifier
-                    let classifier = map
-                        .buffer_snapshot
-                        .char_classifier_at(selection.head().to_point(map));
+                    let classifier = map.buffer_snapshot.char_classifier_at(head.to_point(map));
 
-                    let mut last_selection = selection.clone();
                     for _ in 0..times {
-                        let (new_tail, new_head) = movement::find_preceding_boundary_trail(
-                            map,
-                            selection.head(),
-                            |left, right| is_boundary(left, right, &classifier),
-                        );
-
-                        selection.set_head(new_head, selection.goal);
-                        if let Some(new_tail) = new_tail {
-                            selection.set_tail(new_tail, selection.goal);
-                        }
+                        let (maybe_next_tail, next_head) =
+                            movement::find_preceding_boundary_trail(map, head, |left, right| {
+                                is_boundary(left, right, &classifier)
+                            });
 
-                        if selection.head() == last_selection.head()
-                            && selection.tail() == last_selection.tail()
-                        {
+                        if next_head == head && maybe_next_tail.unwrap_or(next_head) == tail {
                             break;
                         }
-                        last_selection = selection.clone();
+
+                        head = next_head;
+                        if let Some(next_tail) = maybe_next_tail {
+                            tail = next_tail;
+                        }
                     }
+
+                    selection.set_tail(tail, new_goal);
+                    selection.set_head(head, new_goal);
                 });
             })
         });