kanata_parser/
trie.rs

1//! Wrapper around a trie type for (hopefully) easier swapping of libraries if desired.
2
3use bytemuck::cast_slice;
4use patricia_tree::map::PatriciaMap;
5
6pub type TrieKeyElement = u16;
7
8#[derive(Debug, Clone)]
9pub struct Trie<T> {
10    inner: patricia_tree::map::PatriciaMap<T>,
11}
12
13#[derive(Debug, Copy, Clone, PartialEq, Eq)]
14pub enum GetOrDescendentExistsResult<T> {
15    NotInTrie,
16    InTrie,
17    HasValue(T),
18}
19
20use GetOrDescendentExistsResult::*;
21
22impl<T> Default for Trie<T> {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28fn key_len(k: impl AsRef<[u16]>) -> usize {
29    debug_assert!(std::mem::size_of::<TrieKeyElement>() == 2 * std::mem::size_of::<u8>());
30    k.as_ref().len() * 2
31}
32
33impl<T> Trie<T> {
34    pub fn new() -> Self {
35        Self {
36            inner: PatriciaMap::new(),
37        }
38    }
39
40    pub fn ancestor_exists(&self, key: impl AsRef<[u16]>) -> bool {
41        self.inner
42            .get_longest_common_prefix(cast_slice(key.as_ref()))
43            .is_some()
44    }
45
46    pub fn descendant_exists(&self, key: impl AsRef<[u16]>) -> bool {
47        // Length of the [u8] interpretation of the [u16] key is doubled.
48        self.inner
49            .longest_common_prefix_len(cast_slice(key.as_ref()))
50            == key_len(key)
51    }
52
53    pub fn insert(&mut self, key: impl AsRef<[u16]>, val: T) {
54        self.inner.insert(cast_slice(key.as_ref()), val);
55    }
56
57    pub fn get_or_descendant_exists(&self, key: impl AsRef<[u16]>) -> GetOrDescendentExistsResult<T>
58    where
59        T: Clone,
60    {
61        let mut descendants = self.inner.iter_prefix(cast_slice(key.as_ref()));
62        match descendants.next() {
63            None => NotInTrie,
64            Some(descendant) => {
65                if descendant.0.len() == key_len(key.as_ref()) {
66                    HasValue(descendant.1.clone())
67                } else {
68                    InTrie
69                }
70            }
71        }
72    }
73
74    pub fn is_empty(&self) -> bool {
75        self.inner.is_empty()
76    }
77}
kanata_parser/trie.rs

kanata_parser/
trie.rs
