Add an Amp-like jump command

Co-authored-by: Michael Davis <mcarsondavis@gmail.com>

2 files changed, 189 insertions, 3 deletions

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 6de05651..d545480b 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -10,9 +10,12 @@ use tui::widgets::Row;
 pub use typed::*;
 
 use helix_core::{
-    char_idx_at_visual_offset, comment,
+    char_idx_at_visual_offset,
+    chars::char_is_word,
+    comment,
     doc_formatter::TextFormat,
-    encoding, find_workspace, graphemes,
+    encoding, find_workspace,
+    graphemes::{self, next_grapheme_boundary, RevRopeGraphemes},
     history::UndoKind,
     increment, indent,
     indent::IndentStyle,
@@ -24,7 +27,7 @@ use helix_core::{
     search::{self, CharMatcher},
     selection, shellwords, surround,
     syntax::{BlockCommentToken, LanguageServerFeature},
-    text_annotations::TextAnnotations,
+    text_annotations::{Overlay, TextAnnotations},
     textobject,
     unicode::width::UnicodeWidthChar,
     visual_offset_from_block, Deletion, LineEnding, Position, Range, Rope, RopeGraphemes,
@@ -502,6 +505,8 @@ impl MappableCommand {
         record_macro, "Record macro",
         replay_macro, "Replay macro",
         command_palette, "Open command palette",
+        goto_word, "Jump to a two-character label",
+        extend_to_word, "Extend to a two-character label",
     );
 }
 
@@ -5814,3 +5819,182 @@ fn replay_macro(cx: &mut Context) {
         cx.editor.macro_replaying.pop();
     }));
 }
+
+fn goto_word(cx: &mut Context) {
+    jump_to_word(cx, Movement::Move)
+}
+
+fn extend_to_word(cx: &mut Context) {
+    jump_to_word(cx, Movement::Extend)
+}
+
+fn jump_to_label(cx: &mut Context, labels: Vec<Range>, behaviour: Movement) {
+    let doc = doc!(cx.editor);
+    let alphabet = &cx.editor.config().jump_label_alphabet;
+    if labels.is_empty() {
+        return;
+    }
+    let alphabet_char = |i| {
+        let mut res = Tendril::new();
+        res.push(alphabet[i]);
+        res
+    };
+
+    // Add label for each jump candidate to the View as virtual text.
+    let text = doc.text().slice(..);
+    let mut overlays: Vec<_> = labels
+        .iter()
+        .enumerate()
+        .flat_map(|(i, range)| {
+            [
+                Overlay::new(range.from(), alphabet_char(i / alphabet.len())),
+                Overlay::new(
+                    graphemes::next_grapheme_boundary(text, range.from()),
+                    alphabet_char(i % alphabet.len()),
+                ),
+            ]
+        })
+        .collect();
+    overlays.sort_unstable_by_key(|overlay| overlay.char_idx);
+    let (view, doc) = current!(cx.editor);
+    doc.set_jump_labels(view.id, overlays);
+
+    // Accept two characters matching a visible label. Jump to the candidate
+    // for that label if it exists.
+    let primary_selection = doc.selection(view.id).primary();
+    let view = view.id;
+    let doc = doc.id();
+    cx.on_next_key(move |cx, event| {
+        let alphabet = &cx.editor.config().jump_label_alphabet;
+        let Some(i ) = event.char().and_then(|ch| alphabet.iter().position(|&it| it == ch)) else {
+            doc_mut!(cx.editor, &doc).remove_jump_labels(view);
+            return;
+        };
+        let outer = i * alphabet.len();
+        // Bail if the given character cannot be a jump label.
+        if outer > labels.len() {
+            doc_mut!(cx.editor, &doc).remove_jump_labels(view);
+            return;
+        }
+        cx.on_next_key(move |cx, event| {
+            doc_mut!(cx.editor, &doc).remove_jump_labels(view);
+            let alphabet = &cx.editor.config().jump_label_alphabet;
+            let Some(inner ) = event.char().and_then(|ch| alphabet.iter().position(|&it| it == ch)) else {
+                return;
+            };
+            if let Some(mut range) = labels.get(outer + inner).copied() {
+                range = if behaviour == Movement::Extend {
+                    let anchor = if range.anchor < range.head {
+                        let from = primary_selection.from();
+                        if range.anchor < from {
+                            range.anchor
+                        } else {
+                            from
+                        }
+                    } else {
+                        let to = primary_selection.to();
+                        if range.anchor > to {
+                            range.anchor
+                        } else {
+                            to
+                        }
+                    };
+                     Range::new(anchor, range.head)
+                }else{
+                    range.with_direction(Direction::Forward)
+                };
+                doc_mut!(cx.editor, &doc).set_selection(view, range.into());
+            }
+        });
+    });
+}
+
+fn jump_to_word(cx: &mut Context, behaviour: Movement) {
+    // Calculate the jump candidates: ranges for any visible words with two or
+    // more characters.
+    let alphabet = &cx.editor.config().jump_label_alphabet;
+    let jump_label_limit = alphabet.len() * alphabet.len();
+    let mut words = Vec::with_capacity(jump_label_limit);
+    let (view, doc) = current_ref!(cx.editor);
+    let text = doc.text().slice(..);
+
+    // This is not necessarily exact if there is virtual text like soft wrap.
+    // It's ok though because the extra jump labels will not be rendered.
+    let start = text.line_to_char(text.char_to_line(view.offset.anchor));
+    let end = text.line_to_char(view.estimate_last_doc_line(doc) + 1);
+
+    let primary_selection = doc.selection(view.id).primary();
+    let cursor = primary_selection.cursor(text);
+    let mut cursor_fwd = Range::point(cursor);
+    let mut cursor_rev = Range::point(cursor);
+    if text.get_char(cursor).is_some_and(|c| !c.is_whitespace()) {
+        let cursor_word_end = movement::move_next_word_end(text, cursor_fwd, 1);
+        //  single grapheme words need a specical case
+        if cursor_word_end.anchor == cursor {
+            cursor_fwd = cursor_word_end;
+        }
+        let cursor_word_start = movement::move_prev_word_start(text, cursor_rev, 1);
+        if cursor_word_start.anchor == next_grapheme_boundary(text, cursor) {
+            cursor_rev = cursor_word_start;
+        }
+    }
+    'outer: loop {
+        let mut changed = false;
+        while cursor_fwd.head < end {
+            cursor_fwd = movement::move_next_word_end(text, cursor_fwd, 1);
+            // The cursor is on a word that is atleast two graphemes long and
+            // madeup of word characters. The latter condition is needed because
+            // move_next_word_end simply treats a sequence of characters from
+            // the same char class as a word so `=<` would also count as a word.
+            let add_label = RevRopeGraphemes::new(text.slice(..cursor_fwd.head))
+                .take(2)
+                .take_while(|g| g.chars().all(char_is_word))
+                .count()
+                == 2;
+            if !add_label {
+                continue;
+            }
+            changed = true;
+            // skip any leading whitespace
+            cursor_fwd.anchor += text
+                .chars_at(cursor_fwd.anchor)
+                .take_while(|&c| !char_is_word(c))
+                .count();
+            words.push(cursor_fwd);
+            if words.len() == jump_label_limit {
+                break 'outer;
+            }
+            break;
+        }
+        while cursor_rev.head > start {
+            cursor_rev = movement::move_prev_word_start(text, cursor_rev, 1);
+            // The cursor is on a word that is atleast two graphemes long and
+            // madeup of word characters. The latter condition is needed because
+            // move_prev_word_start simply treats a sequence of characters from
+            // the same char class as a word so `=<` would also count as a word.
+            let add_label = RopeGraphemes::new(text.slice(cursor_rev.head..))
+                .take(2)
+                .take_while(|g| g.chars().all(char_is_word))
+                .count()
+                == 2;
+            if !add_label {
+                continue;
+            }
+            changed = true;
+            cursor_rev.anchor -= text
+                .chars_at(cursor_rev.anchor)
+                .reversed()
+                .take_while(|&c| !char_is_word(c))
+                .count();
+            words.push(cursor_rev);
+            if words.len() == jump_label_limit {
+                break 'outer;
+            }
+            break;
+        }
+        if !changed {
+            break;
+        }
+    }
+    jump_to_label(cx, words, behaviour)
+}
diff --git a/helix-term/src/keymap/default.rs b/helix-term/src/keymap/default.rs
index bab662b0..ca5a21d2 100644
--- a/helix-term/src/keymap/default.rs
+++ b/helix-term/src/keymap/default.rs
@@ -58,6 +58,7 @@ pub fn default() -> HashMap<Mode, KeyTrie> {
             "k" => move_line_up,
             "j" => move_line_down,
             "." => goto_last_modification,
+            "w" => goto_word,
         },
         ":" => command_mode,
 
@@ -360,6 +361,7 @@ pub fn default() -> HashMap<Mode, KeyTrie> {
         "g" => { "Goto"
             "k" => extend_line_up,
             "j" => extend_line_down,
+            "w" => extend_to_word,
         },
     }));
     let insert = keymap!({ "Insert mode"