From b5f5959585a6c14bd0ece28138594cbb198b684a Mon Sep 17 00:00:00 2001
From: Thomas Ip <thomas@ipthomas.com>
Date: Sun, 6 Feb 2022 14:21:25 +0800
Subject: [PATCH] dex: Add level 2 iterator

Fixes #208
---
 dex/src/critbit.rs |   2 +-
 dex/src/lib.rs     |   1 +
 dex/src/utils.rs   | 129 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 131 insertions(+), 1 deletion(-)
 create mode 100644 dex/src/utils.rs

diff --git a/dex/src/critbit.rs b/dex/src/critbit.rs
index 73ce800b..fa09567a 100644
--- a/dex/src/critbit.rs
+++ b/dex/src/critbit.rs
@@ -1023,7 +1023,7 @@ mod tests {
 
         let mut rng = StdRng::from_entropy();
 
-        for i in 0..100 {
+        for _i in 0..100 {
             let offset = rng.gen();
             let key = rng.gen();
             let owner = rng.gen();
diff --git a/dex/src/lib.rs b/dex/src/lib.rs
index 7c58d40a..e52c6827 100644
--- a/dex/src/lib.rs
+++ b/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;
diff --git a/dex/src/utils.rs b/dex/src/utils.rs
new file mode 100644
index 00000000..9a2d52a4
--- /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, 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()));
+    }
+}