Add get, get_mut

voltrevo · Aug 17, 2023 · 0eb7a51 · 0eb7a51
1 parent 949c4ee
commit 0eb7a51
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 4 deletions.
diff --git 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::<usize, 4>::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
@@ -222,15 +222,52 @@ impl<T: Clone, const N: usize> RadixTree<T, N> {
     }
   }
 
-  fn index_path(&self, mut i: usize) -> Vec<usize> {
+  fn index_path(&self, mut i: usize) -> Option<Vec<usize>> {
     let mut path = vec![0; self.depth()];
 
     for p in path.iter_mut().rev() {
       *p = i % N;
       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<T: Clone, const N: usize> Index<usize> for RadixTree<T, N> {
   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<T: Clone, const N: usize> IndexMut<usize> for RadixTree<T, N> {
   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];