Skip to main content

asciimath_parser/prefix_map/
trie.rs

1use super::PrefixMap;
2use qp_trie::Trie;
3use std::borrow::Borrow;
4
5#[derive(Debug, PartialEq, Eq, Clone)]
6struct Wrapper<K>(K);
7
8impl<K> Borrow<[u8]> for Wrapper<K>
9where
10    K: Borrow<str>,
11{
12    fn borrow(&self) -> &[u8] {
13        self.0.borrow().as_bytes()
14    }
15}
16
17/// A prefix map backed by a qp-trie
18///
19/// This is the [fastest][crate::prefix_map] and default token map. This prefix map requires the `qp-trie`
20/// feature (enabled by default). This takes time `O(longest_prefix)`.
21///
22/// # Example
23/// ```
24/// use asciimath_parser::prefix_map::QpTriePrefixMap;
25/// use asciimath_parser::{ASCIIMATH_TOKENS};
26///
27/// let token_map = QpTriePrefixMap::from_iter(ASCIIMATH_TOKENS);
28/// ```
29#[derive(Debug, Clone)]
30pub struct QpTriePrefixMap<K: Clone, V>(Trie<Wrapper<K>, V>);
31
32impl<K, V> FromIterator<(K, V)> for QpTriePrefixMap<K, V>
33where
34    K: Borrow<str> + Clone,
35{
36    fn from_iter<I>(iter: I) -> Self
37    where
38        I: IntoIterator<Item = (K, V)>,
39    {
40        QpTriePrefixMap(
41            iter.into_iter()
42                .map(|(key, val)| (Wrapper(key), val))
43                .collect(),
44        )
45    }
46}
47
48impl<K, V> PrefixMap<V> for QpTriePrefixMap<K, V>
49where
50    K: Borrow<str> + Clone,
51{
52    fn get_longest_prefix<P: AsRef<str>>(&self, inp: P) -> Option<(usize, &V)> {
53        let bytes = inp.as_ref().as_bytes();
54        let empty = &[][..];
55        let mut subtrie = self.0.subtrie(empty);
56        let mut res = subtrie.get(empty).map(|val| (0, val));
57        for len in 1..=bytes.len() {
58            let slice = &bytes[..len];
59            subtrie = subtrie.subtrie(slice);
60            if subtrie.is_empty() {
61                break;
62            } else if let Some(val) = subtrie.get(slice) {
63                res = Some((len, val));
64            }
65        }
66        res
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::{PrefixMap, QpTriePrefixMap};
73
74    #[test]
75    fn correct_prefixes() {
76        let map = QpTriePrefixMap::from_iter([("a", 0), ("abc", 1), ("bc", 2), ("bc", 3)]);
77        assert_eq!(map.get_longest_prefix("abcd"), Some((3, &1)));
78        assert_eq!(map.get_longest_prefix("ab"), Some((1, &0)));
79        assert_eq!(map.get_longest_prefix("bcd"), Some((2, &3)));
80        assert_eq!(map.get_longest_prefix("bd"), None);
81        assert_eq!(map.get_longest_prefix("💖"), None);
82    }
83
84    #[test]
85    fn works_for_perverse() {
86        let map = QpTriePrefixMap::from_iter([("", 0), (" 3", 1)]);
87        assert_eq!(map.get_longest_prefix(" 3 "), Some((2, &1)));
88        assert_eq!(map.get_longest_prefix("ab"), Some((0, &0)));
89    }
90}