dex: Add leaf node iterator

dex/src/critbit.rs

@@ -124,6 +124,35 @@ impl LeafNode {
     }
 }
 
+pub struct LeafNodeIterator<'a> {
+    descending: bool,
+    slab: &'a Slab,
+    stack: Vec<NodeHandle>,
+}
+
+impl<'a> Iterator for LeafNodeIterator<'a> {
+    type Item = &'a LeafNode;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(node_handle) = self.stack.pop() {
+            let node_ref = self.slab.get(node_handle).unwrap().case().unwrap();
+            match node_ref {
+                NodeRef::Inner(inner) => {
+                    if self.descending {
+                        self.stack.push(inner.children[0]);
+                        self.stack.push(inner.children[1]);
+                    } else {
+                        self.stack.push(inner.children[1]);
+                        self.stack.push(inner.children[0]);
+                    }
+                }
+                NodeRef::Leaf(leaf) => return Some(leaf),
+            }
+        }
+        None
+    }
+}
+
 #[derive(Copy, Clone)]
 #[repr(packed)]
 #[allow(dead_code)]
@@ -726,6 +755,18 @@ impl Slab {
         self.remove_by_key(self.get(self.find_max()?)?.key()?)
     }
 
+    pub fn iter(&self, descending: bool) -> LeafNodeIterator<'_> {
+        let mut stack = vec![];
+        if let Some(node_handle) = self.root() {
+            stack.push(node_handle);
+        }
+        LeafNodeIterator {
+            descending,
+            slab: self,
+            stack,
+        }
+    }
+
     #[cfg(test)]
     fn traverse(&self) -> Vec<&LeafNode> {
         fn walk_rec<'a>(slab: &'a Slab, sub_root: NodeHandle, buf: &mut Vec<&'a LeafNode>) {
@@ -827,11 +868,10 @@ mod tests {
     use super::*;
     use bytemuck::bytes_of;
     use rand::prelude::*;
+    use std::collections::BTreeMap;
 
     #[test]
     fn simulate_find_min() {
-        use std::collections::BTreeMap;
-
         for trial in 0..10u64 {
             let mut aligned_buf = vec![0u64; 10_000];
             let bytes: &mut [u8] = cast_slice_mut(aligned_buf.as_mut_slice());
@@ -997,4 +1037,28 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_leaf_node_iter() {
+        let mut aligned_buf = vec![0u64; 10_000];
+        let bytes: &mut [u8] = cast_slice_mut(aligned_buf.as_mut_slice());
+
+        let slab: &mut Slab = Slab::new(bytes);
+        let mut model: BTreeMap<u128, LeafNode> = BTreeMap::new();
+
+        let mut rng = StdRng::from_entropy();
+
+        for i in 0..100 {
+            let offset = rng.gen();
+            let key = rng.gen();
+            let owner = rng.gen();
+            let qty = rng.gen();
+            let leaf = LeafNode::new(offset, key, owner, qty, FeeTier::Base, 0);
+
+            slab.insert_leaf(&leaf).unwrap();
+            model.insert(key, leaf).ok_or(()).unwrap_err();
+        }
+
+        assert!(slab.iter(false).eq(model.values()));
+    }
 }