From d0c5784faf59b932297a4a58e3173aab27d7ec0b Mon Sep 17 00:00:00 2001
From: Roman Perehonchuk <romanp@meta.com>
Date: Mon, 1 Apr 2024 13:11:05 +0100
Subject: [PATCH] Amp-like jump: remove not relevant overlays after keypress

---
 helix-term/src/commands.rs | 62 ++++++++++++++++++++++++++------------
 1 file changed, 43 insertions(+), 19 deletions(-)

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index d0b9047c8..2d85cbf53 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -5940,28 +5940,10 @@ fn jump_to_label(cx: &mut Context, labels: Vec<Range>, behaviour: Movement) {
     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 overlays = labels_to_overlays(labels.iter(), alphabet, text, None);
     let (view, doc) = current!(cx.editor);
     doc.set_jump_labels(view.id, overlays);
 
@@ -5982,6 +5964,15 @@ fn jump_to_label(cx: &mut Context, labels: Vec<Range>, behaviour: Movement) {
             doc_mut!(cx.editor, &doc).remove_jump_labels(view);
             return;
         }
+        if let Some(ch) = event.char() {
+            let mut grapheme = Tendril::new();
+            grapheme.push(ch);
+
+            let doc = doc_mut!(cx.editor, &doc);
+            let text = doc.text().slice(..);
+            let overlays = labels_to_overlays(labels.iter().skip(outer).take(alphabet.len()), alphabet, text, Some(grapheme));
+            doc.set_jump_labels(view, overlays);
+        }
         cx.on_next_key(move |cx, event| {
             doc_mut!(cx.editor, &doc).remove_jump_labels(view);
             let alphabet = &cx.editor.config().jump_label_alphabet;
@@ -6015,6 +6006,39 @@ fn jump_to_label(cx: &mut Context, labels: Vec<Range>, behaviour: Movement) {
     });
 }
 
+fn alphabet_char(alphabet: &[char], index: usize) -> Tendril {
+    let mut res = Tendril::new();
+    res.push(alphabet[index]);
+    res
+}
+
+fn labels_to_overlays<'a>(
+    it: impl Iterator<Item = &'a Range>,
+    alphabet: &[char],
+    text: RopeSlice,
+    grapheme: Option<Tendril>,
+) -> Vec<Overlay> {
+    let mut overlays: Vec<_> = it
+        .enumerate()
+        .flat_map(|(i, range)| {
+            [
+                Overlay::new(
+                    range.from(),
+                    grapheme
+                        .clone()
+                        .unwrap_or_else(|| alphabet_char(alphabet, i / alphabet.len())),
+                ),
+                Overlay::new(
+                    graphemes::next_grapheme_boundary(text, range.from()),
+                    alphabet_char(alphabet, i % alphabet.len()),
+                ),
+            ]
+        })
+        .collect();
+    overlays.sort_unstable_by_key(|overlay| overlay.char_idx);
+    overlays
+}
+
 fn jump_to_word(cx: &mut Context, behaviour: Movement) {
     // Calculate the jump candidates: ranges for any visible words with two or
     // more characters.