diff --git a/dex/src/critbit.rs b/dex/src/critbit.rs index f29ce9d..ff8ac53 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)] @@ -724,29 +753,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)] @@ -825,11 +841,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()); @@ -931,7 +946,7 @@ mod tests { for i in 0..100_000 { slab.check_invariants(); let model_state = model.values().collect::>(); - 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 { @@ -995,4 +1010,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())); + } } diff --git a/dex/src/lib.rs b/dex/src/lib.rs index 7c58d40..382cb06 100644 --- a/dex/src/lib.rs +++ b/dex/src/lib.rs @@ -8,10 +8,11 @@ pub mod error; mod tests; pub mod critbit; -mod fees; +pub 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; diff --git a/dex/src/utils.rs b/dex/src/utils.rs new file mode 100644 index 0000000..e4450ff --- /dev/null +++ b/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, Clone, Copy, PartialEq, Eq)] +pub struct Level { + pub price: u64, + pub quantity: u64, +} + +pub struct LevelsIterator<'a> { + leafs: Peekable>, +} + +impl Iterator for LevelsIterator<'_> { + type Item = Level; + + fn next(&mut self) -> Option { + let mut level: Option = 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())); + } +}