diff --git a/radix_tree/src/lib.rs b/radix_tree/src/lib.rs index cafe9d3..9d4f6df 100644 --- a/radix_tree/src/lib.rs +++ b/radix_tree/src/lib.rs @@ -49,5 +49,28 @@ mod tests { assert_eq!(tree.first(), Some(&1000)); assert_eq!(tree.last(), Some(&1099)); + + assert_eq!(tree.get(100), None); + assert_eq!(tree.get_mut(100), None); + } + + #[test] + fn push_64() { + let mut tree = RadixTree::::new(); + + for i in 0..64 { + tree.push(i); + assert_eq!(tree.len(), i + 1); + } + + for i in 0..64 { + assert_eq!(tree.get(i), Some(&i)); + assert_eq!(tree.get_mut(i), Some(&mut i.clone())); + } + + for i in 64..256 { + assert_eq!(tree.get(i), None); + assert_eq!(tree.get_mut(i), None); + } } } diff --git a/radix_tree/src/radix_tree.rs b/radix_tree/src/radix_tree.rs index 63aa2b4..9147f50 100644 --- a/radix_tree/src/radix_tree.rs +++ b/radix_tree/src/radix_tree.rs @@ -222,7 +222,7 @@ impl RadixTree { } } - fn index_path(&self, mut i: usize) -> Vec { + fn index_path(&self, mut i: usize) -> Option> { let mut path = vec![0; self.depth()]; for p in path.iter_mut().rev() { @@ -230,7 +230,44 @@ impl RadixTree { i /= N; } - path + match i { + 0 => Some(path), + _ => None, + } + } + + pub fn get(&self, i: usize) -> Option<&T> { + let mut tree = self; + + for p in tree.index_path(i)? { + match tree.data() { + RadixTreeData::Meta(meta) => { + tree = meta.get(p)?; + } + RadixTreeData::Leaves(leaves) => { + return leaves.get(p); + } + } + } + + None + } + + pub fn get_mut(&mut self, i: usize) -> Option<&mut T> { + let mut tree = self; + + for p in tree.index_path(i)? { + match tree.data_mut() { + RadixTreeData::Meta(meta) => { + tree = meta.get_mut(p)?; + } + RadixTreeData::Leaves(leaves) => { + return leaves.get_mut(p); + } + } + } + + None } } @@ -246,7 +283,7 @@ impl Index for RadixTree { fn index(&self, i: usize) -> &T { let mut tree = self; - for p in tree.index_path(i) { + for p in tree.index_path(i).expect("Out of bounds") { match tree.data() { RadixTreeData::Meta(meta) => { tree = &meta[p]; @@ -265,7 +302,7 @@ impl IndexMut for RadixTree { fn index_mut(&mut self, i: usize) -> &mut T { let mut tree = self; - for p in tree.index_path(i) { + for p in tree.index_path(i).expect("Out of bounds") { match tree.data_mut() { RadixTreeData::Meta(meta) => { tree = &mut meta[p];