Add public API to iterate over L2/L3 data

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,29 +755,16 @@ impl Slab {
         self.remove_by_key(self.get(self.find_max()?)?.key()?)
     }
 
-    #[cfg(test)]
-    fn traverse(&self) -> Vec<&LeafNode> {
-        fn walk_rec<'a>(slab: &'a Slab, sub_root: NodeHandle, buf: &mut Vec<&'a LeafNode>) {
-            match slab.get(sub_root).unwrap().case().unwrap() {
-                NodeRef::Leaf(leaf) => {
-                    buf.push(leaf);
-                }
-                NodeRef::Inner(inner) => {
-                    walk_rec(slab, inner.children[0], buf);
-                    walk_rec(slab, inner.children[1], buf);
-                }
-            }
+    pub fn iter(&self, descending: bool) -> LeafNodeIterator<'_> {
+        let mut stack = vec![];
+        if let Some(node_handle) = self.root() {
+            stack.push(node_handle);
         }
-
-        let mut buf = Vec::with_capacity(self.header().leaf_count as usize);
-        if let Some(r) = self.root() {
-            walk_rec(self, r, &mut buf);
-        }
-        if buf.len() != buf.capacity() {
-            self.hexdump();
+        LeafNodeIterator {
+            descending,
+            slab: self,
+            stack,
         }
-        assert_eq!(buf.len(), buf.capacity());
-        buf
     }
 
     #[cfg(test)]
@@ -827,11 +843,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());
@@ -933,7 +948,7 @@ mod tests {
             for i in 0..100_000 {
                 slab.check_invariants();
                 let model_state = model.values().collect::<Vec<_>>();
-                let slab_state = slab.traverse();
+                let slab_state: Vec<&LeafNode> = slab.iter(false).collect();
                 assert_eq!(model_state, slab_state);
 
                 match weights[dist.sample(&mut rng)].0 {
@@ -997,4 +1012,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()));
+    }
 }

dex/src/lib.rs

@@ -12,6 +12,7 @@ mod fees;
 pub mod instruction;
 pub mod matching;
 pub mod state;
+pub mod utils;
 
 #[cfg(all(feature = "program", not(feature = "no-entrypoint")))]
 use solana_program::entrypoint;

dex/src/utils.rs

@@ -0,0 +1,129 @@
+use std::iter::Peekable;
+
+use crate::{
+    critbit::{LeafNodeIterator, Slab},
+    matching::Side,
+};
+
+pub struct OrderBook<'a> {
+    side: Side,
+    slab: &'a Slab,
+}
+
+impl<'a> OrderBook<'a> {
+    pub fn new(side: Side, slab: &'a Slab) -> Self {
+        OrderBook { side, slab }
+    }
+
+    /// Iterate over level 2 data
+    pub fn levels(&'a self) -> LevelsIterator<'a> {
+        let is_descending = self.side == Side::Bid;
+        LevelsIterator {
+            leafs: self.slab.iter(is_descending).peekable(),
+        }
+    }
+}
+
+/// Level 2 data
+///
+/// Values are in lot sizes.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub struct Level {
+    pub price: u64,
+    pub quantity: u64,
+}
+
+pub struct LevelsIterator<'a> {
+    leafs: Peekable<LeafNodeIterator<'a>>,
+}
+
+impl Iterator for LevelsIterator<'_> {
+    type Item = Level;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut level: Option<Level> = None;
+        while let Some(leaf) = self.leafs.peek() {
+            let price = leaf.price().get();
+            let quantity = leaf.quantity();
+            match &mut level {
+                Some(Level {
+                    price: cur_price,
+                    quantity: cur_quantity,
+                }) => {
+                    if price == *cur_price {
+                        *cur_quantity += quantity;
+                    } else {
+                        return level;
+                    }
+                }
+                None => {
+                    level = Some(Level { price, quantity });
+                }
+            }
+            self.leafs.next();
+        }
+        level
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{critbit::LeafNode, fees::FeeTier};
+
+    use bytemuck::cast_slice_mut;
+    use rand::prelude::*;
+
+    #[test]
+    fn test_levels() {
+        let mut aligned_buf = vec![0u64; 10_000];
+        let bytes: &mut [u8] = cast_slice_mut(aligned_buf.as_mut_slice());
+        let slab = Slab::new(bytes);
+
+        let mut rng = StdRng::from_entropy();
+
+        let orders: Vec<(u64, u64)> = vec![
+            (1000, 50),
+            (1000, 25),
+            (1300, 100),
+            (1300, 40),
+            (1300, 250),
+            (1480, 36),
+            (1495, 3000),
+            (1495, 340),
+        ];
+        for (i, (price, qty)) in orders.into_iter().enumerate() {
+            let offset = rng.gen();
+            let owner = rng.gen();
+            let key = ((price as u128) << 64) | (!(i as u64) as u128);
+
+            slab.insert_leaf(&LeafNode::new(offset, key, owner, qty, FeeTier::Base, 0))
+                .unwrap();
+        }
+
+        let ask_levels = vec![
+            Level {
+                price: 1000,
+                quantity: 75,
+            },
+            Level {
+                price: 1300,
+                quantity: 390,
+            },
+            Level {
+                price: 1480,
+                quantity: 36,
+            },
+            Level {
+                price: 1495,
+                quantity: 3340,
+            },
+        ];
+
+        let orderbook_asks = OrderBook::new(Side::Ask, slab);
+        assert!(ask_levels.clone().into_iter().eq(orderbook_asks.levels()));
+
+        let orderbook_bids = OrderBook::new(Side::Bid, slab);
+        assert!(ask_levels.into_iter().rev().eq(orderbook_bids.levels()));
+    }
+}