ui-agregator/crates/evil-keys/src/which_key.rs

use crate::trie::{KeyTrie, KeyTrieNode};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WhichKeyEntry {
    pub key: String,
    pub description: String,
    pub is_group: bool,
}

impl<A> KeyTrieNode<A> {
    pub fn which_key_entries(&self) -> Vec<WhichKeyEntry> {
        let mut groups = Vec::new();
        let mut leaves = Vec::new();

        for (key, trie) in &self.map {
            match trie {
                KeyTrie::Node(node) => {
                    groups.push(WhichKeyEntry {
                        key: key.to_string(),
                        description: node.name.clone(),
                        is_group: true,
                    });
                }
                KeyTrie::Leaf(leaf) => {
                    leaves.push(WhichKeyEntry {
                        key: key.to_string(),
                        description: leaf.description.clone().unwrap_or_default(),
                        is_group: false,
                    });
                }
            }
        }

        groups.sort_by(|a, b| a.key.cmp(&b.key));
        leaves.sort_by(|a, b| a.key.cmp(&b.key));
        groups.extend(leaves);
        groups
    }
}

#[cfg(test)]
mod tests {
    use crate::trie::KeyTrie;

    #[test]
    fn test_empty_node() {
        let trie: KeyTrie<&str> = KeyTrie::new("root");
        if let KeyTrie::Node(node) = trie {
            let entries = node.which_key_entries();
            assert!(entries.is_empty());
        }
    }

    #[test]
    fn test_leaves_only_sorted() {
        let mut trie: KeyTrie<&str> = KeyTrie::new("root");
        trie.bind_desc("c", "action_c", "C action").unwrap();
        trie.bind_desc("a", "action_a", "A action").unwrap();
        trie.bind_desc("b", "action_b", "B action").unwrap();

        if let KeyTrie::Node(node) = trie {
            let entries = node.which_key_entries();
            assert_eq!(entries.len(), 3);
            assert_eq!(entries[0].key, "a");
            assert_eq!(entries[1].key, "b");
            assert_eq!(entries[2].key, "c");
        }
    }

    #[test]
    fn test_groups_before_leaves() {
        let mut trie: KeyTrie<&str> = KeyTrie::new("root");
        trie.bind_desc("z", "action", "Z action").unwrap();
        trie.group("g", "goto", |node| {
            node.bind("g", "goto_top")?;
            Ok(())
        })
        .unwrap();
        trie.bind_desc("a", "action", "A action").unwrap();

        if let KeyTrie::Node(node) = trie {
            let entries = node.which_key_entries();
            assert_eq!(entries.len(), 3);
            assert!(entries[0].is_group);
            assert_eq!(entries[0].key, "g");
            assert!(!entries[1].is_group);
            assert_eq!(entries[1].key, "a");
            assert!(!entries[2].is_group);
            assert_eq!(entries[2].key, "z");
        }
    }

    #[test]
    fn test_descriptions_populated() {
        let mut trie: KeyTrie<&str> = KeyTrie::new("root");
        trie.bind_desc("j", "down", "Move down").unwrap();
        trie.group("g", "goto", |_| Ok(())).unwrap();

        if let KeyTrie::Node(node) = trie {
            let entries = node.which_key_entries();

            let group = entries.iter().find(|e| e.key == "g").unwrap();
            assert_eq!(group.description, "goto");
            assert!(group.is_group);

            let leaf = entries.iter().find(|e| e.key == "j").unwrap();
            assert_eq!(leaf.description, "Move down");
            assert!(!leaf.is_group);
        }
    }
}