├── .github └── workflows │ └── test.yaml ├── .gitignore ├── CHANGELOG.md ├── Cargo.toml ├── LICENSE ├── README.md ├── appveyor.yml ├── benches └── trie_benches.rs ├── data ├── 1984.txt └── sun-rising.txt ├── examples ├── child_iter.rs ├── debug.rs ├── opt.rs └── string_frequency.rs └── src ├── iter.rs ├── keys.rs ├── lib.rs ├── macros.rs ├── qc_test.rs ├── serde.rs ├── subtrie.rs ├── test.rs ├── traversal.rs ├── trie.rs ├── trie_common.rs └── trie_node.rs /.github/workflows/test.yaml: -------------------------------------------------------------------------------- 1 | name: Test suite 2 | on: 3 | push: 4 | branches: 5 | - master 6 | pull_request: 7 | jobs: 8 | test: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - uses: actions/checkout@v2 12 | - uses: actions-rs/toolchain@v1 13 | with: 14 | toolchain: stable 15 | - run: cargo test --release --all-features 16 | - run: cargo doc --no-deps --all-features 17 | - run: cargo fmt -- --check 18 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | Cargo.lock 3 | *.swp 4 | *.swo 5 | .DS_Store 6 | .cargo 7 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | Changelog 2 | ==== 3 | 4 | 0.2.1: 5 | * Implement `TrieKey` for vectors of integers (#63) 6 | * Released 22/10/2020 7 | 8 | 0.2.0: 9 | * Update `nibble_vec` to v0.1.0: faster and uses less memory! 10 | * Released 28/06/2020 11 | 12 | 0.1.6: 13 | 14 | * Update to Rust 2018, add benchmarking code (#52) 15 | * Appveyor CI for Windows (#42) 16 | * No API changes or bugfixes 17 | 18 | 0.1.5: 19 | 20 | * Fix another bug related to the removal of non-existent keys (#50) 21 | * Implement `Clone` for `Trie` 22 | 23 | 0.1.4: 24 | 25 | * Fix a panic that occurred when removing non-existent keys (#40) 26 | * Reformat the code using the latest `rustfmt` 27 | 28 | 0.1.3: 29 | 30 | * Add `prefix()` method for fetching the `NibbleVec` of a trie or subtrie 31 | * Update `nibble_vec` to v0.0.4 32 | * Make better use of lifetime elision (type signatures will look different, but are the same) 33 | 34 | 0.1.2: 35 | 36 | * Display README on crates.io (no code changes) 37 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "radix_trie" 3 | version = "0.2.1" 4 | description = "Generic radix trie data-structure." 5 | edition = "2018" 6 | license = "MIT" 7 | authors = ["Michael Sproul "] 8 | 9 | repository = "https://github.com/michaelsproul/rust_radix_trie" 10 | documentation = "https://docs.rs/radix_trie/" 11 | readme = "README.md" 12 | 13 | keywords = ["trie", "patricia", "collection", "generic", "prefix"] 14 | categories = ["data-structures", "text-processing"] 15 | 16 | [dependencies] 17 | nibble_vec = "0.1" 18 | endian-type = "0.1.2" 19 | serde = { version = "1.0", optional = true } 20 | 21 | [dev-dependencies] 22 | criterion = "0.3" 23 | quickcheck = "1.0" 24 | rand = "0.8" 25 | serde_test = "1.0" 26 | 27 | [[bench]] 28 | name = "trie_benches" 29 | harness = false 30 | 31 | [lib] 32 | bench = false 33 | 34 | [badges] 35 | appveyor = { repository = "michaelsproul/rust_radix_trie" } 36 | maintenance = { status = "as-is" } 37 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Michael Sproul 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | THE SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Rust Radix Trie 2 | ==== 3 | 4 | [![Windows Build Status](https://ci.appveyor.com/api/projects/status/d2voj1te0m7agfne/branch/master?svg=true)](https://ci.appveyor.com/project/michaelsproul/rust-radix-trie/branch/master) 5 | 6 | This is a [Radix Trie][radix-wiki] implementation in Rust, building on the lessons learnt from `TrieMap` and [Sequence Trie][seq-trie]. You can read about my experience implementing this data structure [here][radix-paper]. 7 | 8 | # Help Wanted, Enquire Within 9 | 10 | *Since writing this code I haven't used it in anger (or production) so it is no doubt in need of some maintenance, testing and optimisation love. If you would like to help out, try solving an open issue, optimise something, or just have a poke around! Thanks :)* 11 | 12 | # Features 13 | 14 | * Compressed nodes. Common key prefixes are stored only once. 15 | * Trie-specific methods to look-up closest ancestors and descendants. 16 | * Key Generic. Any type that can be serialised as a vector of bytes can be used as a key. 17 | * Safe - no unsafe code. 18 | 19 | # Documentation 20 | 21 | https://docs.rs/radix_trie/ 22 | 23 | # Usage 24 | 25 | Available on [Crates.io][] as [`radix_trie`][radix-crate]. 26 | 27 | Just add `radix_trie` to the dependencies section of your `Cargo.toml`, like so: 28 | 29 | ```toml 30 | radix_trie = "0.2" 31 | ``` 32 | 33 | # Contributors 34 | 35 | Made by: 36 | 37 | * Allan Simon ([@allan-simon](https://github.com/allan-simon)) 38 | * Andrew Smith ([@andrewcsmith](https://github.com/andrewcsmith)) 39 | * Arthur Carcano ([@NougatRillettes](https://github.com/NougatRillettes)) 40 | * Devin Ragotzy ([@DevinR528](https://github.com/DevinR528)) 41 | * [@hanabi1224](https://github.com/hanabi1224) 42 | * Jakob Dalsgaard ([@jakobdalsgaard](https://github.com/jakobdalsgaard)) 43 | * Michael Sproul ([@michaelsproul](https://github.com/michaelsproul)) 44 | * Robin Lambertz ([@roblabla](https://github.com/roblabla)) 45 | * Sergey ([@Albibek](https://github.com/Albibek)) 46 | * Stuart Hinson ([@stuarth](https://github.com/stuarth)) 47 | * Vinzent Steinberg ([@vks](https://github.com/vks)) 48 | 49 | # License 50 | 51 | MIT License. Copyright © Michael Sproul and contributors 2015-present. 52 | 53 | [radix-wiki]: http://en.wikipedia.org/wiki/Radix_tree 54 | [seq-trie]: https://github.com/michaelsproul/rust_sequence_trie 55 | [radix-paper]: https://michaelsproul.github.io/rust_radix_paper/ 56 | [crates.io]: https://crates.io/ 57 | [radix-crate]: https://crates.io/crates/radix_trie 58 | -------------------------------------------------------------------------------- /appveyor.yml: -------------------------------------------------------------------------------- 1 | image: 2 | - Visual Studio 2017 3 | environment: 4 | matrix: 5 | - RUST_TOOL_CHAIN: stable 6 | - RUST_TOOL_CHAIN: nightly 7 | before_build: 8 | - rustc -V 9 | - cargo -V 10 | - rustup component add rustfmt-preview 11 | build_script: 12 | - cargo build --release 13 | test_script: 14 | - cargo test --release --verbose --features serde 15 | - cargo doc --no-deps --features serde 16 | - cargo fmt --all -- --check 17 | branches: 18 | only: 19 | - master 20 | pull_requests: 21 | do_not_increment_build_number: true 22 | matrix: 23 | fast_finish: true 24 | for: 25 | - 26 | matrix: 27 | only: 28 | - image: Visual Studio 2017 29 | install: 30 | - ps: iex (new-object net.webclient).downloadstring('https://get.scoop.sh') 31 | - scoop install rustup 32 | - refreshenv 33 | - rustup default %RUST_TOOL_CHAIN% 34 | -------------------------------------------------------------------------------- /benches/trie_benches.rs: -------------------------------------------------------------------------------- 1 | use criterion::{criterion_group, criterion_main, Criterion}; 2 | use radix_trie::Trie; 3 | 4 | fn get_text() -> Vec { 5 | use std::fs::File; 6 | use std::io::Read; 7 | const DATA: &[&str] = &["data/1984.txt", "data/sun-rising.txt"]; 8 | let mut contents = String::new(); 9 | File::open(&DATA[1]) 10 | .unwrap() 11 | .read_to_string(&mut contents) 12 | .unwrap(); 13 | contents 14 | .split(|c: char| c.is_whitespace()) 15 | .map(|s| s.to_string()) 16 | .collect() 17 | } 18 | 19 | fn make_trie(words: &[String]) -> Trie<&str, usize> { 20 | let mut trie = Trie::new(); 21 | for w in words { 22 | trie.insert(&w[..], w.len()); 23 | } 24 | trie 25 | } 26 | 27 | fn trie_insert(b: &mut Criterion) { 28 | let words = get_text(); 29 | b.bench_function("trie insert", |b| b.iter(|| make_trie(&words))); 30 | } 31 | 32 | fn trie_get(b: &mut Criterion) { 33 | let words = get_text(); 34 | let trie = make_trie(&words); 35 | b.bench_function("trie get", |b| { 36 | b.iter(|| { 37 | words 38 | .iter() 39 | .map(|w| trie.get(&&w[..])) 40 | .collect::>>() 41 | }) 42 | }); 43 | } 44 | 45 | fn trie_insert_remove(b: &mut Criterion) { 46 | let words = get_text(); 47 | 48 | b.bench_function("trie remove", |b| { 49 | b.iter(|| { 50 | let mut trie = make_trie(&words); 51 | for w in &words { 52 | trie.remove(&&w[..]); 53 | } 54 | }); 55 | }); 56 | } 57 | 58 | criterion_group!(benches, trie_insert, trie_get, trie_insert_remove); 59 | 60 | criterion_main!(benches); 61 | -------------------------------------------------------------------------------- /data/sun-rising.txt: -------------------------------------------------------------------------------- 1 | BUSY old fool, unruly Sun, 2 | Why dost thou thus, 3 | Through windows, and through curtains, call on us ? 4 | Must to thy motions lovers' seasons run ? 5 | Saucy pedantic wretch, go chide 6 | Late school-boys and sour prentices, 7 | Go tell court-huntsmen that the king will ride, 8 | Call country ants to harvest offices ; 9 | Love, all alike, no season knows nor clime, 10 | Nor hours, days, months, which are the rags of time. 11 | 12 | Thy beams so reverend, and strong 13 | Why shouldst thou think ? 14 | I could eclipse and cloud them with a wink, 15 | But that I would not lose her sight so long. 16 | If her eyes have not blinded thine, 17 | Look, and to-morrow late tell me, 18 | Whether both th' Indias of spice and mine 19 | Be where thou left'st them, or lie here with me. 20 | Ask for those kings whom thou saw'st yesterday, 21 | And thou shalt hear, "All here in one bed lay." 22 | 23 | She's all states, and all princes I ; 24 | Nothing else is ; 25 | Princes do but play us ; compared to this, 26 | All honour's mimic, all wealth alchemy. 27 | Thou, Sun, art half as happy as we, 28 | In that the world's contracted thus ; 29 | Thine age asks ease, and since thy duties be 30 | To warm the world, that's done in warming us. 31 | Shine here to us, and thou art everywhere ; 32 | This bed thy center is, these walls thy sphere. 33 | -------------------------------------------------------------------------------- /examples/child_iter.rs: -------------------------------------------------------------------------------- 1 | extern crate radix_trie; 2 | 3 | use radix_trie::{Trie, TrieCommon}; 4 | 5 | fn main() { 6 | let mut t = Trie::new(); 7 | t.insert("z", 2); 8 | t.insert("aba", 5); 9 | t.insert("abb", 6); 10 | t.insert("abc", 50); 11 | 12 | // This is a bit of a hack that relies on knowing the binary representation of 13 | // strings... "abd" works, but "abz" doesn't... 14 | let ab_sum = t.get_raw_ancestor(&"abd").children().fold(0, |acc, c| { 15 | println!("Iterating over child with value: {:?}", c.value()); 16 | acc + *c.value().unwrap_or(&0) 17 | }); 18 | println!("{}", ab_sum); 19 | assert_eq!(ab_sum, 5 + 6 + 50); 20 | } 21 | -------------------------------------------------------------------------------- /examples/debug.rs: -------------------------------------------------------------------------------- 1 | extern crate radix_trie; 2 | 3 | use radix_trie::Trie; 4 | 5 | fn main() { 6 | let mut trie = Trie::new(); 7 | trie.insert("hello", 19u32); 8 | trie.insert("hellcat", 35u32); 9 | trie.insert("not related", 1u32); 10 | trie.insert("handle nested", 5u32); 11 | 12 | println!("{:#?}", trie); 13 | } 14 | -------------------------------------------------------------------------------- /examples/opt.rs: -------------------------------------------------------------------------------- 1 | extern crate radix_trie; 2 | 3 | use radix_trie::*; 4 | 5 | fn main() { 6 | let mut trie = Trie::new(); 7 | let mut key = vec![]; 8 | for i in 0..10_000 { 9 | key.push(0); 10 | trie.insert(key.clone(), i); 11 | key.pop(); 12 | key.push(1); 13 | trie.insert(key.clone(), i); 14 | } 15 | //let res = trie.remove(&blow_stack); 16 | //println!("{}", res.unwrap()); 17 | } 18 | -------------------------------------------------------------------------------- /examples/string_frequency.rs: -------------------------------------------------------------------------------- 1 | extern crate radix_trie; 2 | 3 | use radix_trie::{Trie, TrieCommon}; 4 | 5 | fn main() { 6 | let example = "bananaphone".to_string(); 7 | // Our frequency trie will store various strings with the frequency that they are used 8 | let mut trie: Trie = Trie::new(); 9 | let v: Vec = example.chars().collect(); 10 | 11 | // We want the frequencies of all strings from 1 to 3 characters long 12 | for window_length in 1..4 { 13 | for a in v.windows(window_length) { 14 | // Create a new String to hold our key 15 | let mut s = String::new(); 16 | // Append all chars in a to the String 17 | s.extend(a); 18 | // If the value at s exists, add 1. Otherwise, assign 1 as the key's value. 19 | trie.map_with_default(s, |v| *v += 1, 1); 20 | } 21 | } 22 | 23 | // Iterate through all the values in the trie and print them with their frequencies. 24 | // Iterates and prints in lexicographic order. 25 | println!("All trie nodes"); 26 | for (k, v) in trie.iter() { 27 | println!("{}: {}", k, v); 28 | } 29 | 30 | println!("All children of 'a'"); 31 | for n in trie.subtrie(&"a".to_string()).unwrap().children() { 32 | println!("{}: {}", n.key().unwrap(), n.value().unwrap()); 33 | } 34 | } 35 | -------------------------------------------------------------------------------- /src/iter.rs: -------------------------------------------------------------------------------- 1 | //! Iterators over key-value pairs, keys, values and child subtries. 2 | 3 | use std::iter::{FilterMap, FromIterator, Map}; 4 | use std::slice; 5 | 6 | use crate::TrieNode; 7 | use crate::{SubTrie, Trie, TrieKey}; 8 | 9 | use nibble_vec::Nibblet; 10 | 11 | // MY EYES. 12 | type Child = Box>; 13 | type RawChildIter<'a, K, V> = slice::Iter<'a, Option>>; 14 | type ChildMapFn<'a, K, V> = fn(&'a Option>) -> Option<&'a Child>; 15 | type ChildIter<'a, K, V> = FilterMap, ChildMapFn<'a, K, V>>; 16 | 17 | /// Iterator over the keys and values of a Trie. 18 | pub struct Iter<'a, K: 'a, V: 'a> { 19 | root: &'a TrieNode, 20 | root_visited: bool, 21 | stack: Vec>, 22 | } 23 | 24 | impl<'a, K, V> Iter<'a, K, V> { 25 | // TODO: make this private somehow (and same for the other iterators). 26 | pub fn new(root: &'a TrieNode) -> Iter<'a, K, V> { 27 | Iter { 28 | root: root, 29 | root_visited: false, 30 | stack: vec![], 31 | } 32 | } 33 | } 34 | 35 | /// Iterator over the keys of a Trie. 36 | pub struct Keys<'a, K: 'a, V: 'a> { 37 | inner: Map, KeyMapFn<'a, K, V>>, 38 | } 39 | 40 | type KeyMapFn<'a, K, V> = fn((&'a K, &'a V)) -> &'a K; 41 | 42 | impl<'a, K, V> Keys<'a, K, V> { 43 | pub fn new(iter: Iter<'a, K, V>) -> Keys<'a, K, V> { 44 | fn first<'b, K, V>((k, _): (&'b K, &'b V)) -> &'b K { 45 | k 46 | } 47 | Keys { 48 | inner: iter.map(first), 49 | } 50 | } 51 | } 52 | 53 | impl<'a, K, V> Iterator for Keys<'a, K, V> { 54 | type Item = &'a K; 55 | 56 | fn next(&mut self) -> Option<&'a K> { 57 | self.inner.next() 58 | } 59 | } 60 | 61 | /// Iterator over the values of a Trie. 62 | pub struct Values<'a, K: 'a, V: 'a> { 63 | inner: Map, ValueMapFn<'a, K, V>>, 64 | } 65 | 66 | type ValueMapFn<'a, K, V> = fn((&'a K, &'a V)) -> &'a V; 67 | 68 | impl<'a, K, V> Values<'a, K, V> { 69 | pub fn new(iter: Iter<'a, K, V>) -> Values<'a, K, V> { 70 | fn second<'b, K, V>((_, v): (&'b K, &'b V)) -> &'b V { 71 | v 72 | } 73 | Values { 74 | inner: iter.map(second), 75 | } 76 | } 77 | } 78 | 79 | impl<'a, K, V> Iterator for Values<'a, K, V> { 80 | type Item = &'a V; 81 | 82 | fn next(&mut self) -> Option<&'a V> { 83 | self.inner.next() 84 | } 85 | } 86 | 87 | /// Iterator over the child subtries of a trie. 88 | pub struct Children<'a, K: 'a, V: 'a> { 89 | prefix: Nibblet, 90 | inner: ChildIter<'a, K, V>, 91 | } 92 | 93 | impl<'a, K, V> Children<'a, K, V> { 94 | pub fn new(key: Nibblet, node: &'a TrieNode) -> Self { 95 | Children { 96 | prefix: key, 97 | inner: node.child_iter(), 98 | } 99 | } 100 | } 101 | 102 | impl<'a, K, V> Iterator for Children<'a, K, V> { 103 | type Item = SubTrie<'a, K, V>; 104 | 105 | fn next(&mut self) -> Option> { 106 | self.inner.next().map(|node| SubTrie { 107 | prefix: self.prefix.clone().join(&node.key), 108 | node: node, 109 | }) 110 | } 111 | } 112 | 113 | impl TrieNode { 114 | /// Helper function to get all the non-empty children of a node. 115 | fn child_iter(&self) -> ChildIter { 116 | fn id(x: &Option>) -> Option<&Child> { 117 | x.as_ref() 118 | } 119 | 120 | self.children.iter().filter_map(id) 121 | } 122 | 123 | /// Get the key and value of a node as a pair. 124 | fn kv_as_pair(&self) -> Option<(&K, &V)> { 125 | self.key_value.as_ref().map(|kv| (&kv.key, &kv.value)) 126 | } 127 | } 128 | 129 | enum IterAction<'a, K: 'a, V: 'a> { 130 | Push(&'a TrieNode), 131 | Pop, 132 | } 133 | 134 | impl<'a, K, V> Iterator for Iter<'a, K, V> { 135 | type Item = (&'a K, &'a V); 136 | 137 | fn next(&mut self) -> Option { 138 | use self::IterAction::*; 139 | 140 | // Visit each node as it is reached from its parent (with special root handling). 141 | if !self.root_visited { 142 | self.root_visited = true; 143 | self.stack.push(self.root.child_iter()); 144 | if let Some(kv) = self.root.kv_as_pair() { 145 | return Some(kv); 146 | } 147 | } 148 | 149 | loop { 150 | let action = match self.stack.last_mut() { 151 | Some(stack_top) => match stack_top.next() { 152 | Some(child) => Push(child), 153 | None => Pop, 154 | }, 155 | None => return None, 156 | }; 157 | 158 | match action { 159 | Push(trie) => { 160 | self.stack.push(trie.child_iter()); 161 | if let Some(kv) = trie.kv_as_pair() { 162 | return Some(kv); 163 | } 164 | } 165 | Pop => { 166 | self.stack.pop(); 167 | } 168 | } 169 | } 170 | } 171 | } 172 | 173 | impl FromIterator<(K, V)> for Trie 174 | where 175 | K: TrieKey, 176 | { 177 | fn from_iter(iter: T) -> Trie 178 | where 179 | T: IntoIterator, 180 | { 181 | let mut trie = Trie::new(); 182 | for (k, v) in iter { 183 | trie.insert(k, v); 184 | } 185 | trie 186 | } 187 | } 188 | -------------------------------------------------------------------------------- /src/keys.rs: -------------------------------------------------------------------------------- 1 | use endian_type::{BigEndian, LittleEndian}; 2 | use std::ffi::OsString; 3 | use std::path::{Path, PathBuf}; 4 | 5 | use nibble_vec::Nibblet; 6 | 7 | /// Trait for types which can be used to key a Radix Trie. 8 | /// 9 | /// Types that implement this trait should be convertible to a vector of half-bytes (nibbles) 10 | /// such that no two instances of the type convert to the same vector. 11 | /// To protect against faulty behaviour, the trie will **panic** if it finds two distinct keys 12 | /// of type `K` which encode to the same `Nibblet`, so be careful! 13 | /// 14 | /// If you would like to implement this trait for your own type, you need to implement 15 | /// *either* `encode_bytes` or `encode`. You only need to implement one of the two. 16 | /// If you don't implement one, your code will **panic** as soon you use the trie. 17 | /// There is no performance penalty for implementing `encode_bytes` instead of `encode`, 18 | /// so it is preferred except in the case where you require half-byte precision. 19 | /// 20 | /// Many standard types implement this trait already. Integer types are encoded *big-endian* 21 | /// by default but can be encoded little-endian using the `LittleEndian` wrapper type. 22 | pub trait TrieKey: PartialEq + Eq { 23 | /// Encode a value as a vector of bytes. 24 | fn encode_bytes(&self) -> Vec { 25 | panic!("implement this method or TrieKey::encode"); 26 | } 27 | 28 | /// Encode a value as a NibbleVec. 29 | #[inline] 30 | fn encode(&self) -> Nibblet { 31 | Nibblet::from_byte_vec(self.encode_bytes()) 32 | } 33 | } 34 | 35 | /// Key comparison result. 36 | #[derive(Debug)] 37 | pub enum KeyMatch { 38 | /// The keys match up to the given index. 39 | Partial(usize), 40 | /// The first key is a prefix of the second. 41 | FirstPrefix, 42 | /// The second key is a prefix of the first. 43 | SecondPrefix, 44 | /// The keys match exactly. 45 | Full, 46 | } 47 | 48 | /// Compare two Trie keys. 49 | /// 50 | /// Compares `first[start_idx .. ]` to `second`, i.e. only looks at a slice of the first key. 51 | #[inline] 52 | pub fn match_keys(start_idx: usize, first: &Nibblet, second: &Nibblet) -> KeyMatch { 53 | let first_len = first.len() - start_idx; 54 | let min_length = ::std::cmp::min(first_len, second.len()); 55 | 56 | for i in 0..min_length { 57 | if first.get(start_idx + i) != second.get(i) { 58 | return KeyMatch::Partial(i); 59 | } 60 | } 61 | 62 | match (first_len, second.len()) { 63 | (x, y) if x < y => KeyMatch::FirstPrefix, 64 | (x, y) if x == y => KeyMatch::Full, 65 | _ => KeyMatch::SecondPrefix, 66 | } 67 | } 68 | 69 | /// Check two keys for equality and panic if they differ. 70 | #[inline] 71 | pub fn check_keys(key1: &K, key2: &K) 72 | where 73 | K: TrieKey, 74 | { 75 | if *key1 != *key2 { 76 | panic!("multiple-keys with the same bit representation."); 77 | } 78 | } 79 | 80 | // --- TrieKey Implementations for standard types --- /// 81 | 82 | // This blanket implementation goes into play when specialization is stabilized 83 | // impl TrieKey for T where T: Into> + Clone + Eq + PartialEq { 84 | // fn encode_bytes(&self) -> Vec { 85 | // self.clone().into() 86 | // } 87 | // } 88 | 89 | impl TrieKey for Vec { 90 | #[inline] 91 | fn encode_bytes(&self) -> Vec { 92 | self.clone() 93 | } 94 | } 95 | 96 | impl TrieKey for [u8] { 97 | #[inline] 98 | fn encode_bytes(&self) -> Vec { 99 | self.to_vec() 100 | } 101 | } 102 | 103 | impl TrieKey for String { 104 | #[inline] 105 | fn encode_bytes(&self) -> Vec { 106 | self.as_bytes().encode_bytes() 107 | } 108 | } 109 | 110 | impl TrieKey for str { 111 | #[inline] 112 | fn encode_bytes(&self) -> Vec { 113 | self.as_bytes().encode_bytes() 114 | } 115 | } 116 | 117 | impl<'a, T: ?Sized + TrieKey> TrieKey for &'a T { 118 | #[inline] 119 | fn encode_bytes(&self) -> Vec { 120 | (**self).encode_bytes() 121 | } 122 | } 123 | 124 | impl<'a, T: ?Sized + TrieKey> TrieKey for &'a mut T { 125 | #[inline] 126 | fn encode_bytes(&self) -> Vec { 127 | (**self).encode_bytes() 128 | } 129 | } 130 | 131 | impl TrieKey for i8 { 132 | #[inline] 133 | fn encode_bytes(&self) -> Vec { 134 | let mut v: Vec = Vec::with_capacity(1); 135 | v.push(*self as u8); 136 | v 137 | } 138 | } 139 | 140 | impl TrieKey for u8 { 141 | #[inline] 142 | fn encode_bytes(&self) -> Vec { 143 | let mut v: Vec = Vec::with_capacity(1); 144 | v.push(*self); 145 | v 146 | } 147 | } 148 | 149 | #[cfg(unix)] 150 | impl TrieKey for PathBuf { 151 | fn encode_bytes(&self) -> Vec { 152 | use std::os::unix::ffi::OsStringExt; 153 | let str: OsString = self.clone().into(); 154 | str.into_vec() 155 | } 156 | } 157 | 158 | #[cfg(unix)] 159 | impl TrieKey for Path { 160 | fn encode_bytes(&self) -> Vec { 161 | use std::os::unix::ffi::OsStrExt; 162 | self.as_os_str().as_bytes().encode_bytes() 163 | } 164 | } 165 | 166 | impl TrieKey for LittleEndian 167 | where 168 | T: Eq + Copy, 169 | { 170 | fn encode_bytes(&self) -> Vec { 171 | self.as_bytes().encode_bytes() 172 | } 173 | } 174 | 175 | impl TrieKey for BigEndian 176 | where 177 | T: Eq + Copy, 178 | { 179 | fn encode_bytes(&self) -> Vec { 180 | self.as_bytes().to_vec() 181 | } 182 | } 183 | 184 | macro_rules! int_keys { 185 | ( $( $t:ty ),* ) => { 186 | $( 187 | impl TrieKey for $t { 188 | fn encode_bytes(&self) -> Vec { 189 | let be: BigEndian<$t> = From::from(*self); 190 | be.encode_bytes() 191 | } 192 | } 193 | )* 194 | }; 195 | } 196 | 197 | int_keys!(u16, u32, u64, i16, i32, i64, usize, isize); 198 | 199 | macro_rules! vec_int_keys { 200 | ( $( $t:ty ),* ) => { 201 | $( 202 | impl TrieKey for Vec<$t> { 203 | fn encode_bytes(&self) -> Vec { 204 | let mut v = Vec::::with_capacity(self.len() * std::mem::size_of::<$t>()); 205 | for u in self { 206 | v.extend_from_slice(&u.to_be_bytes()); 207 | } 208 | v 209 | } 210 | } 211 | )* 212 | }; 213 | } 214 | 215 | vec_int_keys!(u16, u32, u64, i16, i32, i64, usize, isize); 216 | 217 | #[cfg(test)] 218 | mod test { 219 | pub trait DefaultTrieKey { 220 | fn encode_bytes(&self) -> Vec; 221 | } 222 | 223 | impl> + Clone + PartialEq + Eq> DefaultTrieKey for T { 224 | #[inline] 225 | fn encode_bytes(&self) -> Vec { 226 | self.clone().into() 227 | } 228 | } 229 | 230 | pub trait AsTrieKey { 231 | fn encode_bytes(&self) -> Vec; 232 | } 233 | 234 | impl + Clone + PartialEq + Eq> AsTrieKey for &T { 235 | #[inline] 236 | fn encode_bytes(&self) -> Vec { 237 | self.as_ref().to_vec() 238 | } 239 | } 240 | 241 | macro_rules! encode_bytes { 242 | ($e:expr) => { 243 | (&$e).encode_bytes() 244 | }; 245 | } 246 | 247 | #[test] 248 | fn test_autoref_specialization() { 249 | let _ = encode_bytes!([0_u8]); 250 | let _ = encode_bytes!("hello"); 251 | let _ = encode_bytes!("hello".to_string()); 252 | } 253 | } 254 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! A wonderful, fast, safe, generic radix trie implementation. 2 | //! 3 | //! To get started, see the docs for `Trie` below. 4 | 5 | // #![warn(missing_docs)] 6 | 7 | extern crate endian_type; 8 | extern crate nibble_vec; 9 | #[cfg(test)] 10 | extern crate quickcheck; 11 | #[cfg(test)] 12 | extern crate rand; 13 | 14 | pub use keys::TrieKey; 15 | pub use nibble_vec::NibbleVec; 16 | pub use trie_common::TrieCommon; 17 | use trie_node::TrieNode; 18 | 19 | use nibble_vec::Nibblet; 20 | 21 | #[macro_use] 22 | mod macros; 23 | pub mod iter; 24 | mod keys; 25 | #[cfg(feature = "serde")] 26 | mod serde; 27 | mod subtrie; 28 | mod traversal; 29 | mod trie; 30 | mod trie_common; 31 | mod trie_node; 32 | 33 | #[cfg(test)] 34 | mod qc_test; 35 | #[cfg(test)] 36 | mod test; 37 | 38 | const BRANCH_FACTOR: usize = 16; 39 | 40 | /// Data-structure for storing and querying string-like keys and associated values. 41 | /// 42 | /// Any keys which share a common *prefix* are stored below a single copy of that prefix. 43 | /// This saves space, and also allows the longest prefix of any given key to be found. 44 | /// 45 | /// You can read more about Radix Tries on [Wikipedia][radix-wiki]. 46 | /// 47 | /// Lots of the methods on `Trie` return optional values - they can be composed 48 | /// nicely using `Option::and_then`. 49 | /// 50 | /// [radix-wiki]: http://en.wikipedia.org/wiki/Radix_tree 51 | #[derive(Debug, Clone)] 52 | pub struct Trie { 53 | /// The number of values stored in this sub-trie (this node and all descendants). 54 | length: usize, 55 | /// The main content of this trie. 56 | node: TrieNode, 57 | } 58 | 59 | /// Immutable view of a sub-tree a larger trie. 60 | #[derive(Debug)] 61 | pub struct SubTrie<'a, K: 'a, V: 'a> { 62 | prefix: Nibblet, 63 | node: &'a TrieNode, 64 | } 65 | 66 | /// Mutable view of a sub-tree of a larger trie. 67 | #[derive(Debug)] 68 | pub struct SubTrieMut<'a, K: 'a, V: 'a> { 69 | prefix: Nibblet, 70 | length: &'a mut usize, 71 | node: &'a mut TrieNode, 72 | } 73 | 74 | /// Wrapper for subtrie lookup results. 75 | /// 76 | /// When fetching from a subtrie, if the prefix is wrong you'll get an `Err(())`. 77 | /// Otherwise you'll get an `Ok(_)`, where the contained option value is what would ordinarily 78 | /// be returned by get/insert/whatever. 79 | pub type SubTrieResult = Result, ()>; 80 | -------------------------------------------------------------------------------- /src/macros.rs: -------------------------------------------------------------------------------- 1 | // Identity macro to allow expansion of the "mutability" token tree. 2 | macro_rules! id { 3 | ($e:item) => { 4 | $e 5 | }; 6 | } 7 | -------------------------------------------------------------------------------- /src/qc_test.rs: -------------------------------------------------------------------------------- 1 | //! Proper testing, with QuickCheck. 2 | 3 | use crate::{Trie, TrieCommon, TrieKey}; 4 | use quickcheck::{quickcheck, Arbitrary, Gen}; 5 | use std::collections::{HashMap, HashSet}; 6 | use std::iter::FromIterator; 7 | 8 | #[derive(Clone, Debug, PartialEq, Eq, Hash)] 9 | struct Key(Vec); 10 | 11 | #[derive(Clone, Debug)] 12 | struct RandomKeys(HashSet); 13 | 14 | const MAX_KEYS: usize = 512; 15 | const KEY_RUN_LEN: usize = 8; 16 | const KEY_MAX_VAL: u8 = 4; 17 | 18 | impl Arbitrary for Key { 19 | fn arbitrary(g: &mut Gen) -> Key { 20 | let len = usize::arbitrary(g) % KEY_RUN_LEN; 21 | let mut key = Vec::with_capacity(len); 22 | for _ in 0..len { 23 | key.push(u8::arbitrary(g) % KEY_MAX_VAL); 24 | } 25 | Key(key) 26 | } 27 | } 28 | 29 | impl Key { 30 | fn extend_random(&self, g: &mut Gen) -> Key { 31 | self.extend(Key::arbitrary(g)) 32 | } 33 | 34 | fn extend(&self, other: Key) -> Key { 35 | let mut key = self.clone(); 36 | key.0.extend(other.0); 37 | key 38 | } 39 | 40 | fn len(&self) -> usize { 41 | self.0.len() 42 | } 43 | } 44 | 45 | impl TrieKey for Key { 46 | fn encode_bytes(&self) -> Vec { 47 | self.0.clone() 48 | } 49 | } 50 | 51 | impl Arbitrary for RandomKeys { 52 | fn arbitrary(g: &mut Gen) -> RandomKeys { 53 | let num_keys = usize::arbitrary(g) % MAX_KEYS; 54 | let mut keys = Vec::with_capacity(num_keys); 55 | keys.push(Key::arbitrary(g)); 56 | 57 | for _ in 0..num_keys { 58 | match u8::arbitrary(g) % 10 { 59 | // Generate a new random key. 60 | 1 => keys.push(Key::arbitrary(g)), 61 | // Extend an existing key. 62 | _ => { 63 | let i = usize::arbitrary(g) % keys.len(); 64 | let key = keys[i].extend_random(g); 65 | keys.push(key); 66 | } 67 | } 68 | } 69 | 70 | RandomKeys(HashSet::from_iter(keys)) 71 | } 72 | } 73 | 74 | #[test] 75 | fn insert_all_remove_all() { 76 | fn prop(RandomKeys(keys): RandomKeys) -> bool { 77 | let mut trie = Trie::new(); 78 | let mut length = 0; 79 | 80 | for k in &keys { 81 | if trie.insert(k.clone(), k.len()).is_some() { 82 | return false; 83 | } 84 | length += 1; 85 | if trie.len() != length { 86 | return false; 87 | } 88 | } 89 | 90 | if !trie.check_integrity() { 91 | return false; 92 | } 93 | 94 | for k in &keys { 95 | if trie.get(&k) != Some(&k.len()) { 96 | return false; 97 | } 98 | if trie.remove(&k) != Some(k.len()) { 99 | return false; 100 | } 101 | length -= 1; 102 | if trie.len() != length { 103 | return false; 104 | } 105 | if trie.get(&k).is_some() { 106 | return false; 107 | } 108 | } 109 | trie.check_integrity() 110 | } 111 | 112 | quickcheck(prop as fn(RandomKeys) -> bool); 113 | } 114 | 115 | #[test] 116 | fn subtrie() { 117 | fn prop(RandomKeys(keys): RandomKeys) -> bool { 118 | let half = keys.len() / 2; 119 | let first_half = keys.iter().take(half).map(|k| (k.clone(), k.len())); 120 | let trie = Trie::from_iter(first_half); 121 | 122 | // Check node existence for inserted keys. 123 | for k in keys.iter().take(half) { 124 | match trie.subtrie(&k) { 125 | Some(node) => { 126 | if node.value() != Some(&k.len()) { 127 | return false; 128 | } 129 | } 130 | None => return false, 131 | } 132 | } 133 | 134 | // Check that nodes for non-inserted keys don't have values. 135 | for k in keys.iter().skip(half) { 136 | if let Some(node) = trie.subtrie(&k) { 137 | if node.value().is_some() { 138 | return false; 139 | } 140 | } 141 | } 142 | 143 | trie.check_integrity() 144 | } 145 | 146 | quickcheck(prop as fn(RandomKeys) -> bool); 147 | } 148 | 149 | // trie.subtrie(k1).get(k2) should be the same as trie.get(k2) if k1 is a prefix of k2. 150 | #[test] 151 | fn subtrie_get() { 152 | fn prop(trie_keys: RandomKeys, k1: Key, k2: Key) -> bool { 153 | let mut trie = length_trie(trie_keys.0); 154 | trie.insert(k1.clone(), k1.len()); 155 | 156 | let subtrie = trie.subtrie(&k1).unwrap(); 157 | 158 | if k2.0.starts_with(&k1.0) { 159 | subtrie.get(&k2).unwrap() == trie.get(&k2) 160 | } else { 161 | subtrie.get(&k2).is_err() 162 | } 163 | } 164 | 165 | quickcheck(prop as fn(RandomKeys, Key, Key) -> bool); 166 | } 167 | 168 | #[test] 169 | fn subtrie_mut_get() { 170 | fn prop(trie_keys: RandomKeys, k1: Key, k2: Key) -> bool { 171 | let mut trie = length_trie(trie_keys.0); 172 | trie.insert(k1.clone(), k1.len()); 173 | 174 | let subtrie = trie.subtrie_mut(&k1).unwrap(); 175 | 176 | let ok = if k2.0.starts_with(&k1.0) { 177 | subtrie.get(&k2).is_ok() 178 | } else { 179 | subtrie.get(&k2).is_err() 180 | }; 181 | ok && trie.check_integrity() 182 | } 183 | 184 | quickcheck(prop as fn(RandomKeys, Key, Key) -> bool); 185 | } 186 | 187 | #[test] 188 | fn subtrie_insert() { 189 | fn prop(trie_keys: RandomKeys, key_suffixes: RandomKeys, k1: Key) -> bool { 190 | let mut trie = length_trie(trie_keys.0); 191 | trie.insert(k1.clone(), k1.len()); 192 | 193 | { 194 | let mut subtrie = trie.subtrie_mut(&k1).unwrap(); 195 | 196 | let insert_keys = key_suffixes 197 | .0 198 | .into_iter() 199 | .map(|x| k1.extend(x)) 200 | .collect::>(); 201 | 202 | for k in insert_keys.iter() { 203 | assert!(subtrie.insert(k.clone(), k.len()).is_ok()); 204 | } 205 | 206 | for k in insert_keys.iter() { 207 | match subtrie.get(k) { 208 | Ok(Some(_)) => (), 209 | _ => return false, 210 | } 211 | } 212 | } 213 | 214 | trie.check_integrity() 215 | } 216 | 217 | quickcheck(prop as fn(RandomKeys, RandomKeys, Key) -> bool); 218 | } 219 | 220 | // Construct a trie from a set of keys, with each key mapped to its length. 221 | fn length_trie(keys: HashSet) -> Trie { 222 | let mut t = Trie::new(); 223 | for k in keys { 224 | let len = k.len(); 225 | t.insert(k, len); 226 | } 227 | t 228 | } 229 | 230 | #[test] 231 | fn remove_non_existent() { 232 | fn prop(RandomKeys(insert_keys): RandomKeys, RandomKeys(remove_keys): RandomKeys) -> bool { 233 | let mut trie = length_trie(insert_keys.clone()); 234 | 235 | for k in remove_keys { 236 | if !insert_keys.contains(&k) && trie.remove(&k).is_some() { 237 | return false; 238 | } 239 | } 240 | trie.check_integrity() 241 | } 242 | quickcheck(prop as fn(RandomKeys, RandomKeys) -> bool); 243 | } 244 | 245 | #[test] 246 | fn keys_iter() { 247 | fn prop(RandomKeys(keys): RandomKeys) -> bool { 248 | let trie = length_trie(keys.clone()); 249 | let trie_keys: HashSet = trie.keys().cloned().collect(); 250 | trie_keys == keys 251 | } 252 | quickcheck(prop as fn(RandomKeys) -> bool); 253 | } 254 | 255 | #[test] 256 | fn values_iter() { 257 | // Create a map of values to frequencies. 258 | fn frequency_map>(values: I) -> HashMap { 259 | let mut map = HashMap::new(); 260 | for v in values { 261 | let current_val = map.entry(v).or_insert(0); 262 | *current_val += 1; 263 | } 264 | map 265 | } 266 | 267 | fn prop(RandomKeys(keys): RandomKeys) -> bool { 268 | let trie = length_trie(keys.clone()); 269 | let trie_values: HashMap = frequency_map(trie.values().cloned()); 270 | let key_values = frequency_map(keys.into_iter().map(|k| k.len())); 271 | trie_values == key_values 272 | } 273 | quickcheck(prop as fn(RandomKeys) -> bool); 274 | } 275 | -------------------------------------------------------------------------------- /src/serde.rs: -------------------------------------------------------------------------------- 1 | extern crate serde; 2 | 3 | use self::serde::ser::SerializeMap; 4 | use self::serde::{de, Deserialize, Deserializer, Serialize, Serializer}; 5 | use super::{Trie, TrieCommon, TrieKey}; 6 | use std::fmt::{self, Formatter}; 7 | use std::marker::PhantomData; 8 | 9 | impl Serialize for Trie 10 | where 11 | K: Serialize + TrieKey, 12 | V: Serialize, 13 | { 14 | fn serialize(&self, serializer: S) -> Result 15 | where 16 | S: Serializer, 17 | { 18 | let mut map = serializer.serialize_map(Some(self.len()))?; 19 | for (k, v) in self.iter() { 20 | map.serialize_entry(k, v)?; 21 | } 22 | map.end() 23 | } 24 | } 25 | 26 | struct TrieVisitor { 27 | marker: PhantomData>, 28 | } 29 | 30 | impl TrieVisitor { 31 | fn new() -> Self { 32 | TrieVisitor { 33 | marker: PhantomData, 34 | } 35 | } 36 | } 37 | 38 | impl<'a, K, V> de::Visitor<'a> for TrieVisitor 39 | where 40 | K: Deserialize<'a> + Clone + Eq + PartialEq + TrieKey, 41 | V: Deserialize<'a>, 42 | { 43 | type Value = Trie; 44 | 45 | fn expecting(&self, formatter: &mut Formatter) -> fmt::Result { 46 | write!(formatter, "a serialized trie") 47 | } 48 | 49 | fn visit_map(self, mut visitor: M) -> Result 50 | where 51 | M: de::MapAccess<'a>, 52 | { 53 | let mut values = Trie::new(); 54 | 55 | while let Some((key, value)) = visitor.next_entry()? { 56 | values.insert(key, value); 57 | } 58 | 59 | Ok(values) 60 | } 61 | 62 | fn visit_unit(self) -> Result 63 | where 64 | E: de::Error, 65 | { 66 | Ok(Trie::new()) 67 | } 68 | } 69 | 70 | impl<'a, K, V> Deserialize<'a> for Trie 71 | where 72 | K: Deserialize<'a> + Clone + Eq + PartialEq + TrieKey, 73 | V: Deserialize<'a>, 74 | { 75 | fn deserialize(deserializer: D) -> Result 76 | where 77 | D: Deserializer<'a>, 78 | { 79 | // Instantiate our Visitor and ask the Deserializer to drive 80 | // it over the input data, resulting in an instance of MyMap. 81 | deserializer.deserialize_map(TrieVisitor::new()) 82 | } 83 | } 84 | 85 | #[cfg(test)] 86 | mod test { 87 | extern crate serde_test; 88 | use self::serde_test::Token; 89 | use super::super::Trie; 90 | 91 | macro_rules! tests_de { 92 | ($($name:ident => $value:expr => $tokens:expr,)+) => { 93 | $(#[test] 94 | fn $name() { 95 | // Test ser/de roundtripping 96 | serde_test::assert_de_tokens(&$value, $tokens); 97 | })+ 98 | } 99 | } 100 | 101 | macro_rules! tests_ser { 102 | ($($name:ident => $value:expr => $tokens:expr,)+) => { 103 | $(#[test] 104 | fn $name() { 105 | serde_test::assert_ser_tokens(&$value, $tokens); 106 | })+ 107 | } 108 | } 109 | 110 | macro_rules! trie { 111 | () => { 112 | Trie::new() 113 | }; 114 | ($($key:expr => $value:expr),+) => { 115 | { 116 | let mut map = Trie::new(); 117 | $(map.insert($key, $value);)+ 118 | map 119 | } 120 | } 121 | } 122 | 123 | tests_ser! { 124 | test_ser_empty_trie => Trie::<&str, isize>::new() => &[ 125 | Token::Map { len: Some(0) }, 126 | Token::MapEnd, 127 | ], 128 | test_ser_single_element_trie => trie!["1" => 2] => &[ 129 | Token::Map { len: Some(1) }, 130 | 131 | Token::Str("1"), 132 | Token::I32(2), 133 | Token::MapEnd, 134 | ], 135 | test_ser_multiple_element_trie => trie!["1" => 2, "3" => 4] => &[ 136 | Token::Map { len: Some(2) }, 137 | Token::Str("1"), 138 | Token::I32(2), 139 | 140 | Token::Str("3"), 141 | Token::I32(4), 142 | Token::MapEnd, 143 | ], 144 | test_ser_deep_trie => trie!["1" => trie![], "2" => trie!["3" => 4, "5" => 6]] => &[ 145 | Token::Map { len: Some(2) }, 146 | Token::Str("1"), 147 | Token::Map { len: Some(0) }, 148 | Token::MapEnd, 149 | 150 | Token::Str("2"), 151 | Token::Map { len: Some(2) }, 152 | Token::Str("3"), 153 | Token::I32(4), 154 | 155 | Token::Str("5"), 156 | Token::I32(6), 157 | Token::MapEnd, 158 | Token::MapEnd, 159 | ], 160 | } 161 | 162 | tests_de! { 163 | test_de_empty_trie1 => Trie::::new() => &[ 164 | Token::Unit, 165 | ], 166 | test_de_empty_trie2 => Trie::::new() => &[ 167 | Token::Map { len: Some(0) }, 168 | Token::MapEnd, 169 | ], 170 | test_de_single_element_trie => trie!["1".to_string() => 2] => &[ 171 | Token::Map { len: Some(1) }, 172 | Token::Str("1"), 173 | Token::I32(2), 174 | Token::MapEnd, 175 | ], 176 | test_de_multiple_element_trie => trie!["1".to_string() => 2, "3".to_string() => 4] => &[ 177 | Token::Map { len: Some(2) }, 178 | Token::Str("1"), 179 | Token::I32(2), 180 | 181 | Token::Str("3"), 182 | Token::I32(4), 183 | Token::MapEnd, 184 | ], 185 | test_de_deep_trie => trie!["1".to_string() => trie![], "2".to_string() => trie!["3".to_string() => 4, "5".to_string() => 6]] => &[ 186 | Token::Map { len: Some(2) }, 187 | Token::Str("1"), 188 | Token::Map { len: Some(0) }, 189 | Token::MapEnd, 190 | 191 | Token::Str("2"), 192 | Token::Map { len: Some(2) }, 193 | Token::Str("3"), 194 | Token::I32(4), 195 | 196 | Token::Str("5"), 197 | Token::I32(6), 198 | Token::MapEnd, 199 | Token::MapEnd, 200 | ], 201 | test_de_empty_trie3 => Trie::::new() => &[ 202 | Token::UnitStruct { name: "Anything" }, 203 | ], 204 | test_de_empty_trie4 => Trie::::new() => &[ 205 | Token::Struct { 206 | name: "Anything", 207 | len: 0, 208 | }, 209 | Token::StructEnd, 210 | ], 211 | } 212 | } 213 | -------------------------------------------------------------------------------- /src/subtrie.rs: -------------------------------------------------------------------------------- 1 | use crate::keys::*; 2 | use crate::TrieNode; 3 | use crate::{SubTrie, SubTrieMut, SubTrieResult}; 4 | use std::borrow::Borrow; 5 | 6 | use nibble_vec::Nibblet; 7 | 8 | impl<'a, K, V> SubTrie<'a, K, V> 9 | where 10 | K: TrieKey, 11 | { 12 | /// Look up the value for the given key, which should be an extension of this subtrie's key. 13 | /// 14 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 15 | /// form *must* match those for the key type 16 | pub fn get(&self, key: &Q) -> SubTrieResult<&V> 17 | where 18 | K: Borrow, 19 | Q: TrieKey, 20 | { 21 | subtrie_get(&self.prefix, self.node, key) 22 | } 23 | } 24 | 25 | fn subtrie_get<'a, K, Q: ?Sized, V>( 26 | prefix: &Nibblet, 27 | node: &'a TrieNode, 28 | key: &Q, 29 | ) -> SubTrieResult<&'a V> 30 | where 31 | K: TrieKey, 32 | K: Borrow, 33 | Q: TrieKey, 34 | { 35 | let key_enc = key.encode(); 36 | match match_keys(0, prefix, &key_enc) { 37 | KeyMatch::Full => Ok(node.value()), 38 | KeyMatch::FirstPrefix => Ok(node 39 | .get(&stripped(key_enc, prefix)) 40 | .and_then(TrieNode::value)), 41 | _ => Err(()), 42 | } 43 | } 44 | 45 | impl<'a, K, V> SubTrieMut<'a, K, V> 46 | where 47 | K: TrieKey, 48 | { 49 | /// Mutable reference to the node's value. 50 | pub fn value_mut(&mut self) -> Option<&mut V> { 51 | self.node.value_mut() 52 | } 53 | 54 | /// Look up the value for the given key, which should be an extension of this subtrie's key. 55 | /// 56 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 57 | /// form *must* match those for the key type 58 | pub fn get(&self, key: &Q) -> SubTrieResult<&V> 59 | where 60 | K: Borrow, 61 | Q: TrieKey, 62 | { 63 | subtrie_get(&self.prefix, &*self.node, key) 64 | } 65 | 66 | /// Insert a value in this subtrie. The key should be an extension of this subtrie's key. 67 | pub fn insert(&mut self, key: K, value: V) -> SubTrieResult { 68 | let key_enc = key.encode(); 69 | let previous = match match_keys(0, &self.prefix, &key_enc) { 70 | KeyMatch::Full => self.node.replace_value(key, value), 71 | KeyMatch::FirstPrefix => self 72 | .node 73 | .insert(key, value, stripped(key_enc, &self.prefix)), 74 | _ => { 75 | return Err(()); 76 | } 77 | }; 78 | 79 | if previous.is_none() { 80 | *self.length += 1; 81 | } 82 | 83 | Ok(previous) 84 | } 85 | 86 | /// Remove a value from this subtrie. The key should be an extension of this subtrie's key. 87 | /// 88 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 89 | /// form *must* match those for the key type 90 | pub fn remove(&mut self, key: &Q) -> SubTrieResult 91 | where 92 | K: Borrow, 93 | Q: TrieKey, 94 | { 95 | let key_enc = key.encode(); 96 | let removed = match match_keys(0, &self.prefix, &key_enc) { 97 | KeyMatch::Full => self.node.take_value(key), 98 | KeyMatch::FirstPrefix => self.node.remove(key), 99 | _ => { 100 | return Err(()); 101 | } 102 | }; 103 | 104 | if removed.is_some() { 105 | *self.length -= 1; 106 | } 107 | 108 | Ok(removed) 109 | } 110 | } 111 | 112 | fn stripped(mut key: Nibblet, prefix: &Nibblet) -> Nibblet { 113 | key.split(prefix.len()) 114 | } 115 | -------------------------------------------------------------------------------- /src/test.rs: -------------------------------------------------------------------------------- 1 | use crate::keys::TrieKey; 2 | use crate::{Trie, TrieCommon}; 3 | use std::collections::HashSet; 4 | use std::iter::FromIterator; 5 | 6 | const TEST_DATA: [(&'static str, u32); 7] = [ 7 | ("abcdefgh", 19), 8 | ("abcdef", 18), 9 | ("abcd", 17), 10 | ("ab", 16), 11 | ("a", 15), 12 | ("acbdef", 30), 13 | ("bcdefgh", 29), 14 | ]; 15 | 16 | fn test_trie() -> Trie<&'static str, u32> { 17 | let mut trie = Trie::new(); 18 | 19 | for &(key, val) in &TEST_DATA { 20 | trie.insert(key, val); 21 | assert!(trie.check_integrity()); 22 | } 23 | 24 | trie 25 | } 26 | 27 | #[test] 28 | fn get_nonexistant() { 29 | let trie = test_trie(); 30 | assert!(trie.get(&"nonexistant").is_none()); 31 | assert!(trie.get(&"").is_none()); 32 | } 33 | 34 | #[test] 35 | fn subtrie_nonexistant() { 36 | let trie = test_trie(); 37 | assert!(trie.subtrie(&"nonexistant").is_none()); 38 | assert!(trie.subtrie(&"").is_some()); 39 | } 40 | 41 | #[test] 42 | fn unicode() { 43 | let mut trie = Trie::new(); 44 | trie.insert("bär", 1); 45 | trie.insert("bären", 2); 46 | 47 | assert_eq!(*trie.get("bär").unwrap(), 1); 48 | let values = trie 49 | .get_raw_descendant("bä") 50 | .unwrap() 51 | .values() 52 | .collect::>(); 53 | assert_eq!([1, 2].iter().collect::>(), values); 54 | } 55 | 56 | #[test] 57 | fn subtrie() { 58 | let mut trie = Trie::new(); 59 | trie.insert("hello", 55); 60 | assert!(trie.subtrie(&"h").is_none()); 61 | assert!(trie.subtrie(&"hello").is_some()); 62 | } 63 | #[test] 64 | fn subtrie_string() { 65 | let mut trie = Trie::new(); 66 | trie.insert("hello".to_string(), 55); 67 | 68 | let h = "h".to_string(); 69 | let hello = "hello".to_string(); 70 | 71 | assert!(trie.subtrie(&h).is_none()); 72 | assert!(trie.subtrie(&hello).is_some()); 73 | } 74 | 75 | #[test] 76 | fn empty_key() { 77 | let mut trie = test_trie(); 78 | trie.insert(&"", 99); 79 | assert_eq!(*trie.get(&"").unwrap(), 99); 80 | assert_eq!(trie.remove(&""), Some(99)); 81 | } 82 | 83 | #[test] 84 | fn insert() { 85 | let trie = test_trie(); 86 | 87 | for &(key, val) in &TEST_DATA { 88 | assert_eq!(*trie.get(&key).unwrap(), val); 89 | } 90 | 91 | assert!(trie.check_integrity()); 92 | assert_eq!(trie.len(), TEST_DATA.len()); 93 | } 94 | 95 | #[test] 96 | fn insert_replace() { 97 | let mut trie = Trie::new(); 98 | assert_eq!(trie.insert("haskell", 18), None); 99 | let length = trie.len(); 100 | assert_eq!(trie.insert("haskell", 36), Some(18)); 101 | assert_eq!(trie.len(), length); 102 | } 103 | 104 | #[test] 105 | fn map_with_default() { 106 | let mut trie = test_trie(); 107 | trie.map_with_default(&"abcd", |x| *x = *x + 1, 42); 108 | assert_eq!(*trie.get(&"abcd").unwrap(), 17 + 1); 109 | trie.map_with_default(&"zzz", |x| *x = *x + 1, 42); 110 | assert_eq!(*trie.get(&"zzz").unwrap(), 42); 111 | } 112 | 113 | #[test] 114 | fn remove() { 115 | let mut trie = test_trie(); 116 | 117 | // Remove. 118 | for &(key, val) in &TEST_DATA { 119 | let res = trie.remove(&key); 120 | assert_eq!(res, Some(val)); 121 | assert!(trie.check_integrity()); 122 | } 123 | 124 | // Check non-existance. 125 | for &(key, _) in &TEST_DATA { 126 | assert!(trie.get(&key).is_none()); 127 | } 128 | } 129 | 130 | #[test] 131 | fn remove_simple() { 132 | let mut trie = Trie::new(); 133 | 134 | trie.insert("HELL", 66); 135 | trie.insert("HELLO", 77); 136 | let val = trie.remove(&"HELLO"); 137 | assert_eq!(val, Some(77)); 138 | } 139 | 140 | #[test] 141 | fn remove_non_existent() { 142 | let mut trie = Trie::new(); 143 | 144 | trie.insert("acab", true); 145 | 146 | assert_eq!(trie.remove(&"abc"), None); 147 | assert_eq!(trie.remove(&"acaba"), None); 148 | assert_eq!(trie.remove(&"a"), None); 149 | assert_eq!(trie.remove(&""), None); 150 | assert_eq!(trie.len(), 1); 151 | 152 | trie.insert("acaz", true); 153 | 154 | assert_eq!(trie.remove(&"acb"), None); 155 | assert_eq!(trie.remove(&"acaca"), None); 156 | assert_eq!(trie.remove(&"aca"), None); 157 | assert_eq!(trie.len(), 2); 158 | } 159 | 160 | #[test] 161 | fn nearest_ancestor_root() { 162 | let mut trie = Trie::new(); 163 | trie.insert("", 55); 164 | assert_eq!(trie.get_ancestor_value(&""), Some(&55)); 165 | } 166 | 167 | #[test] 168 | fn nearest_ancestor() { 169 | let trie = test_trie(); 170 | assert_eq!(trie.get_ancestor_value(&""), None); 171 | 172 | // Test identity prefixes. 173 | for &(key, val) in &TEST_DATA { 174 | assert_eq!(trie.get_ancestor_value(&key), Some(&val)); 175 | } 176 | 177 | assert_eq!(trie.get_ancestor_value(&"abcdefg"), trie.get(&"abcdef")); 178 | assert_eq!(trie.get_ancestor_value(&"abcde"), trie.get(&"abcd")); 179 | assert_eq!(trie.get_ancestor_value(&"aauksdjk"), trie.get(&"a")); 180 | } 181 | 182 | #[test] 183 | fn nearest_ancestor_no_child_fn() { 184 | let mut t = Trie::new(); 185 | t.insert("ab", 5); 186 | let anc = t.get_ancestor(&"abc").unwrap(); 187 | assert_eq!(*anc.value().unwrap(), 5); 188 | } 189 | 190 | #[test] 191 | fn raw_ancestor() { 192 | let mut t = Trie::new(); 193 | 194 | for &(key, _) in &TEST_DATA { 195 | assert_eq!(t.get_raw_ancestor(&key).key(), t.key()); 196 | } 197 | 198 | t.insert("wow", 0); 199 | t.insert("hella", 1); 200 | t.insert("hellb", 2); 201 | 202 | // Ancestor should be "hell" node. 203 | let anc = t.get_raw_ancestor(&"hello"); 204 | assert_eq!(anc.len(), 2); 205 | } 206 | 207 | // Check that the subtrie prefix is correct for raw_ancestor. 208 | #[test] 209 | fn raw_ancestor_prefix() { 210 | let mut t = Trie::new(); 211 | 212 | t.insert("abac", ()); 213 | t.insert("abaz", ()); 214 | 215 | let anc = t.get_raw_ancestor(&"aba"); 216 | assert_eq!(anc.prefix, "aba".encode()); 217 | } 218 | 219 | #[test] 220 | fn iter() { 221 | type Set = HashSet<(&'static str, u32)>; 222 | let trie = test_trie(); 223 | let expected = TEST_DATA.iter().map(|&x| x).collect::(); 224 | let observed = trie.iter().map(|(&k, &v)| (k, v)).collect::(); 225 | assert_eq!(expected, observed); 226 | } 227 | 228 | #[test] 229 | fn get_raw_descendant() { 230 | let trie = test_trie(); 231 | assert_eq!( 232 | trie.get_raw_descendant(&"abcdefgh").and_then(|t| t.value()), 233 | Some(&19) 234 | ); 235 | assert_eq!( 236 | trie.get_raw_descendant(&"abcdefg").and_then(|t| t.value()), 237 | Some(&19) 238 | ); 239 | assert!(trie.get_raw_descendant(&"acbg").is_none()); 240 | } 241 | 242 | #[test] 243 | fn raw_descendant_prefix() { 244 | let mut t = Trie::new(); 245 | 246 | t.insert("abczzzz", ()); 247 | t.insert("abcaaaa", ()); 248 | 249 | assert_eq!(t.get_raw_descendant(&"a").unwrap().prefix, "abc".encode()); 250 | assert_eq!(t.get_raw_descendant(&"abc").unwrap().prefix, "abc".encode()); 251 | assert_eq!( 252 | t.get_raw_descendant(&"abca").unwrap().prefix, 253 | "abcaaaa".encode() 254 | ); 255 | } 256 | 257 | #[test] 258 | fn get_prefix_bug() { 259 | let mut trie = Trie::new(); 260 | trie.insert("abdc", 5); 261 | trie.insert("abde", 6); 262 | assert!(trie.get(&"abc").is_none()); 263 | } 264 | 265 | #[test] 266 | fn get_ancestor_bug() { 267 | let mut trie = Trie::new(); 268 | trie.insert("abc", 1); 269 | trie.insert("abcde", 2); 270 | assert_eq!(trie.get_ancestor_value(&"abcdz"), Some(&1)); 271 | } 272 | 273 | #[test] 274 | fn root_replace_bug() { 275 | let mut trie = Trie::new(); 276 | trie.insert("a", ()); 277 | trie.insert("p", ()); 278 | trie.remove(&"a"); 279 | assert_eq!(trie.len(), 1); 280 | trie.remove(&"p"); 281 | assert_eq!(trie.len(), 0); 282 | } 283 | 284 | #[test] 285 | fn subtrie_insert() { 286 | let mut trie = Trie::new(); 287 | trie.insert("abc", 3); 288 | { 289 | let mut subtrie = trie.subtrie_mut(&"abc").unwrap(); 290 | assert_eq!(subtrie.insert("somerandomshit", 666), Err(())); 291 | assert_eq!(subtrie.insert("abcdef", 6), Ok(None)); 292 | assert_eq!(subtrie.insert("abc", 9), Ok(Some(3))); 293 | } 294 | assert_eq!(trie.get(&"abc"), Some(&9)); 295 | assert_eq!(trie.get(&"abcdef"), Some(&6)); 296 | assert_eq!(trie.len(), 2); 297 | } 298 | 299 | #[test] 300 | fn subtrie_len() { 301 | let trie = test_trie(); 302 | assert_eq!(trie.subtrie(&"ab").unwrap().len(), 4); 303 | assert_eq!(trie.subtrie(&"a").unwrap().len(), 6); 304 | assert_eq!(trie.subtrie(&"").unwrap().len(), trie.len()); 305 | assert_eq!(trie.subtrie(&"bcdefgh").unwrap().len(), 1); 306 | } 307 | 308 | // Subtrie functions that return references should return references valid for 309 | // the lifetime of the *original trie* from which they were borrowed, NOT 310 | // just the lifetime of the subtrie (which may be shorter). 311 | #[test] 312 | fn subtrie_lifetime() { 313 | let trie = test_trie(); 314 | let subtrie_value = { 315 | let subtrie = trie.subtrie(&"ab").unwrap(); 316 | subtrie.value() 317 | }; 318 | assert_eq!(*subtrie_value.unwrap(), 16); 319 | } 320 | 321 | #[test] 322 | fn subtrie_mut_lifetime() { 323 | let mut trie = test_trie(); 324 | let subtrie_value = { 325 | let mut subtrie = trie.subtrie_mut(&"ab").unwrap(); 326 | *subtrie.value_mut().unwrap() = 999; 327 | subtrie.value() 328 | }; 329 | assert_eq!(*subtrie_value.unwrap(), 999); 330 | } 331 | 332 | #[test] 333 | fn ancestor_key() { 334 | let trie = test_trie(); 335 | let subtrie = trie.get_ancestor(&"abcde").unwrap(); 336 | assert_eq!(*subtrie.key().unwrap(), "abcd"); 337 | assert_eq!(subtrie.prefix, "abcd".encode()); 338 | assert_eq!(*subtrie.get(&"abcdef").unwrap().unwrap(), 18); 339 | assert_eq!(*subtrie.get(&"abcdefgh").unwrap().unwrap(), 19); 340 | } 341 | 342 | #[test] 343 | fn child_subtrie_keys() { 344 | let trie = test_trie(); 345 | let subtrie = trie.subtrie(&"abcd").unwrap(); 346 | for subsubtrie in subtrie.children() { 347 | // This subtrie should be for "abcde". 348 | assert_eq!(*subsubtrie.get(&"abcdef").unwrap().unwrap(), 18); 349 | assert_eq!(*subsubtrie.get(&"abcdefgh").unwrap().unwrap(), 19); 350 | } 351 | } 352 | 353 | #[test] 354 | fn int_keys() { 355 | let mut trie = Trie::new(); 356 | trie.insert(0x00ffu64, "asdf"); 357 | trie.insert(0xdeadbeefu64, "asdf"); 358 | assert!(trie.check_integrity()); 359 | } 360 | 361 | #[test] 362 | fn from_iter() { 363 | let trie: Trie<&str, u32> = Trie::from_iter(vec![("test", 10), ("hello", 12)]); 364 | assert_eq!(*trie.get(&"test").unwrap(), 10); 365 | assert_eq!(*trie.get(&"hello").unwrap(), 12); 366 | assert_eq!(trie.len(), 2); 367 | } 368 | 369 | #[test] 370 | fn test_get_borrow() { 371 | let mut trie = Trie::new(); 372 | trie.insert("/boot".to_string(), "dir"); 373 | assert_eq!(*trie.get("/boot").unwrap(), "dir"); 374 | } 375 | 376 | #[test] 377 | fn test_get_mut_borrow() { 378 | let mut trie = Trie::new(); 379 | trie.insert("/boot".to_string(), "dir"); 380 | assert_eq!(*trie.get_mut("/boot").unwrap(), "dir"); 381 | } 382 | 383 | #[test] 384 | fn test_remove_borrow() { 385 | let mut trie = Trie::new(); 386 | trie.insert("/boot".to_string(), "dir"); 387 | assert_eq!(trie.remove("/boot").unwrap(), "dir"); 388 | } 389 | 390 | #[test] 391 | fn test_subtrie_borrow() { 392 | let mut trie = Trie::new(); 393 | trie.insert("/boot".to_string(), "dir"); 394 | trie.insert("/boot/lol".to_string(), "dir"); 395 | trie.insert("/bleh".to_string(), "ohi"); 396 | let subtrie = trie.subtrie("/boot").unwrap(); 397 | assert_eq!(*subtrie.value().unwrap(), "dir"); 398 | } 399 | 400 | #[test] 401 | fn test_subtrie_mut_borrow() { 402 | let mut trie = Trie::new(); 403 | trie.insert("/boot".to_string(), "dir"); 404 | trie.insert("/boot/lol".to_string(), "dir"); 405 | trie.insert("/bleh".to_string(), "ohi"); 406 | let subtrie = trie.subtrie_mut("/boot").unwrap(); 407 | assert_eq!(*subtrie.value().unwrap(), "dir"); 408 | } 409 | 410 | #[test] 411 | fn test_get_ancestor_borrow() { 412 | let mut trie = Trie::new(); 413 | trie.insert("/boot".to_string(), "ancestor"); 414 | trie.insert("/boot/lol".to_string(), "dir"); 415 | trie.insert("/bleh".to_string(), "ohi"); 416 | let subtrie = trie.get_ancestor("/boot/lo").unwrap(); 417 | assert_eq!(*subtrie.value().unwrap(), "ancestor"); 418 | } 419 | 420 | #[test] 421 | fn test_get_ancestor_value_borrow() { 422 | let mut trie = Trie::new(); 423 | trie.insert("/boot".to_string(), "ancestor"); 424 | trie.insert("/boot/lol".to_string(), "dir"); 425 | trie.insert("/bleh".to_string(), "ohi"); 426 | let ancestor_val = trie.get_ancestor_value("/boot/lo").unwrap(); 427 | assert_eq!(*ancestor_val, "ancestor"); 428 | } 429 | 430 | #[test] 431 | fn test_get_raw_ancestor_borrow() { 432 | let mut trie = Trie::new(); 433 | trie.insert("/boot".to_string(), "boot"); 434 | trie.insert("/boot/lol".to_string(), "dir"); 435 | trie.insert("/bleh".to_string(), "ohi"); 436 | let subtrie = trie.get_raw_ancestor("/boot/lo"); 437 | assert_eq!(*subtrie.value().unwrap(), "boot"); 438 | } 439 | 440 | #[test] 441 | fn test_get_raw_descendant_borrow() { 442 | let mut trie = Trie::new(); 443 | trie.insert("/boot".to_string(), "dir"); 444 | trie.insert("/boot/lol".to_string(), "lol"); 445 | trie.insert("/bleh".to_string(), "ohi"); 446 | let subtrie = trie.get_raw_descendant("/boot/lo").unwrap(); 447 | assert_eq!(*subtrie.value().unwrap(), "lol"); 448 | } 449 | 450 | #[test] 451 | fn test_prefix() { 452 | let mut t = Trie::::new(); 453 | t.insert(0xf1, ()); 454 | t.remove(&0xf2); 455 | t.insert(0xf2, ()); 456 | println!("{:#?}", t); 457 | assert_eq!(t.prefix(), [].as_ref()); 458 | let first = t.children().next().unwrap(); 459 | assert_eq!(first.prefix(), [0xf].as_ref()); 460 | let mut c = first.children(); 461 | let second = c.next().unwrap(); 462 | let third = c.next().unwrap(); 463 | assert!(c.next().is_none()); 464 | assert_eq!(second.prefix(), [0x1].as_ref()); 465 | assert_eq!(third.prefix(), [0x2].as_ref()); 466 | } 467 | 468 | #[test] 469 | fn clone() { 470 | let mut t1 = test_trie(); 471 | let mut t2 = t1.clone(); 472 | 473 | assert_eq!(t1, t2); 474 | t1.insert("abc", 22); 475 | 476 | assert_ne!(t1, t2); 477 | t2.insert("abc", 22); 478 | 479 | assert_eq!(t1, t2); 480 | } 481 | -------------------------------------------------------------------------------- /src/traversal.rs: -------------------------------------------------------------------------------- 1 | //! This module contains the core algorithms. 2 | 3 | use crate::keys::{match_keys, KeyMatch}; 4 | use crate::TrieKey; 5 | use crate::TrieNode; 6 | use std::borrow::Borrow; 7 | 8 | use nibble_vec::Nibblet; 9 | 10 | use self::DescendantResult::*; 11 | 12 | impl TrieNode 13 | where 14 | K: TrieKey, 15 | { 16 | #[inline] 17 | pub fn get(&self, nv: &Nibblet) -> Option<&TrieNode> { 18 | iterative_get(self, nv) 19 | } 20 | #[inline] 21 | pub fn get_mut(&mut self, nv: &Nibblet) -> Option<&mut TrieNode> { 22 | iterative_get_mut(self, nv) 23 | } 24 | #[inline] 25 | pub fn insert(&mut self, key: K, value: V, nv: Nibblet) -> Option { 26 | iterative_insert(self, key, value, nv) 27 | } 28 | #[inline] 29 | pub fn remove(&mut self, key: &Q) -> Option 30 | where 31 | K: Borrow, 32 | Q: TrieKey, 33 | { 34 | recursive_remove(self, key) 35 | } 36 | #[inline] 37 | pub fn get_ancestor(&self, nv: &Nibblet) -> Option<(&TrieNode, usize)> { 38 | get_ancestor(self, nv) 39 | } 40 | #[inline] 41 | pub fn get_raw_ancestor(&self, nv: &Nibblet) -> (&TrieNode, usize) { 42 | get_raw_ancestor(self, nv) 43 | } 44 | #[inline] 45 | pub fn get_raw_descendant<'a>(&'a self, nv: &Nibblet) -> Option> { 46 | get_raw_descendant(self, nv) 47 | } 48 | } 49 | 50 | macro_rules! get_func { 51 | ( 52 | name: $name:ident, 53 | trie_type: $trie_type:ty, 54 | mutability: $($mut_:tt)* 55 | ) => {id!{ 56 | #[inline] 57 | fn $name<'a, K, V>(trie: $trie_type, nv: &Nibblet) -> Option<$trie_type> { 58 | if nv.len() == 0 { 59 | return Some(trie); 60 | } 61 | 62 | let mut prev = trie; 63 | let mut depth = 0; 64 | 65 | loop { 66 | let bucket = nv.get(depth) as usize; 67 | let current = prev; 68 | if let Some(ref $($mut_)* child) = current.children[bucket] { 69 | match match_keys(depth, nv, &child.key) { 70 | KeyMatch::Full => { 71 | return Some(child); 72 | } 73 | KeyMatch::SecondPrefix => { 74 | depth += child.key.len(); 75 | prev = child; 76 | } 77 | _ => { 78 | return None; 79 | } 80 | } 81 | } else { 82 | return None; 83 | } 84 | } 85 | } 86 | }} 87 | } 88 | 89 | get_func!(name: iterative_get, trie_type: &'a TrieNode, mutability: ); 90 | get_func!(name: iterative_get_mut, trie_type: &'a mut TrieNode, mutability: mut); 91 | 92 | #[inline] 93 | fn iterative_insert(trie: &mut TrieNode, key: K, value: V, mut nv: Nibblet) -> Option 94 | where 95 | K: TrieKey, 96 | { 97 | if nv.len() == 0 { 98 | return trie.replace_value(key, value); 99 | } 100 | 101 | let mut prev = trie; 102 | let mut depth = 0; 103 | 104 | loop { 105 | let bucket = nv.get(depth) as usize; 106 | let current = prev; 107 | if let Some(ref mut child) = current.children[bucket] { 108 | match match_keys(depth, &nv, &child.key) { 109 | KeyMatch::Full => { 110 | return child.replace_value(key, value); 111 | } 112 | KeyMatch::Partial(idx) => { 113 | // Split the existing child. 114 | child.split(idx); 115 | 116 | // Insert the new key below the prefix node. 117 | let new_key = nv.split(depth + idx); 118 | let new_key_bucket = new_key.get(0) as usize; 119 | 120 | child.add_child( 121 | new_key_bucket, 122 | Box::new(TrieNode::with_key_value(new_key, key, value)), 123 | ); 124 | 125 | return None; 126 | } 127 | KeyMatch::FirstPrefix => { 128 | child.split(nv.len() - depth); 129 | child.add_key_value(key, value); 130 | return None; 131 | } 132 | KeyMatch::SecondPrefix => { 133 | depth += child.key.len(); 134 | prev = child; 135 | } 136 | } 137 | } else { 138 | let node_key = nv.split(depth); 139 | current.add_child( 140 | bucket, 141 | Box::new(TrieNode::with_key_value(node_key, key, value)), 142 | ); 143 | return None; 144 | } 145 | } 146 | } 147 | 148 | // TODO: clean this up and make it iterative. 149 | #[inline] 150 | fn recursive_remove(trie: &mut TrieNode, key: &Q) -> Option 151 | where 152 | K: TrieKey, 153 | K: Borrow, 154 | Q: TrieKey, 155 | { 156 | let nv = key.encode(); 157 | 158 | if nv.len() == 0 { 159 | return trie.take_value(key); 160 | } 161 | 162 | let bucket = nv.get(0) as usize; 163 | 164 | let child = trie.take_child(bucket); 165 | 166 | match child { 167 | Some(mut child) => { 168 | match match_keys(0, &nv, &child.key) { 169 | KeyMatch::Full => { 170 | let result = child.take_value(key); 171 | if child.child_count != 0 { 172 | // If removing this node's value has made it a value-less node with a 173 | // single child, then merge its child. 174 | let repl = if child.child_count == 1 { 175 | get_merge_child(&mut child) 176 | } else { 177 | child 178 | }; 179 | trie.add_child(bucket, repl); 180 | } 181 | result 182 | } 183 | KeyMatch::SecondPrefix => { 184 | let depth = child.key.len(); 185 | rec_remove(trie, child, bucket, key, depth, &nv) 186 | } 187 | KeyMatch::FirstPrefix | KeyMatch::Partial(_) => { 188 | trie.add_child(bucket, child); 189 | None 190 | } 191 | } 192 | } 193 | None => None, 194 | } 195 | } 196 | #[inline] 197 | fn get_merge_child(trie: &mut TrieNode) -> Box> 198 | where 199 | K: TrieKey, 200 | { 201 | let mut child = trie.take_only_child(); 202 | 203 | // Join the child's key onto the existing one. 204 | child.key = trie.key.clone().join(&child.key); 205 | 206 | child 207 | } 208 | 209 | // Tail-recursive remove function used by `recursive_remove`. 210 | #[inline] 211 | fn rec_remove( 212 | parent: &mut TrieNode, 213 | mut middle: Box>, 214 | prev_bucket: usize, 215 | key: &Q, 216 | depth: usize, 217 | nv: &Nibblet, 218 | ) -> Option 219 | where 220 | K: TrieKey, 221 | K: Borrow, 222 | Q: TrieKey, 223 | { 224 | let bucket = nv.get(depth) as usize; 225 | 226 | let child = middle.take_child(bucket); 227 | parent.add_child(prev_bucket, middle); 228 | 229 | match child { 230 | Some(mut child) => { 231 | let middle = parent.children[prev_bucket].as_mut().unwrap(); 232 | match match_keys(depth, nv, &child.key) { 233 | KeyMatch::Full => { 234 | let result = child.take_value(key); 235 | 236 | // If this node has children, keep it. 237 | if child.child_count != 0 { 238 | // If removing this node's value has made it a value-less node with a 239 | // single child, then merge its child. 240 | let repl = if child.child_count == 1 { 241 | get_merge_child(&mut *child) 242 | } else { 243 | child 244 | }; 245 | middle.add_child(bucket, repl); 246 | } 247 | // Otherwise, if the parent node now only has a single child, merge it. 248 | else if middle.child_count == 1 && middle.key_value.is_none() { 249 | let repl = get_merge_child(middle); 250 | *middle = repl; 251 | } 252 | 253 | result 254 | } 255 | KeyMatch::SecondPrefix => { 256 | let new_depth = depth + child.key.len(); 257 | rec_remove(middle, child, bucket, key, new_depth, nv) 258 | } 259 | KeyMatch::FirstPrefix | KeyMatch::Partial(_) => { 260 | middle.add_child(bucket, child); 261 | None 262 | } 263 | } 264 | } 265 | None => None, 266 | } 267 | } 268 | #[inline] 269 | fn get_ancestor<'a, K, V>( 270 | trie: &'a TrieNode, 271 | nv: &Nibblet, 272 | ) -> Option<(&'a TrieNode, usize)> 273 | where 274 | K: TrieKey, 275 | { 276 | if nv.len() == 0 { 277 | return trie.as_value_node().map(|node| (node, 0)); 278 | } 279 | 280 | let mut prev = trie; 281 | // The ancestor is such that all nodes upto and including `prev` have 282 | // already been considered. 283 | let mut ancestor = prev.as_value_node(); 284 | let mut depth = 0; 285 | 286 | loop { 287 | let bucket = nv.get(depth) as usize; 288 | let current = prev; 289 | if let Some(ref child) = current.children[bucket] { 290 | match match_keys(depth, nv, &child.key) { 291 | KeyMatch::Full => { 292 | return child 293 | .as_value_node() 294 | .map(|node| (node, depth + node.key.len())) 295 | .or_else(|| ancestor.map(|anc| (anc, depth))); 296 | } 297 | KeyMatch::FirstPrefix | KeyMatch::Partial(_) => { 298 | return ancestor.map(|anc| (anc, depth)); 299 | } 300 | KeyMatch::SecondPrefix => { 301 | depth += child.key.len(); 302 | ancestor = child.as_value_node().or(ancestor); 303 | prev = child; 304 | } 305 | } 306 | } else { 307 | return ancestor.map(|anc| (anc, depth)); 308 | } 309 | } 310 | } 311 | #[inline] 312 | fn get_raw_ancestor<'a, K, V>(trie: &'a TrieNode, nv: &Nibblet) -> (&'a TrieNode, usize) 313 | where 314 | K: TrieKey, 315 | { 316 | if nv.len() == 0 { 317 | return (trie, 0); 318 | } 319 | 320 | let mut prev = trie; 321 | // The ancestor is such that all nodes upto and including `prev` have 322 | // already been considered. 323 | let mut ancestor = prev; 324 | let mut depth = 0; 325 | 326 | loop { 327 | let bucket = nv.get(depth) as usize; 328 | let current = prev; 329 | if let Some(ref child) = current.children[bucket] { 330 | match match_keys(depth, nv, &child.key) { 331 | KeyMatch::Full => { 332 | return (child, depth + child.key.len()); 333 | } 334 | KeyMatch::FirstPrefix | KeyMatch::Partial(_) => { 335 | return (ancestor, depth); 336 | } 337 | KeyMatch::SecondPrefix => { 338 | depth += child.key.len(); 339 | ancestor = child; 340 | prev = child; 341 | } 342 | } 343 | } else { 344 | return (ancestor, depth); 345 | } 346 | } 347 | } 348 | 349 | // Type used to propogate subtrie construction instructions to the top-level `get_raw_descendant` 350 | // method. 351 | pub enum DescendantResult<'a, K: 'a, V: 'a> { 352 | NoModification(&'a TrieNode), 353 | ExtendKey(&'a TrieNode, usize, &'a Nibblet), 354 | } 355 | #[inline] 356 | fn get_raw_descendant<'a, K, V>( 357 | trie: &'a TrieNode, 358 | nv: &Nibblet, 359 | ) -> Option> { 360 | if nv.len() == 0 { 361 | return Some(NoModification(trie)); 362 | } 363 | 364 | let mut prev = trie; 365 | let mut depth = 0; 366 | 367 | loop { 368 | let bucket = nv.get(depth) as usize; 369 | let current = prev; 370 | if let Some(ref child) = current.children[bucket] { 371 | match match_keys(depth, nv, &child.key) { 372 | KeyMatch::Full => { 373 | return Some(NoModification(child)); 374 | } 375 | KeyMatch::FirstPrefix => { 376 | return Some(ExtendKey(child, depth, &child.key)); 377 | } 378 | KeyMatch::SecondPrefix => { 379 | depth += child.key.len(); 380 | prev = child; 381 | } 382 | _ => { 383 | return None; 384 | } 385 | } 386 | } else { 387 | return None; 388 | } 389 | } 390 | } 391 | -------------------------------------------------------------------------------- /src/trie.rs: -------------------------------------------------------------------------------- 1 | use crate::traversal::DescendantResult::*; 2 | use crate::TrieNode; 3 | use crate::{SubTrie, SubTrieMut, Trie, TrieCommon, TrieKey}; 4 | use std::borrow::Borrow; 5 | 6 | use nibble_vec::Nibblet; 7 | 8 | impl Trie 9 | where 10 | K: TrieKey, 11 | { 12 | /// Create an empty Trie. 13 | #[inline] 14 | pub fn new() -> Trie { 15 | Trie { 16 | length: 0, 17 | node: TrieNode::new(), 18 | } 19 | } 20 | 21 | /// Fetch a reference to the given key's corresponding value, if any. 22 | /// 23 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 24 | /// form *must* match those for the key type. 25 | #[inline] 26 | pub fn get(&self, key: &Q) -> Option<&V> 27 | where 28 | K: Borrow, 29 | Q: TrieKey, 30 | { 31 | let key_fragments = key.encode(); 32 | self.node 33 | .get(&key_fragments) 34 | .and_then(|t| t.value_checked(key)) 35 | } 36 | 37 | /// Fetch a mutable reference to the given key's corresponding value, if any. 38 | /// 39 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 40 | /// form *must* match those for the key type. 41 | #[inline] 42 | pub fn get_mut(&mut self, key: &Q) -> Option<&mut V> 43 | where 44 | K: Borrow, 45 | Q: TrieKey, 46 | { 47 | let key_fragments = key.encode(); 48 | self.node 49 | .get_mut(&key_fragments) 50 | .and_then(|t| t.value_checked_mut(key)) 51 | } 52 | 53 | /// Insert the given key-value pair, returning any previous value associated with the key. 54 | #[inline] 55 | pub fn insert(&mut self, key: K, value: V) -> Option { 56 | let key_fragments = key.encode(); 57 | let result = self.node.insert(key, value, key_fragments); 58 | if result.is_none() { 59 | self.length += 1; 60 | } 61 | result 62 | } 63 | 64 | /// Remove the value associated with the given key. 65 | /// 66 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 67 | /// form *must* match those for the key type. 68 | #[inline] 69 | pub fn remove(&mut self, key: &Q) -> Option 70 | where 71 | K: Borrow, 72 | Q: TrieKey, 73 | { 74 | let removed = self.node.remove(key); 75 | if removed.is_some() { 76 | self.length -= 1; 77 | } 78 | removed 79 | } 80 | 81 | /// Get a mutable reference to the value stored at this node, if any. 82 | pub fn value_mut(&mut self) -> Option<&mut V> { 83 | self.node.value_mut() 84 | } 85 | 86 | /// Fetch a reference to the subtrie for a given key. 87 | /// 88 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 89 | /// form *must* match those for the key type. 90 | #[inline] 91 | pub fn subtrie<'a, Q: ?Sized>(&'a self, key: &Q) -> Option> 92 | where 93 | K: Borrow, 94 | Q: TrieKey, 95 | { 96 | let key_fragments = key.encode(); 97 | self.node 98 | .get(&key_fragments) 99 | .map(|node| node.as_subtrie(key_fragments)) 100 | } 101 | 102 | /// Fetch a mutable reference to the subtrie for a given key. 103 | /// 104 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 105 | /// form *must* match those for the key type. 106 | #[inline] 107 | pub fn subtrie_mut<'a, Q: ?Sized>(&'a mut self, key: &Q) -> Option> 108 | where 109 | K: Borrow, 110 | Q: TrieKey, 111 | { 112 | let key_fragments = key.encode(); 113 | let length_ref = &mut self.length; 114 | self.node 115 | .get_mut(&key_fragments) 116 | .map(move |node| node.as_subtrie_mut(key_fragments, length_ref)) 117 | } 118 | 119 | /// Fetch a reference to the closest ancestor node of the given key. 120 | /// 121 | /// If `key` is encoded as byte-vector `b`, return the node `n` in the tree 122 | /// such that `n`'s key's byte-vector is the longest possible prefix of `b`, and `n` 123 | /// has a value. 124 | /// 125 | /// Invariant: `result.is_some() => result.key_value.is_some()`. 126 | /// 127 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 128 | /// form *must* match those for the key type. 129 | #[inline] 130 | pub fn get_ancestor<'a, Q: ?Sized>(&'a self, key: &Q) -> Option> 131 | where 132 | K: Borrow, 133 | Q: TrieKey, 134 | { 135 | let mut key_fragments = key.encode(); 136 | self.node 137 | .get_ancestor(&key_fragments) 138 | .map(|(node, node_key_len)| { 139 | key_fragments.split(node_key_len); 140 | node.as_subtrie(key_fragments) 141 | }) 142 | } 143 | 144 | /// Fetch the closest ancestor *value* for a given key. 145 | /// 146 | /// See `get_ancestor` for precise semantics, this is just a shortcut. 147 | /// 148 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 149 | /// form *must* match those for the key type. 150 | #[inline] 151 | pub fn get_ancestor_value(&self, key: &Q) -> Option<&V> 152 | where 153 | K: Borrow, 154 | Q: TrieKey, 155 | { 156 | self.get_ancestor(key).and_then(|t| t.node.value()) 157 | } 158 | 159 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 160 | /// form *must* match those for the key type 161 | #[inline] 162 | pub fn get_raw_ancestor<'a, Q: ?Sized>(&'a self, key: &Q) -> SubTrie<'a, K, V> 163 | where 164 | K: Borrow, 165 | Q: TrieKey, 166 | { 167 | let mut nv = key.encode(); 168 | let (ancestor_node, depth) = self.node.get_raw_ancestor(&nv); 169 | nv.split(depth); 170 | ancestor_node.as_subtrie(nv) 171 | } 172 | 173 | /// Fetch the closest descendant for a given key. 174 | /// 175 | /// If the key is in the trie, this is the same as `subtrie`. 176 | /// 177 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 178 | /// form *must* match those for the key type 179 | #[inline] 180 | pub fn get_raw_descendant<'a, Q: ?Sized>(&'a self, key: &Q) -> Option> 181 | where 182 | K: Borrow, 183 | Q: TrieKey, 184 | { 185 | let mut nv = key.encode(); 186 | self.node.get_raw_descendant(&nv).map(|desc| { 187 | let (node, prefix) = match desc { 188 | NoModification(node) => (node, nv), 189 | ExtendKey(node, depth, extension) => { 190 | nv.split(depth); 191 | (node, nv.join(extension)) 192 | } 193 | }; 194 | node.as_subtrie(prefix) 195 | }) 196 | } 197 | 198 | /// Take a function `f` and apply it to the value stored at `key`. 199 | /// 200 | /// If no value is stored at `key`, store `default`. 201 | #[inline] 202 | pub fn map_with_default(&mut self, key: K, f: F, default: V) 203 | where 204 | F: Fn(&mut V), 205 | { 206 | { 207 | if let Some(v) = self.get_mut(&key) { 208 | f(v); 209 | return; 210 | } 211 | } 212 | self.insert(key, default); 213 | } 214 | 215 | /// Check that the Trie invariants are satisfied - you shouldn't ever have to call this! 216 | /// Quite slow! 217 | #[doc(hidden)] 218 | pub fn check_integrity(&self) -> bool { 219 | let (ok, length) = self.node.check_integrity_recursive(&Nibblet::new()); 220 | ok && length == self.length 221 | } 222 | } 223 | 224 | impl PartialEq for Trie 225 | where 226 | K: TrieKey, 227 | V: PartialEq, 228 | { 229 | #[inline] 230 | fn eq(&self, other: &Trie) -> bool { 231 | if self.len() != other.len() { 232 | return false; 233 | } 234 | 235 | self.iter() 236 | .all(|(key, value)| other.get(key).map_or(false, |v| *value == *v)) 237 | } 238 | } 239 | 240 | impl Default for Trie { 241 | #[inline] 242 | fn default() -> Self { 243 | Self::new() 244 | } 245 | } 246 | -------------------------------------------------------------------------------- /src/trie_common.rs: -------------------------------------------------------------------------------- 1 | use crate::iter::*; 2 | use crate::TrieNode; 3 | use crate::{SubTrie, SubTrieMut, Trie, TrieKey}; 4 | 5 | use nibble_vec::Nibblet; 6 | 7 | /// Common functionality available for tries and subtries. 8 | pub trait TrieCommon<'a, K: 'a, V: 'a>: ContainsTrieNode<'a, K, V> 9 | where 10 | K: TrieKey, 11 | Self: Sized, 12 | { 13 | /// Get the key stored at this node, if any. 14 | #[inline] 15 | fn key(self) -> Option<&'a K> { 16 | self.trie_node().key() 17 | } 18 | 19 | /// Get the value stored at this node, if any. 20 | #[inline] 21 | fn value(self) -> Option<&'a V> { 22 | self.trie_node().value() 23 | } 24 | 25 | /// Number of key/value pairs stored in this trie. 26 | fn len(self) -> usize; 27 | 28 | /// Determine if the Trie contains 0 key-value pairs. 29 | #[inline] 30 | fn is_empty(self) -> bool { 31 | self.len() == 0 32 | } 33 | 34 | /// Determine if the trie is a leaf node (has no children). 35 | #[inline] 36 | fn is_leaf(self) -> bool { 37 | self.trie_node().child_count == 0 38 | } 39 | 40 | /// Return an iterator over the keys and values of the Trie. 41 | #[inline] 42 | fn iter(self) -> Iter<'a, K, V> { 43 | Iter::new(self.trie_node()) 44 | } 45 | 46 | /// Return an iterator over the keys of the Trie. 47 | #[inline] 48 | fn keys(self) -> Keys<'a, K, V> { 49 | Keys::new(self.iter()) 50 | } 51 | 52 | /// Return an iterator over the values of the Trie. 53 | #[inline] 54 | fn values(self) -> Values<'a, K, V> { 55 | Values::new(self.iter()) 56 | } 57 | 58 | /// Return an iterator over the child subtries of this node. 59 | fn children(self) -> Children<'a, K, V>; 60 | 61 | /// Get the prefix of this node. 62 | #[inline] 63 | fn prefix(self) -> &'a Nibblet { 64 | &self.trie_node().key 65 | } 66 | } 67 | 68 | /// Helper trait for Trie/SubTrie/SubTrieMut, which all contain a trie node. 69 | pub trait ContainsTrieNode<'a, K: 'a, V: 'a> 70 | where 71 | K: TrieKey, 72 | { 73 | fn trie_node(self) -> &'a TrieNode; 74 | } 75 | 76 | /// Regular trie. 77 | impl<'a, K: 'a, V: 'a> ContainsTrieNode<'a, K, V> for &'a Trie 78 | where 79 | K: TrieKey, 80 | { 81 | #[inline] 82 | fn trie_node(self) -> &'a TrieNode { 83 | &self.node 84 | } 85 | } 86 | 87 | impl<'a, K: 'a, V: 'a> TrieCommon<'a, K, V> for &'a Trie 88 | where 89 | K: TrieKey, 90 | { 91 | #[inline] 92 | fn len(self) -> usize { 93 | self.length 94 | } 95 | #[inline] 96 | fn children(self) -> Children<'a, K, V> { 97 | Children::new(self.node.key.clone(), &self.node) 98 | } 99 | } 100 | 101 | /// Subtrie. 102 | impl<'a: 'b, 'b, K: 'a, V: 'a> ContainsTrieNode<'a, K, V> for &'b SubTrie<'a, K, V> 103 | where 104 | K: TrieKey, 105 | { 106 | #[inline] 107 | fn trie_node(self) -> &'a TrieNode { 108 | self.node 109 | } 110 | } 111 | 112 | impl<'a: 'b, 'b, K: 'a, V: 'a> TrieCommon<'a, K, V> for &'b SubTrie<'a, K, V> 113 | where 114 | K: TrieKey, 115 | { 116 | #[inline] 117 | fn len(self) -> usize { 118 | self.node.compute_size() 119 | } 120 | #[inline] 121 | fn children(self) -> Children<'a, K, V> { 122 | Children::new(self.prefix.clone(), self.node) 123 | } 124 | } 125 | 126 | /// Mutable subtrie *by value* (consumes the subtrie). 127 | impl<'a, K: 'a, V: 'a> ContainsTrieNode<'a, K, V> for SubTrieMut<'a, K, V> 128 | where 129 | K: TrieKey, 130 | { 131 | #[inline] 132 | fn trie_node(self) -> &'a TrieNode { 133 | self.node 134 | } 135 | } 136 | 137 | impl<'a, K: 'a, V: 'a> TrieCommon<'a, K, V> for SubTrieMut<'a, K, V> 138 | where 139 | K: TrieKey, 140 | { 141 | /// **Computes** from scratch. 142 | #[inline] 143 | fn len(self) -> usize { 144 | self.node.compute_size() 145 | } 146 | #[inline] 147 | fn children(self) -> Children<'a, K, V> { 148 | Children::new(self.prefix.clone(), self.node) 149 | } 150 | } 151 | 152 | /// Mutable subtrie *by reference* (doesn't consume the subtrie, but limited). 153 | impl<'a: 'b, 'b, K: 'a, V: 'a> ContainsTrieNode<'b, K, V> for &'b SubTrieMut<'a, K, V> 154 | where 155 | K: TrieKey, 156 | { 157 | #[inline] 158 | fn trie_node(self) -> &'b TrieNode { 159 | self.node 160 | } 161 | } 162 | 163 | impl<'a: 'b, 'b, K: 'a, V: 'a> TrieCommon<'b, K, V> for &'b SubTrieMut<'a, K, V> 164 | where 165 | K: TrieKey, 166 | { 167 | #[inline] 168 | fn len(self) -> usize { 169 | self.node.compute_size() 170 | } 171 | #[inline] 172 | fn children(self) -> Children<'b, K, V> { 173 | Children::new(self.prefix.clone(), self.node) 174 | } 175 | } 176 | -------------------------------------------------------------------------------- /src/trie_node.rs: -------------------------------------------------------------------------------- 1 | use crate::keys::*; 2 | use crate::{SubTrie, SubTrieMut, BRANCH_FACTOR}; 3 | use std::borrow::Borrow; 4 | use std::default::Default; 5 | 6 | use nibble_vec::Nibblet; 7 | 8 | #[derive(Debug, Clone)] 9 | pub struct TrieNode { 10 | /// Key fragments/bits associated with this node, such that joining the keys from all 11 | /// parent nodes and this node is equal to the bit-encoding of this node's key. 12 | pub key: Nibblet, 13 | 14 | /// The key and value stored at this node. 15 | pub key_value: Option>>, 16 | 17 | /// The number of children which are Some rather than None. 18 | pub child_count: usize, 19 | 20 | /// The children of this node stored such that the first nibble of each child key 21 | /// dictates the child's bucket. 22 | pub children: [Option>>; BRANCH_FACTOR], 23 | } 24 | 25 | #[derive(Debug, Clone)] 26 | pub struct KeyValue { 27 | pub key: K, 28 | pub value: V, 29 | } 30 | 31 | impl TrieNode 32 | where 33 | K: TrieKey, 34 | { 35 | /// Create a value-less, child-less TrieNode. 36 | #[inline] 37 | pub fn new() -> TrieNode { 38 | TrieNode { 39 | key: Nibblet::new(), 40 | key_value: None, 41 | children: Default::default(), 42 | child_count: 0, 43 | } 44 | } 45 | 46 | /// Create a TrieNode with no children. 47 | #[inline] 48 | pub fn with_key_value(key_fragments: Nibblet, key: K, value: V) -> TrieNode { 49 | TrieNode { 50 | key: key_fragments, 51 | key_value: Some(Box::new(KeyValue { 52 | key: key, 53 | value: value, 54 | })), 55 | children: Default::default(), 56 | child_count: 0, 57 | } 58 | } 59 | 60 | /// Get the key stored at this node, if any. 61 | #[inline] 62 | pub fn key(&self) -> Option<&K> { 63 | self.key_value.as_ref().map(|kv| &kv.key) 64 | } 65 | 66 | /// Get the value stored at this node, if any. 67 | #[inline] 68 | pub fn value(&self) -> Option<&V> { 69 | self.key_value.as_ref().map(|kv| &kv.value) 70 | } 71 | 72 | /// Get a mutable reference to the value stored at this node, if any. 73 | #[inline] 74 | pub fn value_mut(&mut self) -> Option<&mut V> { 75 | self.key_value.as_mut().map(|kv| &mut kv.value) 76 | } 77 | 78 | /// Get the value whilst checking a key match. 79 | /// 80 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 81 | /// form *must* match those for the key type. 82 | #[inline] 83 | pub fn value_checked(&self, key: &Q) -> Option<&V> 84 | where 85 | K: Borrow, 86 | Q: TrieKey, 87 | { 88 | self.key_value.as_ref().map(|kv| { 89 | check_keys(kv.key.borrow(), key); 90 | &kv.value 91 | }) 92 | } 93 | 94 | /// Get a mutable value whilst checking a key match. 95 | /// 96 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 97 | /// form *must* match those for the key type. 98 | #[inline] 99 | pub fn value_checked_mut(&mut self, key: &Q) -> Option<&mut V> 100 | where 101 | K: Borrow, 102 | Q: TrieKey, 103 | { 104 | self.key_value.as_mut().map(|kv| { 105 | check_keys(kv.key.borrow(), key); 106 | &mut kv.value 107 | }) 108 | } 109 | 110 | /// Compute the number of keys and values in this node's subtrie. 111 | #[inline] 112 | pub fn compute_size(&self) -> usize { 113 | let mut size = self.key_value.is_some() as usize; 114 | 115 | for child in &self.children { 116 | if let Some(ref child) = *child { 117 | // TODO: could unroll this recursion 118 | size += child.compute_size(); 119 | } 120 | } 121 | 122 | size 123 | } 124 | 125 | /// Add a child at the given index, given that none exists there already. 126 | #[inline] 127 | pub fn add_child(&mut self, idx: usize, node: Box>) { 128 | debug_assert!(self.children[idx].is_none()); 129 | self.child_count += 1; 130 | self.children[idx] = Some(node); 131 | } 132 | 133 | /// Remove a child at the given index, if it exists. 134 | #[inline] 135 | pub fn take_child(&mut self, idx: usize) -> Option>> { 136 | self.children[idx].take().map(|node| { 137 | self.child_count -= 1; 138 | node 139 | }) 140 | } 141 | 142 | /// Helper function for removing the single child of a node. 143 | #[inline] 144 | pub fn take_only_child(&mut self) -> Box> { 145 | debug_assert_eq!(self.child_count, 1); 146 | for i in 0..BRANCH_FACTOR { 147 | if let Some(child) = self.take_child(i) { 148 | return child; 149 | } 150 | } 151 | unreachable!("node with child_count 1 has no actual children"); 152 | } 153 | 154 | /// Set the key and value of a node, given that it currently lacks one. 155 | #[inline] 156 | pub fn add_key_value(&mut self, key: K, value: V) { 157 | debug_assert!(self.key_value.is_none()); 158 | self.key_value = Some(Box::new(KeyValue { key, value })); 159 | } 160 | 161 | /// Move the value out of a node, whilst checking that its key is as expected. 162 | /// Can panic (see check_keys). 163 | /// 164 | /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed 165 | /// form *must* match those for the key type 166 | #[inline] 167 | pub fn take_value(&mut self, key: &Q) -> Option 168 | where 169 | K: Borrow, 170 | Q: TrieKey, 171 | { 172 | self.key_value.take().map(|kv| { 173 | check_keys(kv.key.borrow(), key); 174 | kv.value 175 | }) 176 | } 177 | 178 | /// Replace a value, returning the previous value if there was one. 179 | #[inline] 180 | pub fn replace_value(&mut self, key: K, value: V) -> Option { 181 | // TODO: optimise this? 182 | let previous = self.take_value(&key); 183 | self.add_key_value(key, value); 184 | previous 185 | } 186 | 187 | /// Get a reference to this node if it has a value. 188 | #[inline] 189 | pub fn as_value_node(&self) -> Option<&TrieNode> { 190 | self.key_value.as_ref().map(|_| self) 191 | } 192 | 193 | /// Split a node at a given index in its key, transforming it into a prefix node of its 194 | /// previous self. 195 | #[inline] 196 | pub fn split(&mut self, idx: usize) { 197 | // Extract all the parts of the suffix node, starting with the key. 198 | let key = self.key.split(idx); 199 | 200 | // Key-value. 201 | let key_value = self.key_value.take(); 202 | 203 | // Children. 204 | let mut children: [Option>>; BRANCH_FACTOR] = Default::default(); 205 | 206 | for (i, child) in self.children.iter_mut().enumerate() { 207 | if child.is_some() { 208 | children[i] = child.take(); 209 | } 210 | } 211 | 212 | // Child count. 213 | let child_count = self.child_count; 214 | self.child_count = 1; 215 | 216 | // Insert the collected items below what is now an empty prefix node. 217 | let bucket = key.get(0) as usize; 218 | self.children[bucket] = Some(Box::new(TrieNode { 219 | key: key, 220 | key_value, 221 | children, 222 | child_count, 223 | })); 224 | } 225 | #[inline] 226 | pub fn as_subtrie(&self, prefix: Nibblet) -> SubTrie { 227 | SubTrie { 228 | prefix: prefix, 229 | node: self, 230 | } 231 | } 232 | #[inline] 233 | pub fn as_subtrie_mut<'a>( 234 | &'a mut self, 235 | prefix: Nibblet, 236 | length: &'a mut usize, 237 | ) -> SubTrieMut<'a, K, V> { 238 | SubTrieMut { 239 | prefix: prefix, 240 | length: length, 241 | node: self, 242 | } 243 | } 244 | 245 | /// Check the integrity of a trie subtree (quite costly). 246 | /// Return true and the size of the subtree if all checks are successful, 247 | /// or false and a junk value if any test fails. 248 | pub fn check_integrity_recursive(&self, prefix: &Nibblet) -> (bool, usize) { 249 | let mut sub_tree_size = 0; 250 | let is_root = prefix.len() == 0; 251 | 252 | // Check that no value-less, non-root nodes have only 1 child. 253 | if !is_root && self.child_count == 1 && self.key_value.is_none() { 254 | println!("Value-less node with a single child."); 255 | return (false, sub_tree_size); 256 | } 257 | 258 | // Check that all non-root key vector's have length > 1. 259 | if !is_root && self.key.len() == 0 { 260 | println!("Key length is 0 at non-root node."); 261 | return (false, sub_tree_size); 262 | } 263 | 264 | // Check that the child count matches the actual number of children. 265 | let child_count = self 266 | .children 267 | .iter() 268 | .fold(0, |acc, e| acc + (e.is_some() as usize)); 269 | 270 | if child_count != self.child_count { 271 | println!( 272 | "Child count error, recorded: {}, actual: {}", 273 | self.child_count, child_count 274 | ); 275 | return (false, sub_tree_size); 276 | } 277 | 278 | // Compute the key fragments for this node, according to the trie. 279 | let trie_key = prefix.clone().join(&self.key); 280 | 281 | // Account for this node in the size check, and check its key. 282 | if let Some(ref kv) = self.key_value { 283 | sub_tree_size += 1; 284 | 285 | let actual_key = kv.key.encode(); 286 | 287 | if trie_key != actual_key { 288 | return (false, sub_tree_size); 289 | } 290 | } 291 | 292 | // Recursively check children. 293 | for i in 0..BRANCH_FACTOR { 294 | if let Some(ref child) = self.children[i] { 295 | match child.check_integrity_recursive(&trie_key) { 296 | (false, _) => return (false, sub_tree_size), 297 | (true, child_size) => sub_tree_size += child_size, 298 | } 299 | } 300 | } 301 | 302 | (true, sub_tree_size) 303 | } 304 | } 305 | 306 | impl Default for TrieNode { 307 | fn default() -> Self { 308 | Self::new() 309 | } 310 | } 311 | --------------------------------------------------------------------------------