indigo_core/
keymap.rs

1use crate::key::{Key, Keys};
2use arrayvec::ArrayVec;
3use std::{iter, mem};
4
5struct Trie<K, V> {
6    root: Node<K, V>,
7}
8
9impl<K, V> Trie<K, V> {
10    #[must_use]
11    fn new() -> Self {
12        Self {
13            root: Node::default(),
14        }
15    }
16
17    fn insert(&mut self, key: &[K], value: V) -> Option<V>
18    where
19        K: Clone + Ord,
20    {
21        self.root.insert(key, value)
22    }
23
24    fn remove(&mut self, key: &[K]) -> Option<V>
25    where
26        K: Ord,
27    {
28        self.root.remove(key)
29    }
30
31    #[must_use]
32    fn get(&self, key: &[K]) -> TrieResult<&V>
33    where
34        K: Ord,
35    {
36        self.root.get(key)
37    }
38
39    #[must_use]
40    fn get_mut(&mut self, key: &[K]) -> TrieResult<&mut V>
41    where
42        K: Ord,
43    {
44        self.root.get_mut(key)
45    }
46
47    fn iter<const N: usize>(&self) -> impl Iterator<Item = (ArrayVec<K, N>, &V)>
48    where
49        K: Clone,
50    {
51        let mut pairs = Vec::new();
52        self.root
53            .collect_pairs::<N>(&mut ArrayVec::new(), &mut pairs);
54        pairs.into_iter()
55    }
56}
57
58impl<K, V> Default for Trie<K, V> {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64enum Node<K, V> {
65    Branch {
66        keys: Vec<K>,
67        children: Vec<Box<Self>>,
68    },
69    Leaf {
70        value: V,
71    },
72}
73
74impl<K, V> Node<K, V> {
75    fn insert(&mut self, key: &[K], value: V) -> Option<V>
76    where
77        K: Clone + Ord,
78    {
79        match self {
80            Self::Branch { keys, children } => {
81                let Some((head, tail)) = key.split_first() else {
82                    panic!("Overlapping keys are forbidden")
83                };
84                match keys.binary_search(head) {
85                    Ok(index) => children[index].insert(tail, value),
86                    Err(index) => {
87                        keys.insert(index, head.clone());
88                        let child = if tail.is_empty() {
89                            Box::new(Self::Leaf { value })
90                        } else {
91                            let mut branch = Self::default();
92                            branch.insert(tail, value);
93                            Box::new(branch)
94                        };
95                        children.insert(index, child);
96                        None
97                    }
98                }
99            }
100            Self::Leaf { value: leaf_value } => {
101                assert!(key.is_empty(), "Overlapping keys are forbidden");
102                Some(std::mem::replace(leaf_value, value))
103            }
104        }
105    }
106
107    fn is_empty(&self) -> bool {
108        match self {
109            Self::Branch { keys, .. } => keys.is_empty(),
110            Self::Leaf { .. } => false,
111        }
112    }
113
114    fn remove(&mut self, key: &[K]) -> Option<V>
115    where
116        K: Ord,
117    {
118        match self {
119            Self::Branch { keys, children } => {
120                let (head, tail) = key.split_first()?;
121                let index = keys.binary_search(head).ok()?;
122                let result = children[index].remove(tail);
123                if children[index].is_empty() {
124                    keys.remove(index);
125                    children.remove(index);
126                }
127                result
128            }
129            Self::Leaf { .. } => {
130                if !key.is_empty() {
131                    return None;
132                }
133                match std::mem::take(self) {
134                    Self::Leaf { value } => Some(value),
135                    Self::Branch { .. } => unreachable!(),
136                }
137            }
138        }
139    }
140
141    #[must_use]
142    fn get(&self, key: &[K]) -> TrieResult<&V>
143    where
144        K: Ord,
145    {
146        match self {
147            Self::Branch { keys, children } => {
148                if let Some((head, tail)) = key.split_first() {
149                    if let Ok(index) = keys.binary_search(head) {
150                        children[index].get(tail)
151                    } else {
152                        TrieResult::Missing
153                    }
154                } else {
155                    TrieResult::Partial
156                }
157            }
158            Self::Leaf { value } => {
159                if key.is_empty() {
160                    TrieResult::Found(value)
161                } else {
162                    TrieResult::Missing
163                }
164            }
165        }
166    }
167
168    #[must_use]
169    fn get_mut(&mut self, key: &[K]) -> TrieResult<&mut V>
170    where
171        K: Ord,
172    {
173        match self {
174            Self::Branch { keys, children } => {
175                if let Some((head, tail)) = key.split_first() {
176                    if let Ok(index) = keys.binary_search(head) {
177                        children[index].get_mut(tail)
178                    } else {
179                        TrieResult::Missing
180                    }
181                } else {
182                    TrieResult::Partial
183                }
184            }
185            Self::Leaf { value } => {
186                if key.is_empty() {
187                    TrieResult::Found(value)
188                } else {
189                    TrieResult::Missing
190                }
191            }
192        }
193    }
194
195    fn collect_pairs<'a, const N: usize>(
196        &'a self,
197        prefix: &mut ArrayVec<K, N>,
198        pairs: &mut Vec<(ArrayVec<K, N>, &'a V)>,
199    ) where
200        K: Clone,
201    {
202        match self {
203            Self::Branch { keys, children } => {
204                for (key, child) in iter::zip(keys, children) {
205                    prefix.push(key.clone());
206                    child.collect_pairs(prefix, pairs);
207                    prefix.pop();
208                }
209            }
210            Self::Leaf { value } => {
211                pairs.push((prefix.clone(), value));
212            }
213        }
214    }
215}
216
217impl<K, V> Default for Node<K, V> {
218    fn default() -> Self {
219        Self::Branch {
220            keys: Vec::new(),
221            children: Vec::new(),
222        }
223    }
224}
225
226#[derive(Debug, PartialEq)]
227enum TrieResult<T> {
228    Missing,
229    Partial,
230    Found(T),
231}
232
233struct Entry<V> {
234    value: V,
235    order: u16,
236    description: Option<&'static str>,
237}
238
239pub struct Keymap<V> {
240    mappings: Trie<Key, Entry<V>>,
241    next_order: u16,
242    fallback: for<'a> fn(&'a Self, &[Key]) -> KeymapResult<'a, V>,
243    fallback_description: Option<&'static str>,
244}
245
246impl<V> Keymap<V> {
247    #[must_use]
248    pub fn new() -> Self {
249        Self {
250            mappings: Trie::new(),
251            next_order: 0,
252            fallback: |_, _| KeymapResult::Unmapped,
253            fallback_description: None,
254        }
255    }
256
257    pub fn insert(&mut self, keys: &str, value: V, description: Option<&'static str>) -> Option<V> {
258        let mut keys: Keys = keys.parse().unwrap();
259        for key in &mut keys.0 {
260            key.normalize();
261        }
262        if let TrieResult::Found(entry) = self.mappings.get_mut(&keys.0) {
263            // Keep original place in order, like `IndexMap::insert`
264            entry.description = description;
265            return Some(mem::replace(&mut entry.value, value));
266        }
267        let order = self.next_order;
268        self.next_order = order.checked_add(1).expect("wow that's a lotta keys");
269        self.mappings.insert(
270            &keys.0,
271            Entry {
272                value,
273                order,
274                description,
275            },
276        );
277        None
278    }
279
280    pub fn remove(&mut self, keys: &str) -> Option<V> {
281        let mut keys: Keys = keys.parse().unwrap();
282        for key in &mut keys.0 {
283            key.normalize();
284        }
285        self.mappings.remove(&keys.0).map(|entry| entry.value)
286    }
287
288    pub fn set_fallback(
289        &mut self,
290        fallback: for<'a> fn(&'a Self, &[Key]) -> KeymapResult<'a, V>,
291        description: Option<&'static str>,
292    ) {
293        self.fallback = fallback;
294        self.fallback_description = description;
295    }
296
297    #[must_use]
298    pub fn get(&self, keys: &str) -> KeymapResult<'_, V> {
299        let mut keys: Keys = keys.parse().unwrap();
300        for key in &mut keys.0 {
301            key.normalize();
302        }
303        match self.mappings.get(&keys.0) {
304            TrieResult::Missing => (self.fallback)(self, &keys.0),
305            TrieResult::Partial => KeymapResult::Pending,
306            TrieResult::Found(entry) => KeymapResult::Mapped(&entry.value),
307        }
308    }
309
310    #[must_use]
311    pub fn get_keys(&self, keys: &[Key]) -> KeymapResult<'_, V> {
312        match self.mappings.get(keys) {
313            TrieResult::Missing => (self.fallback)(self, keys),
314            TrieResult::Partial => KeymapResult::Pending,
315            TrieResult::Found(entry) => KeymapResult::Mapped(&entry.value),
316        }
317    }
318
319    fn entries(&self) -> impl Iterator<Item = (ArrayVec<Key, 3>, &Entry<V>)> {
320        let mut pairs: Vec<_> = self.mappings.iter::<3>().collect();
321        pairs.sort_by_key(|(_, entry)| entry.order);
322        pairs.into_iter()
323    }
324
325    pub fn iter(&self) -> impl Iterator<Item = (ArrayVec<Key, 3>, &V)> + '_ {
326        self.entries().map(|(keys, entry)| (keys, &entry.value))
327    }
328
329    pub fn docs(
330        &self,
331    ) -> impl Iterator<Item = (Keys<ArrayVec<Key, 3>>, Option<&'static str>)> + '_ {
332        self.entries()
333            .map(|(keys, entry)| (Keys(keys), entry.description))
334    }
335
336    #[must_use]
337    pub fn fallback_description(&self) -> Option<&'static str> {
338        self.fallback_description
339    }
340}
341
342impl<'a, V> IntoIterator for &'a Keymap<V> {
343    type Item = (ArrayVec<Key, 3>, &'a V);
344    type IntoIter = std::vec::IntoIter<(ArrayVec<Key, 3>, &'a V)>;
345    fn into_iter(self) -> Self::IntoIter {
346        self.iter().collect::<Vec<_>>().into_iter()
347    }
348}
349
350impl<V> Default for Keymap<V> {
351    fn default() -> Self {
352        Self::new()
353    }
354}
355
356macro_rules! keymap {
357    // [x] type [x] fallback (base case)
358    ($value_ty:ty; $($keys:literal => $value:expr, $(#[doc = $desc:literal])?)* $fallback_arg:ident => $fallback_body:block $(,)? $(#[doc = $fallback_desc:literal])?) => {{
359        let mut keymap: $crate::keymap::Keymap<$value_ty> = $crate::keymap::Keymap::new();
360        $(keymap.insert($keys, $value, None $(.or(Some($desc.trim())))?);)*
361        keymap.set_fallback(|_self, $fallback_arg| $fallback_body, None $(.or(Some($fallback_desc.trim())))?);
362        keymap
363    }};
364    // [x] type [ ] fallback
365    ($value_ty:ty; $($keys:literal => $value:expr, $(#[doc = $desc:literal])?)*) => {
366        $crate::keymap::keymap! { $value_ty;
367            $($keys => $value, $(#[doc = $desc])?)*
368            __keys => { $crate::keymap::KeymapResult::Unmapped },
369        }
370    };
371    // [ ] type [x] fallback
372    ($($keys:literal => $value:expr, $(#[doc = $desc:literal])?)* $fallback_arg:ident => $fallback_body:block $(,)? $(#[doc = $fallback_desc:literal])?) => {
373        $crate::keymap::keymap! { _;
374            $($keys => $value, $(#[doc = $desc])?)*
375            $fallback_arg => $fallback_body, $(#[doc = $fallback_desc])?
376        }
377    };
378    // [ ] type [ ] fallback
379    ($($keys:literal => $value:expr, $(#[doc = $desc:literal])?)*) => {
380        $crate::keymap::keymap! { _;
381            $($keys => $value, $(#[doc = $desc])?)*
382            __keys => { $crate::keymap::KeymapResult::Unmapped },
383        }
384    };
385}
386
387pub(crate) use keymap;
388
389#[derive(Debug, PartialEq)]
390pub enum KeymapResult<'a, V> {
391    Unmapped,
392    Pending,
393    Mapped(&'a V),
394    Fallback(V),
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400
401    #[derive(Clone, Debug, PartialEq)]
402    struct Action(&'static str);
403
404    #[test]
405    fn trie_works() {
406        let mut trie: Trie<&str, &str> = Trie::new();
407
408        trie.insert(&["1", "+", "2"], "3");
409        trie.insert(&["1", "+", "9"], "10");
410
411        assert_eq!(trie.get(&["1"]), TrieResult::Partial);
412        assert_eq!(trie.get(&["1", "+"]), TrieResult::Partial);
413        assert_eq!(trie.get(&["1", "+", "2"]), TrieResult::Found(&"3"));
414        assert_eq!(trie.get(&["1", "+", "9"]), TrieResult::Found(&"10"));
415        assert_eq!(trie.get(&["1", "+", "5"]), TrieResult::Missing);
416        assert_eq!(trie.get(&["1", "+", "2", "+", "3"]), TrieResult::Missing);
417
418        // remove returns the value, missing keys return None
419        assert_eq!(trie.remove(&["1", "+", "5"]), None);
420        assert_eq!(trie.remove(&["1"]), None);
421        assert_eq!(trie.remove(&["1", "+", "2"]), Some("3"));
422
423        // after removal: exact key is gone, sibling and its parent are intact
424        assert_eq!(trie.get(&["1", "+", "2"]), TrieResult::Missing);
425        assert_eq!(trie.get(&["1", "+", "9"]), TrieResult::Found(&"10"));
426        assert_eq!(trie.get(&["1"]), TrieResult::Partial);
427
428        // removing the last entry prunes empty branches
429        assert_eq!(trie.remove(&["1", "+", "9"]), Some("10"));
430        assert_eq!(trie.get(&["1"]), TrieResult::Missing);
431    }
432
433    #[test]
434    #[should_panic]
435    fn trie_forbids_overlap() {
436        let mut trie: Trie<&str, ()> = Trie::new();
437        trie.insert(&["1"], ());
438        trie.insert(&["1", "2"], ());
439    }
440
441    #[test]
442    fn keymap_get() {
443        let keymap = keymap! {
444            "gj" => vec![Action("move to bottom")],
445        };
446        assert_eq!(keymap.get("x"), KeymapResult::Unmapped);
447        assert_eq!(keymap.get("g"), KeymapResult::Pending);
448        assert_eq!(
449            keymap.get("gj"),
450            KeymapResult::Mapped(&vec![Action("move to bottom")])
451        );
452    }
453
454    #[test]
455    fn keymap_remove() {
456        let mut keymap = keymap! {
457            "gj" => vec![Action("move to bottom")],
458            "gk" => vec![Action("move to top")],
459        };
460        assert_eq!(
461            keymap.get("gj"),
462            KeymapResult::Mapped(&vec![Action("move to bottom")])
463        );
464        assert_eq!(keymap.remove("gj"), Some(vec![Action("move to bottom")]));
465        assert_eq!(keymap.get("gj"), KeymapResult::Unmapped);
466        // sibling still present
467        assert_eq!(
468            keymap.get("gk"),
469            KeymapResult::Mapped(&vec![Action("move to top")])
470        );
471        // missing key returns None
472        assert_eq!(keymap.remove("x"), None);
473    }
474
475    #[test]
476    fn keymap_fallback() {
477        let keymap = keymap! {
478            "gj" => vec![Action("move to bottom")],
479            keys => {
480                if keys.len() == 4 {
481                    KeymapResult::Fallback(vec![Action("yay 4")])
482                } else {
483                    KeymapResult::Unmapped
484                }
485            } /// Yay 4
486        };
487        assert_eq!(keymap.fallback_description(), Some("Yay 4"));
488        assert_eq!(
489            keymap.get("gj"),
490            KeymapResult::Mapped(&vec![Action("move to bottom")])
491        );
492        assert_eq!(keymap.get("123"), KeymapResult::Unmapped);
493        assert_eq!(
494            keymap.get("1234"),
495            KeymapResult::Fallback(vec![Action("yay 4")])
496        );
497    }
498
499    #[test]
500    #[should_panic]
501    fn keymap_forbids_overlap() {
502        let _keymap = keymap! {
503            "g" => vec![Action("enter goto mode")],
504            "gj" => vec![Action("move to bottom")],
505        };
506    }
507
508    #[test]
509    fn keymap_iter() {
510        let mut keymap = keymap! {
511            "gk" => vec![Action("move to top")],
512            "x" => vec![Action("select line")],
513            "gj" => vec![Action("move to bottom")],
514        };
515        let values = |keymap: &Keymap<Vec<Action>>| {
516            keymap
517                .iter()
518                .map(|(keys, value)| (Keys(keys).to_string(), value[0].0))
519                .collect::<Vec<_>>()
520        };
521        // insertion order, not trie order
522        assert_eq!(
523            values(&keymap),
524            [
525                (String::from("gk"), "move to top"),
526                (String::from("x"), "select line"),
527                (String::from("gj"), "move to bottom"),
528            ]
529        );
530        // reinserting keeps the original position
531        keymap.insert("gk", vec![Action("goto top")], None);
532    }
533
534    #[test]
535    fn keymap_descriptions() {
536        let keymap = keymap! {
537            "gk" => (), /// Move to top
538            "x" => (), // this isn't parsed as a description, because it isn't a doc comment
539        };
540        let descriptions: Vec<_> = keymap.docs().map(|(_, description)| description).collect();
541        assert_eq!(descriptions, [Some("Move to top"), None]);
542        assert_eq!(keymap.fallback_description(), None);
543    }
544}
indigo_core/keymap.rs

indigo_core/
keymap.rs
