radix_tree/
lib.rs

1extern crate alloc;
2
3use alloc::vec::Vec;
4use core::mem;
5
6const fn pos<K>(l: &usize, _: &K, _: &Vec<K>) -> usize {
7    *l
8}
9
10pub trait Vectorable<K>
11where
12    K: Copy + PartialEq + PartialOrd,
13{
14    fn into(&self) -> Vec<K>;
15}
16
17#[macro_use]
18pub mod macros;
19
20pub trait Radix<K, V, P: Vectorable<K>>
21where
22    K: Copy + PartialEq + PartialOrd,
23    V: Clone,
24{
25    fn remove(&mut self, path: P);
26    fn insert(&mut self, path: P, data: V) -> &mut Self;
27    fn find(&self, path: P) -> Option<&Self>;
28    fn add_node(&mut self, path: P, data: V) -> &mut Self;
29    fn find_node(&self, path: P) -> Option<&Self>;
30}
31
32#[derive(Debug, Clone, PartialEq)]
33pub struct Node<K, V> {
34    pub path: Vec<K>,
35    pub data: Option<V>,
36    pub indices: Vec<K>,
37    pub nodes: Vec<Node<K, V>>,
38}
39
40impl<K, V> Node<K, V>
41where
42    K: Copy + PartialEq + PartialOrd,
43    V: Clone,
44{
45    pub fn new<P: Vectorable<K>>(path: P, data: Option<V>) -> Self {
46        Node {
47            data,
48            path: (&path).into(),
49            nodes: Vec::new(),
50            indices: Vec::new(),
51        }
52    }
53
54    pub fn insert_with<F>(
55        &mut self,
56        path: &mut Vec<K>,
57        data: Option<V>,
58        force_update: bool,
59        pos: F,
60    ) -> &mut Self
61    where
62        F: Fn(&usize, &K, &Vec<K>) -> usize,
63    {
64        let pl = path.len();
65        let sl = self.path.len();
66
67        // empty input path
68        if 0 == pl {
69            if force_update {
70                self.data = data;
71            }
72            return self;
73        }
74
75        // empty node
76        if 0 == sl && 0 == self.indices.len() {
77            if force_update {
78                self.data = data;
79            }
80            self.path = path.to_owned();
81            return self;
82        }
83
84        // pl > 0 && sl >= 0
85        let max = pl.min(sl);
86        let mut i = 0;
87        while i < max && path[i] == self.path[i] {
88            i += 1;
89        }
90
91        if i < sl {
92            let child = Node {
93                data: self.data.take(),
94                path: self.path.split_off(i),
95                nodes: mem::replace(&mut self.nodes, Vec::new()),
96                indices: mem::replace(&mut self.indices, Vec::new()),
97            };
98            let c = child.path[0];
99            let index = pos(&self.indices.len(), &c, &self.indices);
100            self.indices.insert(index, c);
101            self.nodes.insert(index, child);
102
103            // self.indices.push(child.path[0]);
104            // self.nodes.push(child);
105        }
106
107        if i == pl {
108            if force_update {
109                self.data = data;
110            }
111            return self;
112        }
113
114        self.add_node_with(path, data, i, force_update, pos)
115    }
116
117    pub fn add_node_with<F>(
118        &mut self,
119        path: &mut Vec<K>,
120        data: Option<V>,
121        i: usize,
122        force_update: bool,
123        pos: F,
124    ) -> &mut Self
125    where
126        F: Fn(&usize, &K, &Vec<K>) -> usize,
127    {
128        let l = self.indices.len();
129        let c = path[i];
130        let mut j = 0;
131        while j < l {
132            if c == self.indices[j] {
133                return self.nodes[j].insert_with(&mut path.split_off(i), data, force_update, pos);
134            }
135            j += 1;
136        }
137
138        let index = pos(&l, &c, &self.indices);
139        self.indices.insert(index, c);
140        self.nodes.insert(
141            index,
142            Node {
143                data,
144                nodes: Vec::new(),
145                indices: Vec::new(),
146                path: path.split_off(i),
147            },
148        );
149
150        &mut self.nodes[index]
151
152        // self.indices.push(c);
153        // self.nodes.push(Node {
154        //     data,
155        //     nodes: Vec::new(),
156        //     indices: Vec::new(),
157        //     path: path.split_off(i),
158        // });
159        // &mut self.nodes[l]
160    }
161
162    pub fn find_with(&self, path: &mut Vec<K>) -> Option<&Self> {
163        let pl = path.len();
164        let sl = self.path.len();
165
166        // "abc" < "abcde"
167        // not found
168        if pl < sl {
169            return None;
170        }
171
172        // "abcde" > "abc" or "abc" == "abc"
173        let mut i = 0;
174        while i < sl && path[i] == self.path[i] {
175            i += 1;
176        }
177
178        // "abc" == "abc"
179        if pl == sl {
180            if i == pl {
181                return Some(self);
182            }
183            // not found
184            return None;
185        }
186
187        // "abcde" > "abc"
188        self.find_node_with(path, i)
189    }
190
191    pub fn find_node_with(&self, path: &mut Vec<K>, i: usize) -> Option<&Self> {
192        let l = self.indices.len();
193        let c = path[i];
194        let mut j = 0;
195        while j < l {
196            if c == self.indices[j] {
197                return self.nodes[j].find_with(&mut path.split_off(i));
198            }
199            j += 1;
200        }
201        // not found
202        None
203    }
204}
205
206impl<K, V, P: Vectorable<K>> Radix<K, V, P> for Node<K, V>
207where
208    K: Copy + PartialEq + PartialOrd,
209    V: Clone,
210{
211    #[allow(unused_variables)]
212    fn remove(&mut self, path: P) {}
213
214    fn insert(&mut self, path: P, data: V) -> &mut Self {
215        self.insert_with(&mut (&path).into(), Some(data), true, pos)
216    }
217
218    fn find(&self, path: P) -> Option<&Self> {
219        self.find_with(&mut (&path).into())
220    }
221
222    fn add_node(&mut self, path: P, data: V) -> &mut Self {
223        self.add_node_with(&mut (&path).into(), Some(data), 0, true, pos)
224    }
225
226    fn find_node(&self, path: P) -> Option<&Self> {
227        self.find_node_with(&mut (&path).into(), 0)
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    macro_rules! find {
236        ($tree:expr, $($path:expr, $data:expr),*,) => {{
237            $(
238                let node = $tree.find($path);
239                assert_eq!(node.is_some(), $data);
240                if node.is_some() {
241                    assert_eq!(node.unwrap().data.unwrap(), $data);
242                }
243            )*
244        }};
245    }
246
247    macro_rules! insert_and_find {
248        ($tree:expr, $($path:expr, $data:expr),*,) => {{
249            $(
250                $tree.insert($path, $data);
251                find!($tree, $path, $data,);
252            )*
253        }};
254    }
255
256    #[test]
257    fn new_any_type_node() {
258        let node = Node::<u8, &str>::new("Hello world!", Some("a"));
259        assert_eq!(
260            node.path,
261            vec![72, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100, 33]
262        );
263        assert_eq!(node.data.unwrap(), "a");
264
265        let node = Node::<u8, &str>::new("Hello 世界！", Some("a"));
266        assert_eq!(
267            node.path,
268            vec![72, 101, 108, 108, 111, 32, 228, 184, 150, 231, 149, 140, 239, 188, 129]
269        );
270        assert_eq!(node.data.unwrap(), "a");
271
272        let node = Node::<char, &str>::new("Hello 世界！", Some("a"));
273        assert_eq!(
274            node.path,
275            vec!['H', 'e', 'l', 'l', 'o', ' ', '世', '界', '！']
276        );
277        assert_eq!(node.data.unwrap(), "a");
278
279        let node = Node::<char, u32>::new("你好，世界！", Some(0));
280        assert_eq!(node.path, vec!['你', '好', '，', '世', '界', '！']);
281        assert_eq!(node.data.unwrap(), 0);
282
283        let node = Node::<u8, u8>::new("abcde", Some(1));
284        assert_eq!(node.path, vec![97, 98, 99, 100, 101]);
285        assert_eq!(node.data.unwrap(), 1);
286
287        let node = Node::new("abcde".as_bytes().to_vec(), Some(97));
288        assert_eq!(node.path, vec![97, 98, 99, 100, 101]);
289        assert_eq!(node.data.unwrap(), 97);
290
291        let node = Node::new("abcde".as_bytes(), Some(97));
292        assert_eq!(node.path, vec![97, 98, 99, 100, 101]);
293        assert_eq!(node.data.unwrap(), 97);
294    }
295
296    #[test]
297    fn node_insert_and_find() {
298        let mut node = Node::<char, bool>::new("你好，世界！", Some(true));
299        node.add_node("Rust", true);
300
301        let n1 = node.find_node("Rust");
302        let n2 = node.find("你好，世界！Rust");
303        assert_eq!(n1, n2);
304    }
305
306    #[test]
307    fn node_insert_then_return_new_node() {
308        let mut tree = Node::<u8, u8>::new("", Some(b' '));
309
310        let a = tree.insert("a", b'a');
311        let b = a.add_node("b", b'b');
312        let c = b.add_node("c", b'c');
313        let d = c.add_node("d", b'd');
314        let _ = d.add_node("e", b'e');
315
316        // println!("{:#?}", tree);
317
318        let node = tree.find("a");
319        assert_eq!(node.is_some(), true);
320        let a = node.unwrap();
321        assert_eq!(a.data.unwrap(), b'a');
322
323        let node = a.find_node("b");
324        assert_eq!(node.is_some(), true);
325        let b = node.unwrap();
326        assert_eq!(b.data.unwrap(), b'b');
327
328        let node = b.find_node("c");
329        assert_eq!(node.is_some(), true);
330        let c = node.unwrap();
331        assert_eq!(c.data.unwrap(), b'c');
332
333        let node = c.find_node("d");
334        assert_eq!(node.is_some(), true);
335        let d = node.unwrap();
336        assert_eq!(d.data.unwrap(), b'd');
337
338        let node = d.find_node("e");
339        assert_eq!(node.is_some(), true);
340        let e = node.unwrap();
341        assert_eq!(e.data.unwrap(), b'e');
342
343        let node = a.find("abcde");
344        assert_eq!(node.is_some(), true);
345        assert_eq!(node.unwrap().data.unwrap(), b'e');
346
347        let node = tree.find("abcdef");
348        assert_eq!(node.is_some(), false);
349
350        let node = tree.find("b");
351        assert_eq!(node.is_some(), false);
352    }
353
354    #[test]
355    fn new_tree() {
356        let mut tree = Node::<char, bool>::new("", Some(false));
357
358        insert_and_find!(
359            tree,
360            "alligator",
361            true,
362            "alien",
363            true,
364            "baloon",
365            true,
366            "chromodynamic",
367            true,
368            "romane",
369            true,
370            "romanus",
371            true,
372            "romulus",
373            true,
374            "rubens",
375            true,
376            "ruber",
377            true,
378            "rubicon",
379            true,
380            "rubicundus",
381            true,
382            "all",
383            true,
384            "rub",
385            true,
386            "ba",
387            true,
388            "你好，世界",
389            true,
390            "你好",
391            true,
392            "你",
393            true,
394        );
395
396        find!(
397            tree, "rpxxx", false, "chro", false, "chromz", false, "zorro", false, "ro", false,
398            "zo", false,
399        );
400
401        let node = tree.find("");
402        assert_eq!(node.is_some(), true);
403        assert_eq!(node.unwrap().data, None);
404
405        tree.insert("", false);
406        let node = tree.find("");
407        assert_eq!(node.is_some(), true);
408        assert_eq!(node.unwrap().data.unwrap(), false);
409
410        let node = tree.find("all");
411        assert_eq!(node.is_some(), true);
412        assert_eq!(node.unwrap().data.unwrap(), true);
413
414        let node = tree.find("dota2");
415        assert_eq!(node.is_none(), true);
416
417        let node = tree.find("你");
418        assert_eq!(node.is_some(), true);
419        assert_eq!(node.unwrap().data.unwrap(), true);
420
421        let node = tree.find("你好");
422        assert_eq!(node.is_some(), true);
423        assert_eq!(node.unwrap().data.unwrap(), true);
424
425        let node = tree.find("语言");
426        assert_eq!(node.is_some(), false);
427
428        let node = tree.find("你好，世界");
429        assert_eq!(node.is_some(), true);
430
431        let node = tree.find("你好，世界 Rust");
432        assert_eq!(node.is_some(), false);
433
434        println!("{:#?}", tree);
435    }
436
437    #[test]
438    fn clone_node() {
439        let mut node = Node::<char, bool>::new("", Some(false));
440        let mut node2 = node.clone();
441        assert_eq!(node, node2);
442
443        node.add_node("/", true);
444        node2.add_node("/", true);
445        assert_eq!(node, node2);
446
447        #[derive(Debug, Clone, PartialEq)]
448        struct NodeMetadata {
449            is_key: bool,
450            params: Option<Vec<&'static str>>,
451        }
452
453        let mut node = Node::<char, NodeMetadata>::new(
454            "/",
455            Some(NodeMetadata {
456                is_key: false,
457                params: Some(vec![]),
458            }),
459        );
460        let mut node2 = node.clone();
461        assert_eq!(node, node2);
462
463        node.add_node(
464            "users",
465            NodeMetadata {
466                is_key: true,
467                params: Some(vec!["tree"]),
468            },
469        );
470        node2.add_node(
471            "users",
472            NodeMetadata {
473                is_key: true,
474                params: Some(vec!["tree"]),
475            },
476        );
477        assert_eq!(node, node2);
478    }
479}
radix_tree/lib.rs

radix_tree/
lib.rs
