Fix keymap async command sequence execution order

If a keymap triggered a command sequence and the sequence contained a
command that needs user input, the following commands were executed
before the command that was waiting for input, which resulted in
commands being executed in the wrong order.

helix-term/src/commands.rs

@@ -75,10 +75,16 @@ pub struct Context<'a> {
     pub editor: &'a mut Editor,
 
     pub callback: Option<crate::compositor::Callback>,
-    pub on_next_key_callback: Option<Box<dyn FnOnce(&mut Context, KeyEvent)>>,
+    pub on_next_key_callback: Option<Box<dyn FnOnce(&mut Context, KeyEvent) -> CommandStatus>>,
     pub jobs: &'a mut Jobs,
 }
 
+#[derive(PartialEq)]
+pub enum CommandStatus {
+    Finished,
+    NeedsInput,
+}
+
 impl<'a> Context<'a> {
     /// Push a new component onto the compositor.
     pub fn push_layer(&mut self, component: Box<dyn Component>) {
@@ -90,7 +96,7 @@ impl<'a> Context<'a> {
     #[inline]
     pub fn on_next_key(
         &mut self,
-        on_next_key_callback: impl FnOnce(&mut Context, KeyEvent) + 'static,
+        on_next_key_callback: impl FnOnce(&mut Context, KeyEvent) -> CommandStatus + 'static,
     ) {
         self.on_next_key_callback = Some(Box::new(on_next_key_callback));
     }
@@ -1121,13 +1127,14 @@ where
                 code: KeyCode::Char(ch),
                 ..
             } => ch,
-            _ => return,
+            _ => return CommandStatus::Finished,
         };
 
         find_char_impl(cx.editor, &search_fn, inclusive, extend, ch, count);
         cx.editor.last_motion = Some(Motion(Box::new(move |editor: &mut Editor| {
             find_char_impl(editor, &search_fn, inclusive, true, ch, 1);
         })));
+        CommandStatus::Finished
     })
 }
 
@@ -1292,6 +1299,7 @@ fn replace(cx: &mut Context) {
 
             apply_transaction(&transaction, doc, view);
         }
+        CommandStatus::Finished
     })
 }
 
@@ -4247,6 +4255,7 @@ fn select_register(cx: &mut Context) {
             cx.editor.autoinfo = None;
             cx.editor.selected_register = Some(ch);
         }
+        CommandStatus::Finished
     })
 }
 
@@ -4258,6 +4267,7 @@ fn insert_register(cx: &mut Context) {
             cx.register = Some(ch);
             paste(cx, Paste::Cursor);
         }
+        CommandStatus::Finished
     })
 }
 
@@ -4435,6 +4445,7 @@ fn select_textobject(cx: &mut Context, objtype: textobject::TextObject) {
             textobject(cx.editor);
             cx.editor.last_motion = Some(Motion(Box::new(textobject)));
         }
+        CommandStatus::Finished
     });
 
     let title = match objtype {
@@ -4462,7 +4473,7 @@ fn surround_add(cx: &mut Context) {
     cx.on_next_key(move |cx, event| {
         let ch = match event.char() {
             Some(ch) => ch,
-            None => return,
+            None => return CommandStatus::Finished,
         };
         let (view, doc) = current!(cx.editor);
         let selection = doc.selection(view.id);
@@ -4480,6 +4491,7 @@ fn surround_add(cx: &mut Context) {
 
         let transaction = Transaction::change(doc.text(), changes.into_iter());
         apply_transaction(&transaction, doc, view);
+        CommandStatus::Finished
     })
 }
 
@@ -4489,7 +4501,7 @@ fn surround_replace(cx: &mut Context) {
         let surround_ch = match event.char() {
             Some('m') => None, // m selects the closest surround pair
             Some(ch) => Some(ch),
-            None => return,
+            None => return CommandStatus::Finished,
         };
         let (view, doc) = current!(cx.editor);
         let text = doc.text().slice(..);
@@ -4499,15 +4511,15 @@ fn surround_replace(cx: &mut Context) {
             Ok(c) => c,
             Err(err) => {
                 cx.editor.set_error(err.to_string());
-                return;
+                return CommandStatus::Finished;
             }
         };
 
         cx.on_next_key(move |cx, event| {
             let (view, doc) = current!(cx.editor);
             let to = match event.char() {
                 Some(to) => to,
-                None => return,
+                None => return CommandStatus::Finished,
             };
             let (open, close) = surround::get_pair(to);
             let transaction = Transaction::change(
@@ -4519,7 +4531,9 @@ fn surround_replace(cx: &mut Context) {
                 }),
             );
             apply_transaction(&transaction, doc, view);
+            CommandStatus::Finished
         });
+        CommandStatus::NeedsInput
     })
 }
 
@@ -4529,7 +4543,7 @@ fn surround_delete(cx: &mut Context) {
         let surround_ch = match event.char() {
             Some('m') => None, // m selects the closest surround pair
             Some(ch) => Some(ch),
-            None => return,
+            None => return CommandStatus::Finished,
         };
         let (view, doc) = current!(cx.editor);
         let text = doc.text().slice(..);
@@ -4539,13 +4553,14 @@ fn surround_delete(cx: &mut Context) {
             Ok(c) => c,
             Err(err) => {
                 cx.editor.set_error(err.to_string());
-                return;
+                return CommandStatus::Finished;
             }
         };
 
         let transaction =
             Transaction::change(doc.text(), change_pos.into_iter().map(|p| (p, p + 1, None)));
         apply_transaction(&transaction, doc, view);
+        CommandStatus::Finished
     })
 }
 

helix-term/src/ui/editor.rs

@@ -24,7 +24,7 @@ use helix_view::{
     keyboard::{KeyCode, KeyModifiers},
     Document, Editor, Theme, View,
 };
-use std::{borrow::Cow, path::PathBuf};
+use std::{borrow::Cow, num::NonZeroUsize, path::PathBuf};
 
 use tui::buffer::Buffer as Surface;
 
@@ -33,31 +33,70 @@ use super::statusline;
 
 pub struct EditorView {
     pub keymaps: Keymaps,
-    on_next_key: Option<Box<dyn FnOnce(&mut commands::Context, KeyEvent)>>,
+    on_next_key:
+        Option<Box<dyn FnOnce(&mut commands::Context, KeyEvent) -> commands::CommandStatus>>,
+    pending_commands: PendingCommands,
     pseudo_pending: Vec<KeyEvent>,
     last_insert: (commands::MappableCommand, Vec<InsertEvent>),
     pub(crate) completion: Option<Completion>,
     spinners: ProgressSpinners,
 }
 
+/// Stores commands that are pending because a previous command is
+/// waiting on input. Also stores context information that was present
+/// when the command chain was first executed.
+pub struct PendingCommands {
+    pub commands: Vec<commands::MappableCommand>,
+    pub register: Option<char>,
+    pub count: Option<NonZeroUsize>,
+}
+
 #[derive(Debug, Clone)]
 pub enum InsertEvent {
     Key(KeyEvent),
     CompletionApply(CompleteAction),
     TriggerCompletion,
 }
 
+enum KeyEventOrResult {
+    Event(KeyEvent),
+    Result(KeymapResult),
+}
+
+impl From<KeyEvent> for KeyEventOrResult {
+    fn from(event: KeyEvent) -> Self {
+        Self::Event(event)
+    }
+}
+
+impl From<KeymapResult> for KeyEventOrResult {
+    fn from(result: KeymapResult) -> Self {
+        Self::Result(result)
+    }
+}
+
 impl Default for EditorView {
     fn default() -> Self {
         Self::new(Keymaps::default())
     }
 }
 
+impl PendingCommands {
+    pub fn new() -> Self {
+        Self {
+            commands: Vec::new(),
+            register: None,
+            count: None,
+        }
+    }
+}
+
 impl EditorView {
     pub fn new(keymaps: Keymaps) -> Self {
         Self {
             keymaps,
             on_next_key: None,
+            pending_commands: PendingCommands::new(),
             pseudo_pending: Vec::new(),
             last_insert: (commands::MappableCommand::normal_mode, Vec::new()),
             completion: None,
@@ -916,53 +955,71 @@ impl EditorView {
         &mut self,
         mode: Mode,
         cxt: &mut commands::Context,
-        event: KeyEvent,
+        event: impl Into<KeyEventOrResult>,
     ) -> Option<KeymapResult> {
         let mut last_mode = mode;
         self.pseudo_pending.extend(self.keymaps.pending());
-        let key_result = self.keymaps.get(mode, event);
+        let key_result = match event.into() {
+            KeyEventOrResult::Event(event) => self.keymaps.get(mode, event),
+            KeyEventOrResult::Result(result) => result,
+        };
         cxt.editor.autoinfo = self.keymaps.sticky().map(|node| node.infobox());
 
-        let mut execute_command = |command: &commands::MappableCommand| {
-            command.execute(cxt);
-            let current_mode = cxt.editor.mode();
-            match (last_mode, current_mode) {
-                (Mode::Normal, Mode::Insert) => {
-                    // HAXX: if we just entered insert mode from normal, clear key buf
-                    // and record the command that got us into this mode.
-
-                    // how we entered insert mode is important, and we should track that so
-                    // we can repeat the side effect.
-                    self.last_insert.0 = command.clone();
-                    self.last_insert.1.clear();
-
-                    commands::signature_help_impl(cxt, commands::SignatureHelpInvoked::Automatic);
-                }
-                (Mode::Insert, Mode::Normal) => {
-                    // if exiting insert mode, remove completion
-                    self.completion = None;
-
-                    // TODO: Use an on_mode_change hook to remove signature help
-                    cxt.jobs.callback(async {
-                        let call: job::Callback = Box::new(|_editor, compositor| {
-                            compositor.remove(SignatureHelp::ID);
+        let mut execute_command =
+            |cxt: &mut commands::Context, command: &commands::MappableCommand| {
+                command.execute(cxt);
+                let current_mode = cxt.editor.mode();
+                match (last_mode, current_mode) {
+                    (Mode::Normal, Mode::Insert) => {
+                        // HAXX: if we just entered insert mode from normal, clear key buf
+                        // and record the command that got us into this mode.
+
+                        // how we entered insert mode is important, and we should track that so
+                        // we can repeat the side effect.
+                        self.last_insert.0 = command.clone();
+                        self.last_insert.1.clear();
+
+                        commands::signature_help_impl(
+                            cxt,
+                            commands::SignatureHelpInvoked::Automatic,
+                        );
+                    }
+                    (Mode::Insert, Mode::Normal) => {
+                        // if exiting insert mode, remove completion
+                        self.completion = None;
+
+                        // TODO: Use an on_mode_change hook to remove signature help
+                        cxt.jobs.callback(async {
+                            let call: job::Callback = Box::new(|_editor, compositor| {
+                                compositor.remove(SignatureHelp::ID);
+                            });
+                            Ok(call)
                         });
-                        Ok(call)
-                    });
+                    }
+                    _ => (),
                 }
-                _ => (),
-            }
-            last_mode = current_mode;
-        };
+                last_mode = current_mode;
+            };
 
-        match &key_result {
+        match key_result {
             KeymapResult::Matched(command) => {
-                execute_command(command);
+                execute_command(cxt, &command);
             }
             KeymapResult::Pending(node) => cxt.editor.autoinfo = Some(node.infobox()),
-            KeymapResult::MatchedSequence(commands) => {
-                for command in commands {
-                    execute_command(command);
+            KeymapResult::MatchedSequence(mut commands) => {
+                while !commands.is_empty() {
+                    // if a previous command is waiting for input, stop executing
+                    // and store the remaining commands
+                    if cxt.on_next_key_callback.is_some() {
+                        self.pending_commands = PendingCommands {
+                            commands,
+                            register: cxt.register.clone(),
+                            count: cxt.count.clone(),
+                        };
+                        break;
+                    } else {
+                        execute_command(cxt, &commands.remove(0));
+                    }
                 }
             }
             KeymapResult::NotFound | KeymapResult::Cancelled(_) => return Some(key_result),
@@ -1339,7 +1396,23 @@ impl Component for EditorView {
 
                 if let Some(on_next_key) = self.on_next_key.take() {
                     // if there's a command waiting input, do that first
-                    on_next_key(&mut cx, key);
+                    let command_status = on_next_key(&mut cx, key);
+                    // if we have pending commands and the current command is finished, execute them
+                    if !self.pending_commands.commands.is_empty()
+                        && command_status == commands::CommandStatus::Finished
+                    {
+                        let pending =
+                            std::mem::replace(&mut self.pending_commands, PendingCommands::new());
+
+                        cx.count = pending.count;
+                        cx.register = pending.register;
+
+                        self.handle_keymap_event(
+                            mode,
+                            &mut cx,
+                            KeymapResult::MatchedSequence(pending.commands),
+                        );
+                    }
                 } else {
                     match mode {
                         Mode::Insert => {