vim: Show 'j' from jk pre-emptively (#32007)

Fixes: #29812
Fixes: #22538

Co-Authored-By: <corentinhenry@gmail.com>

Release Notes:

- vim: Multi-key bindings in insert mode will now show the pending
keystroke in the buffer. For example if you have `jk` mapped to escape,
pressing `j` will immediately show a `j`.

crates/editor/src/editor.rs

@@ -299,6 +299,7 @@ pub enum DebugStackFrameLine {}
 enum DocumentHighlightRead {}
 enum DocumentHighlightWrite {}
 enum InputComposition {}
+pub enum PendingInput {}
 enum SelectedTextHighlight {}
 
 pub enum ConflictsOuter {}
@@ -1776,6 +1777,8 @@ impl Editor {
             .detach();
         cx.on_blur(&focus_handle, window, Self::handle_blur)
             .detach();
+        cx.observe_pending_input(window, Self::observe_pending_input)
+            .detach();
 
         let show_indent_guides = if matches!(mode, EditorMode::SingleLine { .. }) {
             Some(false)
@@ -19553,6 +19556,90 @@ impl Editor {
         cx.notify();
     }
 
+    pub fn observe_pending_input(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        let mut pending: String = window
+            .pending_input_keystrokes()
+            .into_iter()
+            .flatten()
+            .filter_map(|keystroke| {
+                if keystroke.modifiers.is_subset_of(&Modifiers::shift()) {
+                    Some(keystroke.key_char.clone().unwrap_or(keystroke.key.clone()))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        if !self.input_enabled || self.read_only || !self.focus_handle.is_focused(window) {
+            pending = "".to_string();
+        }
+
+        let existing_pending = self
+            .text_highlights::<PendingInput>(cx)
+            .map(|(_, ranges)| ranges.iter().cloned().collect::<Vec<_>>());
+        if existing_pending.is_none() && pending.is_empty() {
+            return;
+        }
+        let transaction =
+            self.transact(window, cx, |this, window, cx| {
+                let selections = this.selections.all::<usize>(cx);
+                let edits = selections
+                    .iter()
+                    .map(|selection| (selection.end..selection.end, pending.clone()));
+                this.edit(edits, cx);
+                this.change_selections(None, window, cx, |s| {
+                    s.select_ranges(selections.into_iter().enumerate().map(|(ix, sel)| {
+                        sel.start + ix * pending.len()..sel.end + ix * pending.len()
+                    }));
+                });
+                if let Some(existing_ranges) = existing_pending {
+                    let edits = existing_ranges.iter().map(|range| (range.clone(), ""));
+                    this.edit(edits, cx);
+                }
+            });
+
+        let snapshot = self.snapshot(window, cx);
+        let ranges = self
+            .selections
+            .all::<usize>(cx)
+            .into_iter()
+            .map(|selection| {
+                snapshot.buffer_snapshot.anchor_after(selection.end)
+                    ..snapshot
+                        .buffer_snapshot
+                        .anchor_before(selection.end + pending.len())
+            })
+            .collect();
+
+        if pending.is_empty() {
+            self.clear_highlights::<PendingInput>(cx);
+        } else {
+            self.highlight_text::<PendingInput>(
+                ranges,
+                HighlightStyle {
+                    underline: Some(UnderlineStyle {
+                        thickness: px(1.),
+                        color: None,
+                        wavy: false,
+                    }),
+                    ..Default::default()
+                },
+                cx,
+            );
+        }
+
+        self.ime_transaction = self.ime_transaction.or(transaction);
+        if let Some(transaction) = self.ime_transaction {
+            self.buffer.update(cx, |buffer, cx| {
+                buffer.group_until_transaction(transaction, cx);
+            });
+        }
+
+        if self.text_highlights::<PendingInput>(cx).is_none() {
+            self.ime_transaction.take();
+        }
+    }
+
     pub fn register_action<A: Action>(
         &mut self,
         listener: impl Fn(&A, &mut Window, &mut App) + 'static,

crates/gpui/src/platform.rs

@@ -839,7 +839,7 @@ impl PlatformInputHandler {
             .ok();
     }
 
-    fn replace_and_mark_text_in_range(
+    pub fn replace_and_mark_text_in_range(
         &mut self,
         range_utf16: Option<Range<usize>>,
         new_text: &str,

crates/gpui/src/window.rs

@@ -3542,6 +3542,7 @@ impl Window {
                         .dispatch_tree
                         .flush_dispatch(currently_pending.keystrokes, &dispatch_path);
 
+                    window.pending_input_changed(cx);
                     window.replay_pending_input(to_replay, cx)
                 })
                 .log_err();

crates/vim/src/test.rs

@@ -7,14 +7,15 @@ use std::time::Duration;
 use collections::HashMap;
 use command_palette::CommandPalette;
 use editor::{
-    DisplayPoint, Editor, EditorMode, MultiBuffer, actions::DeleteLine, display_map::DisplayRow,
-    test::editor_test_context::EditorTestContext,
+    AnchorRangeExt, DisplayPoint, Editor, EditorMode, MultiBuffer, actions::DeleteLine,
+    display_map::DisplayRow, test::editor_test_context::EditorTestContext,
 };
 use futures::StreamExt;
 use gpui::{KeyBinding, Modifiers, MouseButton, TestAppContext};
 use language::Point;
 pub use neovim_backed_test_context::*;
 use settings::SettingsStore;
+use util::test::marked_text_ranges;
 pub use vim_test_context::*;
 
 use indoc::indoc;
@@ -860,6 +861,49 @@ async fn test_jk(cx: &mut gpui::TestAppContext) {
     cx.shared_state().await.assert_eq("jˇohello");
 }
 
+fn assert_pending_input(cx: &mut VimTestContext, expected: &str) {
+    cx.update_editor(|editor, window, cx| {
+        let snapshot = editor.snapshot(window, cx);
+        let highlights = editor
+            .text_highlights::<editor::PendingInput>(cx)
+            .unwrap()
+            .1;
+        let (_, ranges) = marked_text_ranges(expected, false);
+
+        assert_eq!(
+            highlights
+                .iter()
+                .map(|highlight| highlight.to_offset(&snapshot.buffer_snapshot))
+                .collect::<Vec<_>>(),
+            ranges
+        )
+    });
+}
+
+#[gpui::test]
+async fn test_jk_multi(cx: &mut gpui::TestAppContext) {
+    let mut cx = VimTestContext::new(cx, true).await;
+
+    cx.update(|_, cx| {
+        cx.bind_keys([KeyBinding::new(
+            "j k l",
+            NormalBefore,
+            Some("vim_mode == insert"),
+        )])
+    });
+
+    cx.set_state("ˇone ˇone ˇone", Mode::Normal);
+    cx.simulate_keystrokes("i j");
+    cx.simulate_keystrokes("k");
+    cx.assert_state("ˇjkone ˇjkone ˇjkone", Mode::Insert);
+    assert_pending_input(&mut cx, "«jk»one «jk»one «jk»one");
+    cx.simulate_keystrokes("o j k");
+    cx.assert_state("jkoˇjkone jkoˇjkone jkoˇjkone", Mode::Insert);
+    assert_pending_input(&mut cx, "jko«jk»one jko«jk»one jko«jk»one");
+    cx.simulate_keystrokes("l");
+    cx.assert_state("jkˇoone jkˇoone jkˇoone", Mode::Normal);
+}
+
 #[gpui::test]
 async fn test_jk_delay(cx: &mut gpui::TestAppContext) {
     let mut cx = VimTestContext::new(cx, true).await;
@@ -876,7 +920,22 @@ async fn test_jk_delay(cx: &mut gpui::TestAppContext) {
     cx.simulate_keystrokes("i j");
     cx.executor().advance_clock(Duration::from_millis(500));
     cx.run_until_parked();
-    cx.assert_state("ˇhello", Mode::Insert);
+    cx.assert_state("ˇjhello", Mode::Insert);
+    cx.update_editor(|editor, window, cx| {
+        let snapshot = editor.snapshot(window, cx);
+        let highlights = editor
+            .text_highlights::<editor::PendingInput>(cx)
+            .unwrap()
+            .1;
+
+        assert_eq!(
+            highlights
+                .iter()
+                .map(|highlight| highlight.to_offset(&snapshot.buffer_snapshot))
+                .collect::<Vec<_>>(),
+            vec![0..1]
+        )
+    });
     cx.executor().advance_clock(Duration::from_millis(500));
     cx.run_until_parked();
     cx.assert_state("jˇhello", Mode::Insert);