Add a sticky mode for keymaps (#635)

2 files changed, 104 insertions, 45 deletions

diff --git a/helix-term/src/keymap.rs b/helix-term/src/keymap.rs
index a936dccc..f0f980bd 100644
--- a/helix-term/src/keymap.rs
+++ b/helix-term/src/keymap.rs
@@ -4,6 +4,7 @@ use helix_core::hashmap;
 use helix_view::{document::Mode, info::Info, input::KeyEvent};
 use serde::Deserialize;
 use std::{
+    borrow::Cow,
     collections::HashMap,
     ops::{Deref, DerefMut},
 };
@@ -47,13 +48,13 @@ macro_rules! keymap {
     };
 
     (@trie
-        { $label:literal $($($key:literal)|+ => $value:tt,)+ }
+        { $label:literal $(sticky=$sticky:literal)? $($($key:literal)|+ => $value:tt,)+ }
     ) => {
-        keymap!({ $label $($($key)|+ => $value,)+ })
+        keymap!({ $label $(sticky=$sticky)? $($($key)|+ => $value,)+ })
     };
 
     (
-        { $label:literal $($($key:literal)|+ => $value:tt,)+ }
+        { $label:literal $(sticky=$sticky:literal)? $($($key:literal)|+ => $value:tt,)+ }
     ) => {
         // modified from the hashmap! macro
         {
@@ -70,7 +71,9 @@ macro_rules! keymap {
                     _order.push(_key);
                 )+
             )*
-            $crate::keymap::KeyTrie::Node($crate::keymap::KeyTrieNode::new($label, _map, _order))
+            let mut _node = $crate::keymap::KeyTrieNode::new($label, _map, _order);
+            $( _node.is_sticky = $sticky; )?
+            $crate::keymap::KeyTrie::Node(_node)
         }
     };
 }
@@ -84,6 +87,8 @@ pub struct KeyTrieNode {
     map: HashMap<KeyEvent, KeyTrie>,
     #[serde(skip)]
     order: Vec<KeyEvent>,
+    #[serde(skip)]
+    pub is_sticky: bool,
 }
 
 impl KeyTrieNode {
@@ -92,6 +97,7 @@ impl KeyTrieNode {
             name: name.to_string(),
             map,
             order,
+            is_sticky: false,
         }
     }
 
@@ -119,12 +125,10 @@ impl KeyTrieNode {
             }
         }
     }
-}
 
-impl From<KeyTrieNode> for Info {
-    fn from(node: KeyTrieNode) -> Self {
-        let mut body: Vec<(&str, Vec<KeyEvent>)> = Vec::with_capacity(node.len());
-        for (&key, trie) in node.iter() {
+    pub fn infobox(&self) -> Info {
+        let mut body: Vec<(&str, Vec<KeyEvent>)> = Vec::with_capacity(self.len());
+        for (&key, trie) in self.iter() {
             let desc = match trie {
                 KeyTrie::Leaf(cmd) => cmd.doc(),
                 KeyTrie::Node(n) => n.name(),
@@ -136,16 +140,16 @@ impl From<KeyTrieNode> for Info {
             }
         }
         body.sort_unstable_by_key(|(_, keys)| {
-            node.order.iter().position(|&k| k == keys[0]).unwrap()
+            self.order.iter().position(|&k| k == keys[0]).unwrap()
         });
-        let prefix = format!("{} ", node.name());
+        let prefix = format!("{} ", self.name());
         if body.iter().all(|(desc, _)| desc.starts_with(&prefix)) {
             body = body
                 .into_iter()
                 .map(|(desc, keys)| (desc.strip_prefix(&prefix).unwrap(), keys))
                 .collect();
         }
-        Info::new(node.name(), body)
+        Info::new(self.name(), body)
     }
 }
 
@@ -218,7 +222,7 @@ impl KeyTrie {
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub enum KeymapResult {
+pub enum KeymapResultKind {
     /// Needs more keys to execute a command. Contains valid keys for next keystroke.
     Pending(KeyTrieNode),
     Matched(Command),
@@ -229,14 +233,31 @@ pub enum KeymapResult {
     Cancelled(Vec<KeyEvent>),
 }
 
+/// Returned after looking up a key in [`Keymap`]. The `sticky` field has a
+/// reference to the sticky node if one is currently active.
+pub struct KeymapResult<'a> {
+    pub kind: KeymapResultKind,
+    pub sticky: Option<&'a KeyTrieNode>,
+}
+
+impl<'a> KeymapResult<'a> {
+    pub fn new(kind: KeymapResultKind, sticky: Option<&'a KeyTrieNode>) -> Self {
+        Self { kind, sticky }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct Keymap {
     /// Always a Node
     #[serde(flatten)]
     root: KeyTrie,
-    /// Stores pending keys waiting for the next key
+    /// Stores pending keys waiting for the next key. This is relative to a
+    /// sticky node if one is in use.
     #[serde(skip)]
     state: Vec<KeyEvent>,
+    /// Stores the sticky node if one is activated.
+    #[serde(skip)]
+    sticky: Option<KeyTrieNode>,
 }
 
 impl Keymap {
@@ -244,6 +265,7 @@ impl Keymap {
         Keymap {
             root,
             state: Vec::new(),
+            sticky: None,
         }
     }
 
@@ -251,27 +273,60 @@ impl Keymap {
         &self.root
     }
 
+    pub fn sticky(&self) -> Option<&KeyTrieNode> {
+        self.sticky.as_ref()
+    }
+
     /// Returns list of keys waiting to be disambiguated.
     pub fn pending(&self) -> &[KeyEvent] {
         &self.state
     }
 
-    /// Lookup `key` in the keymap to try and find a command to execute
+    /// Lookup `key` in the keymap to try and find a command to execute. Escape
+    /// key cancels pending keystrokes. If there are no pending keystrokes but a
+    /// sticky node is in use, it will be cleared.
     pub fn get(&mut self, key: KeyEvent) -> KeymapResult {
-        let &first = self.state.get(0).unwrap_or(&key);
-        let trie = match self.root.search(&[first]) {
-            Some(&KeyTrie::Leaf(cmd)) => return KeymapResult::Matched(cmd),
-            None => return KeymapResult::NotFound,
+        if let key!(Esc) = key {
+            if self.state.is_empty() {
+                self.sticky = None;
+            }
+            return KeymapResult::new(
+                KeymapResultKind::Cancelled(self.state.drain(..).collect()),
+                self.sticky(),
+            );
+        }
+
+        let first = self.state.get(0).unwrap_or(&key);
+        let trie_node = match self.sticky {
+            Some(ref trie) => Cow::Owned(KeyTrie::Node(trie.clone())),
+            None => Cow::Borrowed(&self.root),
+        };
+
+        let trie = match trie_node.search(&[*first]) {
+            Some(&KeyTrie::Leaf(cmd)) => {
+                return KeymapResult::new(KeymapResultKind::Matched(cmd), self.sticky())
+            }
+            None => return KeymapResult::new(KeymapResultKind::NotFound, self.sticky()),
             Some(t) => t,
         };
+
         self.state.push(key);
         match trie.search(&self.state[1..]) {
-            Some(&KeyTrie::Node(ref map)) => KeymapResult::Pending(map.clone()),
-            Some(&KeyTrie::Leaf(command)) => {
+            Some(&KeyTrie::Node(ref map)) => {
+                if map.is_sticky {
+                    self.state.clear();
+                    self.sticky = Some(map.clone());
+                }
+                KeymapResult::new(KeymapResultKind::Pending(map.clone()), self.sticky())
+            }
+            Some(&KeyTrie::Leaf(cmd)) => {
                 self.state.clear();
-                KeymapResult::Matched(command)
+                return KeymapResult::new(KeymapResultKind::Matched(cmd), self.sticky());
             }
-            None => KeymapResult::Cancelled(self.state.drain(..).collect()),
+            None => KeymapResult::new(
+                KeymapResultKind::Cancelled(self.state.drain(..).collect()),
+                self.sticky(),
+            ),
         }
     }
 
@@ -602,19 +657,19 @@ fn merge_partial_keys() {
 
     let keymap = merged_config.keys.0.get_mut(&Mode::Normal).unwrap();
     assert_eq!(
-        keymap.get(key!('i')),
-        KeymapResult::Matched(Command::normal_mode),
+        keymap.get(key!('i')).kind,
+        KeymapResultKind::Matched(Command::normal_mode),
         "Leaf should replace leaf"
     );
     assert_eq!(
-        keymap.get(key!('无')),
-        KeymapResult::Matched(Command::insert_mode),
+        keymap.get(key!('无')).kind,
+        KeymapResultKind::Matched(Command::insert_mode),
         "New leaf should be present in merged keymap"
     );
     // Assumes that z is a node in the default keymap
     assert_eq!(
-        keymap.get(key!('z')),
-        KeymapResult::Matched(Command::jump_backward),
+        keymap.get(key!('z')).kind,
+        KeymapResultKind::Matched(Command::jump_backward),
         "Leaf should replace node"
     );
     // Assumes that `g` is a node in default keymap
diff --git a/helix-term/src/ui/editor.rs b/helix-term/src/ui/editor.rs
index 4b9c56e7..e8cd40cf 100644
--- a/helix-term/src/ui/editor.rs
+++ b/helix-term/src/ui/editor.rs
@@ -2,7 +2,7 @@ use crate::{
     commands,
     compositor::{Component, Context, EventResult},
     key,
-    keymap::{KeymapResult, Keymaps},
+    keymap::{KeymapResult, KeymapResultKind, Keymaps},
     ui::{Completion, ProgressSpinners},
 };
 
@@ -638,7 +638,7 @@ impl EditorView {
 
     /// Handle events by looking them up in `self.keymaps`. Returns None
     /// if event was handled (a command was executed or a subkeymap was
-    /// activated). Only KeymapResult::{NotFound, Cancelled} is returned
+    /// activated). Only KeymapResultKind::{NotFound, Cancelled} is returned
     /// otherwise.
     fn handle_keymap_event(
         &mut self,
@@ -647,29 +647,32 @@ impl EditorView {
         event: KeyEvent,
     ) -> Option<KeymapResult> {
         self.autoinfo = None;
-        match self.keymaps.get_mut(&mode).unwrap().get(event) {
-            KeymapResult::Matched(command) => command.execute(cxt),
-            KeymapResult::Pending(node) => self.autoinfo = Some(node.into()),
-            k @ KeymapResult::NotFound | k @ KeymapResult::Cancelled(_) => return Some(k),
+        let key_result = self.keymaps.get_mut(&mode).unwrap().get(event);
+        self.autoinfo = key_result.sticky.map(|node| node.infobox());
+
+        match &key_result.kind {
+            KeymapResultKind::Matched(command) => command.execute(cxt),
+            KeymapResultKind::Pending(node) => self.autoinfo = Some(node.infobox()),
+            KeymapResultKind::NotFound | KeymapResultKind::Cancelled(_) => return Some(key_result),
         }
         None
     }
 
     fn insert_mode(&mut self, cx: &mut commands::Context, event: KeyEvent) {
         if let Some(keyresult) = self.handle_keymap_event(Mode::Insert, cx, event) {
-            match keyresult {
-                KeymapResult::NotFound => {
+            match keyresult.kind {
+                KeymapResultKind::NotFound => {
                     if let Some(ch) = event.char() {
                         commands::insert::insert_char(cx, ch)
                     }
                 }
-                KeymapResult::Cancelled(pending) => {
+                KeymapResultKind::Cancelled(pending) => {
                     for ev in pending {
                         match ev.char() {
                             Some(ch) => commands::insert::insert_char(cx, ch),
                             None => {
-                                if let KeymapResult::Matched(command) =
-                                    self.keymaps.get_mut(&Mode::Insert).unwrap().get(ev)
+                                if let KeymapResultKind::Matched(command) =
+                                    self.keymaps.get_mut(&Mode::Insert).unwrap().get(ev).kind
                                 {
                                     command.execute(cx);
                                 }
@@ -976,11 +979,12 @@ impl Component for EditorView {
                         // how we entered insert mode is important, and we should track that so
                         // we can repeat the side effect.
 
-                        self.last_insert.0 = match self.keymaps.get_mut(&mode).unwrap().get(key) {
-                            KeymapResult::Matched(command) => command,
-                            // FIXME: insert mode can only be entered through single KeyCodes
-                            _ => unimplemented!(),
-                        };
+                        self.last_insert.0 =
+                            match self.keymaps.get_mut(&mode).unwrap().get(key).kind {
+                                KeymapResultKind::Matched(command) => command,
+                                // FIXME: insert mode can only be entered through single KeyCodes
+                                _ => unimplemented!(),
+                            };
                         self.last_insert.1.clear();
                     }
                     (Mode::Insert, Mode::Normal) => {