sql_fun_sqlast/
trie_map.rs

1//! A trie-based map used internally by `sql-fun` for efficient prefix-ordered key storage.
2//!
3//! This structure is primarily used to store schema-qualified
4//! SQL object names or generated query identifiers in a prefix-compressed
5//! and lexically sortable form.
6//!
7//! This structure stores key-value pairs in lexicographic order using a trie,
8//! allowing for efficient prefix-based lookups, ordered iteration, and
9//! storage of hierarchical keys like schema paths or query identifiers.
10//!
11//! Keys are byte strings (`&str`), stored and compared by UTF-8 byte order.
12//!
13//! Internally, the trie is implemented as a tree of `TrieNode`s,
14//! with shared prefix compression and stable iteration order.
15//! Empty-string keys are supported via a separate storage slot.
16//!
17//! This module is not part of the public API, but is exposed through source code
18//! and licensed accordingly. Users may extract and reuse this implementation
19//! under the terms of the `sql-fun` license.
20
21mod key_fragment;
22mod leaf;
23mod node;
24
25use self::{
26    key_fragment::{BoxedKeyFragment, KeyFragment},
27    leaf::{TrieLeafIter, TrieLeafs},
28    node::{TrieNode, TrieNodeIter},
29};
30
31/// Map string to value
32///
33#[derive(Debug)]
34pub struct TrieMap<TValue> {
35    /// Value for the empty key
36    empty_node: Option<TValue>,
37    /// Top-level nodes
38    nodes: TrieLeafs<TValue>,
39}
40
41impl<TValue> Default for TrieMap<TValue> {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl<TValue> TrieMap<TValue> {
48    /// create a empty instance
49    #[must_use]
50    pub fn new() -> Self {
51        Self {
52            empty_node: None,
53            nodes: TrieLeafs::new(),
54        }
55    }
56
57    /// insert value to key
58    pub fn insert(&mut self, key: &str, value: TValue) -> Option<TValue> {
59        if key.is_empty() {
60            return self.empty_node.replace(value);
61        }
62        let key = KeyFragment::from(key.as_bytes());
63        let start_byte = key.start_byte();
64        if let Some(match_node) = self.nodes.find_by_start_byte_mut(start_byte) {
65            match_node.insert(&key, value)
66        } else {
67            let new_node = TrieNode::new(key, value);
68            self.nodes.insert(new_node);
69            None
70        }
71    }
72
73    /// get iterator
74    pub fn iter(&self) -> impl Iterator<Item = (String, &TValue)> {
75        TrieMapIter {
76            empty_node: self.empty_node.as_ref(),
77            leaf_iter: self.nodes.iter(Vec::default()),
78        }
79    }
80
81    /// get value by key
82    pub fn get(&self, key: &str) -> Option<&TValue> {
83        if key.is_empty() {
84            return self.empty_node.as_ref();
85        }
86        let mut frag = KeyFragment::from(key.as_bytes());
87        let mut node = self.nodes.find_by_fragment(&frag)?;
88        frag = node.split_key(&frag)?;
89        if frag.is_empty() {
90            return node.value.as_ref();
91        }
92
93        loop {
94            node = node.children.find_by_fragment(&frag)?;
95            frag = node.split_key(&frag)?;
96            if frag.is_empty() {
97                return node.value.as_ref();
98            }
99        }
100    }
101
102    /// get mutable reference by key
103    pub fn get_mut(&mut self, key: &str) -> Option<&mut TValue> {
104        if key.is_empty() {
105            return self.empty_node.as_mut();
106        }
107        let mut frag = KeyFragment::from(key.as_bytes());
108        let mut node = self.nodes.find_by_fragment_mut(&frag)?;
109        frag = node.split_key(&frag)?;
110        if frag.is_empty() {
111            return node.value.as_mut();
112        }
113
114        loop {
115            node = node.children.find_by_fragment_mut(&frag)?;
116            frag = node.split_key(&frag)?;
117            if frag.is_empty() {
118                return node.value.as_mut();
119            }
120        }
121    }
122
123    /// Returns the number of entries in the map.
124    pub fn len(&self) -> usize {
125        if self.empty_node.is_some() {
126            self.nodes.len() + 1
127        } else {
128            self.nodes.len()
129        }
130    }
131
132    /// returns true if empty
133    pub fn is_empty(&self) -> bool {
134        self.empty_node.is_none() && self.nodes.is_empty()
135    }
136}
137
138struct TrieMapIter<'a, TValue> {
139    empty_node: Option<&'a TValue>,
140    leaf_iter: TrieLeafIter<'a, TValue>,
141}
142
143impl<'a, TValue> Iterator for TrieMapIter<'a, TValue> {
144    type Item = (String, &'a TValue);
145
146    fn next(&mut self) -> Option<Self::Item> {
147        if let Some(empty_node) = self.empty_node.take() {
148            return Some((String::new(), empty_node));
149        }
150        let (fragments, item) = self.leaf_iter.next()?;
151
152        // safe to use from_utf8_unchecked, but unsafe
153        Some((KeyFragment::to_string_key(&fragments).unwrap(), item))
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::TrieMap;
160
161    #[test]
162    fn test_insert_and_get() {
163        let mut map = TrieMap::new();
164        assert_eq!(map.len(), 0);
165        assert!(map.get("key1").is_none());
166
167        map.insert("key1", 10);
168        assert_eq!(map.len(), 1);
169        assert_eq!(map.get("key1"), Some(&10));
170        assert!(map.get("key2").is_none());
171
172        // overwrite existing key
173        let old = map.insert("key1", 20);
174        assert_eq!(old, Some(10));
175        assert_eq!(map.get("key1"), Some(&20));
176    }
177
178    #[test]
179    fn test_empty_key() {
180        let mut map = TrieMap::new();
181        map.insert("", 5);
182        assert_eq!(map.len(), 1);
183        assert_eq!(map.get(""), Some(&5));
184        let old = map.insert("", 7);
185        assert_eq!(old, Some(5));
186        assert_eq!(map.get(""), Some(&7));
187    }
188
189    #[test]
190    fn test_iter_order() {
191        let mut map = TrieMap::new();
192        map.insert("b", 2);
193        map.insert("a", 1);
194        map.insert("aa", 11);
195        map.insert("ab", 12);
196        let items: Vec<_> = map.iter().collect();
197        let keys: Vec<String> = items.iter().map(|(k, _)| k.clone()).collect();
198        assert_eq!(
199            keys,
200            vec![
201                "a".to_string(),
202                "aa".to_string(),
203                "ab".to_string(),
204                "b".to_string()
205            ]
206        );
207        let values: Vec<i32> = items.iter().map(|(_, v)| **v).collect();
208        assert_eq!(values, vec![1, 11, 12, 2]);
209    }
210
211    #[test]
212    fn test_get_common_prefix_loop() {
213        let mut map = TrieMap::new();
214        map.insert("foobarbaz", 3);
215        map.insert("foo", 1);
216        map.insert("foobar", 2);
217        map.insert("fool_ploof_needed", 4);
218        assert_eq!(map.get("foo"), Some(&1));
219        assert_eq!(map.get("foobar"), Some(&2));
220        assert_eq!(map.get("foobarbaz"), Some(&3));
221        assert_eq!(map.get("fool_ploof_needed"), Some(&4));
222    }
223
224    #[test]
225    fn test_partial_split() {
226        let mut map = TrieMap::new();
227        // insert in normal order
228        assert_eq!(map.insert("abcd", 1), None);
229        assert_eq!(map.insert("abxy", 2), None);
230        assert_eq!(map.len(), 2);
231        assert_eq!(map.get("abcd"), Some(&1));
232        assert_eq!(map.get("abxy"), Some(&2));
233        // reverse order
234        let mut map2 = TrieMap::new();
235        assert_eq!(map2.insert("abxy", 20), None);
236        assert_eq!(map2.insert("abcd", 10), None);
237        assert_eq!(map2.len(), 2);
238        assert_eq!(map2.get("abcd"), Some(&10));
239        assert_eq!(map2.get("abxy"), Some(&20));
240        // add third partial key
241        assert_eq!(map.insert("abmn", 3), None);
242        assert_eq!(map.len(), 3);
243        assert_eq!(map.get("abmn"), Some(&3));
244    }
245
246    #[test]
247    fn test_iter_empty_key_first() {
248        let mut map = TrieMap::new();
249        // only empty key
250        map.insert("", 100);
251        assert_eq!(map.len(), 1);
252        let mut iter = map.iter();
253        assert_eq!(iter.next(), Some(("".to_string(), &100)));
254        assert_eq!(iter.next(), None);
255
256        // mix empty key with others in fresh map
257        let mut map2 = TrieMap::new();
258        map2.insert("", 100);
259        map2.insert("a", 1);
260        map2.insert("b", 2);
261        assert_eq!(map2.len(), 3);
262        let mut iter2 = map2.iter();
263        assert_eq!(iter2.next(), Some(("".to_string(), &100)));
264        let rest: Vec<_> = iter2.collect();
265        assert_eq!(rest, vec![("a".to_string(), &1), ("b".to_string(), &2)]);
266    }
267}
sql_fun_sqlast/trie_map.rs

sql_fun_sqlast/
trie_map.rs
