Finally: Retain horizontal position when moving vertically.

2 files changed, 56 insertions, 37 deletions

diff --git a/helix-core/src/selection.rs b/helix-core/src/selection.rs
index 121267bf..87216fd9 100644
--- a/helix-core/src/selection.rs
+++ b/helix-core/src/selection.rs
@@ -23,11 +23,16 @@ pub struct Range {
     pub anchor: usize,
     /// The head of the range, moved when extending.
     pub head: usize,
+    pub horiz: Option<u32>,
 } // TODO: might be cheaper to store normalized as from/to and an inverted flag
 
 impl Range {
     pub fn new(anchor: usize, head: usize) -> Self {
-        Self { anchor, head }
+        Self {
+            anchor,
+            head,
+            horiz: None,
+        }
     }
 
     /// Start of the range.
@@ -83,7 +88,11 @@ impl Range {
         if self.anchor == anchor && self.head == head {
             return self;
         }
-        Self { anchor, head }
+        Self {
+            anchor,
+            head,
+            horiz: None,
+        }
     }
 
     /// Extend the range to cover at least `from` `to`.
@@ -93,6 +102,7 @@ impl Range {
             return Range {
                 anchor: from,
                 head: to,
+                horiz: None,
             };
         }
 
@@ -103,6 +113,7 @@ impl Range {
             } else {
                 to
             },
+            horiz: None,
         }
     }
 
@@ -174,7 +185,11 @@ impl Selection {
     /// Constructs a selection holding a single range.
     pub fn single(anchor: usize, head: usize) -> Self {
         Self {
-            ranges: smallvec![Range { anchor, head }],
+            ranges: smallvec![Range {
+                anchor,
+                head,
+                horiz: None
+            }],
             primary_index: 0,
         }
     }
diff --git a/helix-core/src/state.rs b/helix-core/src/state.rs
index 4d531aa0..d2be266c 100644
--- a/helix-core/src/state.rs
+++ b/helix-core/src/state.rs
@@ -19,9 +19,7 @@ pub enum Direction {
 #[derive(Copy, Clone, PartialEq, Eq)]
 pub enum Granularity {
     Character,
-    Word,
     Line,
-    // LineBoundary
 }
 
 impl State {
@@ -87,23 +85,26 @@ impl State {
     // 2. compose onto a ongoing transaction
     // 3. on insert mode leave, that transaction gets stored into undo history
 
-    pub fn move_pos(
+    pub fn move_range(
         &self,
-        pos: usize,
+        range: Range,
         dir: Direction,
         granularity: Granularity,
         count: usize,
-    ) -> usize {
+        extend: bool,
+    ) -> Range {
         let text = &self.doc;
+        let pos = range.head;
         match (dir, granularity) {
             (Direction::Backward, Granularity::Character) => {
                 // Clamp to line
                 let line = text.char_to_line(pos);
                 let start = text.line_to_char(line);
-                std::cmp::max(
+                let pos = std::cmp::max(
                     nth_prev_grapheme_boundary(&text.slice(..), pos, count),
                     start,
-                )
+                );
+                Range::new(if extend { range.anchor } else { pos }, pos)
             }
             (Direction::Forward, Granularity::Character) => {
                 // Clamp to line
@@ -111,16 +112,11 @@ impl State {
                 // Line end is pos at the start of next line - 1
                 // subtract another 1 because the line ends with \n
                 let end = text.line_to_char(line + 1).saturating_sub(2);
-                std::cmp::min(nth_next_grapheme_boundary(&text.slice(..), pos, count), end)
-            }
-            (Direction::Forward, Granularity::Word) => {
-                Self::move_next_word_start(&text.slice(..), pos)
+                let pos =
+                    std::cmp::min(nth_next_grapheme_boundary(&text.slice(..), pos, count), end);
+                Range::new(if extend { range.anchor } else { pos }, pos)
             }
-            (Direction::Backward, Granularity::Word) => {
-                Self::move_prev_word_start(&text.slice(..), pos)
-            }
-            (_, Granularity::Line) => move_vertically(&text.slice(..), dir, pos, count),
-            _ => pos,
+            (_, Granularity::Line) => move_vertically(&text.slice(..), dir, range, count, extend),
         }
     }
 
@@ -205,10 +201,8 @@ impl State {
         // move all selections according to normal cursor move semantics by collapsing it
         // into cursors and moving them vertically
 
-        self.selection.transform(|range| {
-            let pos = self.move_pos(range.head, dir, granularity, count);
-            Range::new(pos, pos)
-        })
+        self.selection
+            .transform(|range| self.move_range(range, dir, granularity, count, false))
     }
 
     pub fn extend_selection(
@@ -217,10 +211,8 @@ impl State {
         granularity: Granularity,
         count: usize,
     ) -> Selection {
-        self.selection.transform(|range| {
-            let pos = self.move_pos(range.head, dir, granularity, count);
-            Range::new(range.anchor, pos)
-        })
+        self.selection
+            .transform(|range| self.move_range(range, dir, granularity, count, true))
     }
 }
 
@@ -239,8 +231,16 @@ pub fn pos_at_coords(text: &RopeSlice, coords: Position) -> usize {
     nth_next_grapheme_boundary(text, line_start, col)
 }
 
-fn move_vertically(text: &RopeSlice, dir: Direction, pos: usize, count: usize) -> usize {
-    let Position { row, col } = coords_at_pos(text, pos);
+fn move_vertically(
+    text: &RopeSlice,
+    dir: Direction,
+    range: Range,
+    count: usize,
+    extend: bool,
+) -> Range {
+    let Position { row, col } = coords_at_pos(text, range.head);
+
+    let horiz = range.horiz.unwrap_or(col as u32);
 
     let new_line = match dir {
         Direction::Backward => row.saturating_sub(count),
@@ -250,14 +250,14 @@ fn move_vertically(text: &RopeSlice, dir: Direction, pos: usize, count: usize) -
     // convert to 0-indexed, subtract another 1 because len_chars() counts \n
     let new_line_len = text.line(new_line).len_chars().saturating_sub(2);
 
-    let new_col = if new_line_len < col {
-        // TODO: preserve horiz here
-        new_line_len
-    } else {
-        col
-    };
+    let new_col = std::cmp::min(horiz as usize, new_line_len);
+
+    let pos = pos_at_coords(text, Position::new(new_line, new_col));
 
-    pos_at_coords(text, Position::new(new_line, new_col))
+    let mut range = Range::new(if extend { range.anchor } else { pos }, pos);
+    use std::convert::TryInto;
+    range.horiz = Some(horiz);
+    range
 }
 
 // used for by-word movement
@@ -346,8 +346,12 @@ mod test {
         let pos = pos_at_coords(&text.slice(..), (0, 4).into());
         let slice = text.slice(..);
 
+        let range = Range::new(pos, pos);
         assert_eq!(
-            coords_at_pos(&slice, move_vertically(&slice, Direction::Forward, pos, 1)),
+            coords_at_pos(
+                &slice,
+                move_vertically(&slice, Direction::Forward, range, 1).head
+            ),
             (1, 2).into()
         );
     }