diff --git a/dex/src/critbit.rs b/dex/src/critbit.rs index 6dfdcb2c..18586fae 100644 --- a/dex/src/critbit.rs +++ b/dex/src/critbit.rs @@ -124,6 +124,35 @@ impl LeafNode { } } +pub struct LeafNodeIterator<'a> { + descending: bool, + slab: &'a Slab, + stack: Vec, +} + +impl<'a> Iterator for LeafNodeIterator<'a> { + type Item = &'a LeafNode; + + fn next(&mut self) -> Option { + 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 = 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())); + } }