Add indexing

radix_tree/src/lib.rs

@@ -7,10 +7,12 @@ mod tests {
   use super::*;
 
   #[test]
-  fn initial_len_0() {
+  fn empty_tree() {
     let tree = RadixTree::<usize, 4>::new();
-    assert_eq!(tree.len(), 0);
     assert!(tree.is_empty());
+    assert_eq!(tree.len(), 0);
+    assert_eq!(tree.first(), None);
+    assert_eq!(tree.last(), None);
   }
 
   #[test]
@@ -29,5 +31,23 @@ mod tests {
       tree.push(i);
       assert_eq!(tree.len(), i + 1);
     }
+
+    assert_eq!(tree.first(), Some(&0));
+    assert_eq!(tree.last(), Some(&99));
+
+    for i in 0..100 {
+      assert_eq!(tree[i], i);
+    }
+
+    for i in 0..100 {
+      tree[i] = 1000 + i;
+    }
+
+    for i in 0..100 {
+      assert_eq!(tree[i], 1000 + i);
+    }
+
+    assert_eq!(tree.first(), Some(&1000));
+    assert_eq!(tree.last(), Some(&1099));
   }
 }

radix_tree/src/radix_tree.rs

@@ -1,9 +1,13 @@
-use std::{mem::swap, rc::Rc};
+use std::{
+  mem::swap,
+  ops::{Index, IndexMut},
+  rc::Rc,
+};
 
 use arrayvec::ArrayVec;
 
 #[derive(Clone)]
-pub enum RadixTreeData<T, const N: usize> {
+enum RadixTreeData<T, const N: usize> {
   Meta(ArrayVec<RadixTree<T, N>, N>),
   Leaves(ArrayVec<T, N>),
 }
@@ -28,6 +32,18 @@ impl<T: Clone, const N: usize> RadixTree<T, N> {
     Rc::make_mut(&mut self.0)
   }
 
+  fn depth(&self) -> usize {
+    let mut res = 1;
+    let mut tree = self;
+
+    while let RadixTreeData::Meta(meta) = tree.data() {
+      tree = &meta[0];
+      res += 1;
+    }
+
+    res
+  }
+
   pub fn is_empty(&self) -> bool {
     match self.data() {
       RadixTreeData::Meta(_) => false,
@@ -212,3 +228,57 @@ impl<T: Clone, const N: usize> Default for RadixTree<T, N> {
     Self::new()
   }
 }
+
+impl<T: Clone, const N: usize> Index<usize> for RadixTree<T, N> {
+  type Output = T;
+
+  fn index(&self, mut i: usize) -> &T {
+    let mut path = vec![0; self.depth()];
+
+    for p in path.iter_mut().rev() {
+      *p = i % N;
+      i /= N;
+    }
+
+    let mut tree = self;
+
+    for p in path {
+      match tree.data() {
+        RadixTreeData::Meta(meta) => {
+          tree = &meta[p];
+        }
+        RadixTreeData::Leaves(leaves) => {
+          return &leaves[p];
+        }
+      }
+    }
+
+    panic!("Out of bounds");
+  }
+}
+
+impl<T: Clone, const N: usize> IndexMut<usize> for RadixTree<T, N> {
+  fn index_mut(&mut self, mut i: usize) -> &mut T {
+    let mut path = vec![0; self.depth()];
+
+    for p in path.iter_mut().rev() {
+      *p = i % N;
+      i /= N;
+    }
+
+    let mut tree = self;
+
+    for p in path {
+      match tree.data_mut() {
+        RadixTreeData::Meta(meta) => {
+          tree = &mut meta[p];
+        }
+        RadixTreeData::Leaves(leaves) => {
+          return &mut leaves[p];
+        }
+      }
+    }
+
+    panic!("Out of bounds");
+  }
+}