pragma_common/orderbook/
mod.rs

1use std::collections::BTreeMap;
2
3use bigdecimal::{BigDecimal, FromPrimitive, ToPrimitive, Zero};
4
5use crate::entries::depth::DepthLevel;
6
7#[derive(Debug, thiserror::Error)]
8pub enum OrderbookError {
9    #[error("Received invalid bid")]
10    InvalidBid,
11    #[error("Received invalid ask")]
12    InvalidAsk,
13}
14
15/// A orderbook order book that handles updates and snapshots.
16///
17/// This struct wraps the core `InnerOrderbook` and manages pending updates before a snapshot is applied.
18#[derive(Debug, Clone)]
19pub struct Orderbook {
20    /// The core order book with bids and asks
21    orderbook: InnerOrderbook,
22    /// Tracks whether a snapshot has been applied
23    snapshot_applied: bool,
24    /// Pending updates before snapshot
25    pending_updates: Option<PendingOrderbookUpdate>,
26}
27
28impl Orderbook {
29    /// Create a new orderbook with the given depth.
30    pub fn new(depth: usize) -> Self {
31        Self {
32            orderbook: InnerOrderbook {
33                depth,
34                bids: BTreeMap::new(),
35                asks: BTreeMap::new(),
36                last_update_id: 0,
37            },
38            snapshot_applied: false,
39            pending_updates: None,
40        }
41    }
42
43    /// Returns the last update id applied to the current orderbook.
44    pub const fn last_update_id(&self) -> u64 {
45        self.orderbook.last_update_id()
46    }
47
48    /// Returns the next update id & increment the state of the inner `Orderbook`.
49    pub const fn next_update_id(&self) -> u64 {
50        self.orderbook.next_update_id()
51    }
52
53    /// Returns true if a snapshot has been applied to the current orderbook, else false.
54    pub const fn snapshot_was_applied(&self) -> bool {
55        self.snapshot_applied
56    }
57
58    /// Returns the current `best_bid` - if any.
59    pub fn best_bid(&self) -> Option<&BigDecimal> {
60        self.orderbook.best_bid()
61    }
62
63    /// Returns the current `best_ask` - if any.
64    pub fn best_ask(&self) -> Option<&BigDecimal> {
65        self.orderbook.best_ask()
66    }
67
68    /// Get the mid price (average of best bid and ask).
69    pub fn mid_price(&self) -> Option<BigDecimal> {
70        self.orderbook.mid_price()
71    }
72
73    /// Compute depth at a given percentage.
74    pub fn depth(&self, percentage: f64) -> Option<DepthLevel> {
75        self.orderbook.depth(percentage)
76    }
77
78    /// Applies an update to the order book.
79    ///
80    /// Before a snapshot, updates are merged into `pending_updates`. After a snapshot, updates are
81    /// applied directly if their `update_id` is newer than the current `last_update_id`.
82    pub fn apply_update(
83        &mut self,
84        bids: Vec<(f64, f64)>,
85        asks: Vec<(f64, f64)>,
86        update_id: u64,
87    ) -> Result<(), OrderbookError> {
88        let bids_bd: Vec<(BigDecimal, f64)> = bids
89            .into_iter()
90            .map(|(p, q)| BigDecimal::from_f64(p).map(|bd| (bd, q)))
91            .collect::<Option<_>>()
92            .ok_or(OrderbookError::InvalidAsk)?;
93
94        let asks_bd: Vec<(BigDecimal, f64)> = asks
95            .into_iter()
96            .map(|(p, q)| BigDecimal::from_f64(p).map(|bd| (bd, q)))
97            .collect::<Option<_>>()
98            .ok_or(OrderbookError::InvalidBid)?;
99
100        if self.snapshot_applied {
101            // After snapshot, apply only if update_id is newer
102            if update_id > self.orderbook.last_update_id {
103                self.orderbook.apply_update(bids_bd, asks_bd, update_id);
104            }
105        } else {
106            // Before snapshot, merge into pending updates
107            if let Some(pending) = self.pending_updates.as_mut() {
108                pending.merge(bids_bd, asks_bd, update_id);
109            } else {
110                self.pending_updates =
111                    Some(PendingOrderbookUpdate::new(bids_bd, asks_bd, update_id));
112            }
113        }
114
115        Ok(())
116    }
117
118    /// Applies a snapshot, replacing the current order book state.
119    ///
120    /// After applying the snapshot, any pending updates with a higher `update_id` are applied.
121    pub fn apply_snapshot(
122        &mut self,
123        bids: Vec<(f64, f64)>,
124        asks: Vec<(f64, f64)>,
125        update_id: u64,
126    ) -> Result<(), OrderbookError> {
127        // Convert f64 prices to BigDecimal for precision
128        let bids_bd: Vec<(BigDecimal, f64)> = bids
129            .into_iter()
130            .map(|(p, q)| BigDecimal::from_f64(p).map(|bd| (bd, q)))
131            .collect::<Option<_>>()
132            .ok_or(OrderbookError::InvalidBid)?;
133        let asks_bd: Vec<(BigDecimal, f64)> = asks
134            .into_iter()
135            .map(|(p, q)| BigDecimal::from_f64(p).map(|bd| (bd, q)))
136            .collect::<Option<_>>()
137            .ok_or(OrderbookError::InvalidAsk)?;
138
139        self.orderbook
140            .clear_and_apply_snapshot(bids_bd, asks_bd, update_id);
141        self.apply_pending_updates_after_snapshot();
142        self.snapshot_applied = true;
143        Ok(())
144    }
145
146    /// Applies pending updates after a snapshot, filtering by update ID.
147    fn apply_pending_updates_after_snapshot(&mut self) {
148        // No pending updates to process
149        let Some(pending) = self.pending_updates.take() else {
150            return;
151        };
152
153        // The pending updates are too old - irrelevant
154        if pending.latest_update_id <= self.orderbook.last_update_id {
155            return;
156        }
157
158        let bids: Vec<(BigDecimal, f64)> = pending
159            .bids
160            .into_iter()
161            .filter_map(|(price, (qty, uid))| {
162                if uid > self.orderbook.last_update_id {
163                    Some((price, qty))
164                } else {
165                    None
166                }
167            })
168            .collect();
169
170        let asks: Vec<(BigDecimal, f64)> = pending
171            .asks
172            .into_iter()
173            .filter_map(|(price, (qty, uid))| {
174                if uid > self.orderbook.last_update_id {
175                    Some((price, qty))
176                } else {
177                    None
178                }
179            })
180            .collect();
181
182        if !bids.is_empty() || !asks.is_empty() {
183            self.orderbook
184                .apply_update(bids, asks, pending.latest_update_id);
185        }
186    }
187}
188
189/// Represents a merged order book update with per-level update IDs.
190#[derive(Debug, Clone)]
191struct PendingOrderbookUpdate {
192    bids: BTreeMap<BigDecimal, (f64, u64)>, // price -> (quantity, update_id)
193    asks: BTreeMap<BigDecimal, (f64, u64)>, // price -> (quantity, update_id)
194    latest_update_id: u64,
195}
196
197impl PendingOrderbookUpdate {
198    /// Creates a new update from bids, asks, and an update ID.
199    fn new(bids: Vec<(BigDecimal, f64)>, asks: Vec<(BigDecimal, f64)>, update_id: u64) -> Self {
200        let mut bids_map = BTreeMap::new();
201        let mut asks_map = BTreeMap::new();
202
203        for (price, qty) in bids {
204            if qty > 0.0 {
205                bids_map.insert(price, (qty, update_id));
206            }
207        }
208        for (price, qty) in asks {
209            if qty > 0.0 {
210                asks_map.insert(price, (qty, update_id));
211            }
212        }
213
214        Self {
215            bids: bids_map,
216            asks: asks_map,
217            latest_update_id: update_id,
218        }
219    }
220
221    /// Merges a new update into this one, keeping the latest data per price level.
222    ///
223    /// Updates are only applied if the new `update_id` is greater than or equal to the existing one.
224    fn merge(
225        &mut self,
226        bids: Vec<(BigDecimal, f64)>,
227        asks: Vec<(BigDecimal, f64)>,
228        update_id: u64,
229    ) {
230        for (price, qty) in bids {
231            if qty == 0.0 {
232                self.bids.remove(&price);
233            } else {
234                self.bids
235                    .entry(price)
236                    .and_modify(|(existing_qty, existing_id)| {
237                        if update_id > *existing_id {
238                            *existing_qty = qty;
239                            *existing_id = update_id;
240                        }
241                    })
242                    .or_insert((qty, update_id));
243            }
244        }
245        for (price, qty) in asks {
246            if qty == 0.0 {
247                self.asks.remove(&price);
248            } else {
249                self.asks
250                    .entry(price)
251                    .and_modify(|(existing_qty, existing_id)| {
252                        if update_id > *existing_id {
253                            *existing_qty = qty;
254                            *existing_id = update_id;
255                        }
256                    })
257                    .or_insert((qty, update_id));
258            }
259        }
260        self.latest_update_id = self.latest_update_id.max(update_id);
261    }
262}
263
264/// The core order book structure.
265#[derive(Debug, Clone)]
266struct InnerOrderbook {
267    depth: usize,                    // Max depth level of the orderbook
268    bids: BTreeMap<BigDecimal, f64>, // price -> quantity
269    asks: BTreeMap<BigDecimal, f64>, // price -> quantity
270    last_update_id: u64,
271}
272
273impl InnerOrderbook {
274    const fn last_update_id(&self) -> u64 {
275        self.last_update_id
276    }
277
278    /// Returns the next update id & increment the state.
279    const fn next_update_id(&self) -> u64 {
280        self.last_update_id + 1
281    }
282
283    /// Clear the current orderbook.
284    fn clear(&mut self) {
285        self.bids.clear();
286        self.asks.clear();
287    }
288
289    /// Clear the current orderbook & replace it with a new snapshot.
290    fn clear_and_apply_snapshot(
291        &mut self,
292        bids: Vec<(BigDecimal, f64)>,
293        asks: Vec<(BigDecimal, f64)>,
294        update_id: u64,
295    ) {
296        self.clear();
297        for (price, qty) in bids {
298            if qty > 0.0 {
299                self.bids.insert(price, qty);
300            }
301        }
302        for (price, qty) in asks {
303            if qty > 0.0 {
304                self.asks.insert(price, qty);
305            }
306        }
307        self.last_update_id = update_id;
308    }
309
310    fn apply_update(
311        &mut self,
312        bids: Vec<(BigDecimal, f64)>,
313        asks: Vec<(BigDecimal, f64)>,
314        update_id: u64,
315    ) {
316        for (price, qty) in bids {
317            if qty.is_zero() {
318                self.bids.remove(&price);
319            } else {
320                self.bids.insert(price, qty);
321            }
322        }
323        for (price, qty) in asks {
324            if qty.is_zero() {
325                self.asks.remove(&price);
326            } else {
327                self.asks.insert(price, qty);
328            }
329        }
330        self.last_update_id = update_id;
331        self.truncate_to_depth();
332    }
333
334    fn best_bid(&self) -> Option<&BigDecimal> {
335        self.bids.iter().next_back().map(|(p, _)| p)
336    }
337
338    fn best_ask(&self) -> Option<&BigDecimal> {
339        self.asks.iter().next().map(|(p, _)| p)
340    }
341
342    fn mid_price(&self) -> Option<BigDecimal> {
343        let best_bid = self.best_bid();
344        let best_ask = self.best_ask();
345        match (best_bid, best_ask) {
346            (Some(bid), Some(ask)) => Some((bid + ask) / BigDecimal::from(2)),
347            _ => None,
348        }
349    }
350
351    fn depth(&self, percentage: f64) -> Option<DepthLevel> {
352        let mid = self.mid_price()?;
353        let mid_price = mid.to_f64()?;
354
355        let percentage_bd = BigDecimal::from_f64(percentage)?;
356        let lower = &mid * (BigDecimal::from(1) - &percentage_bd);
357        let upper = &mid * (BigDecimal::from(1) + &percentage_bd);
358
359        let bid_depth: f64 = self.bids.range(&lower..=&mid).map(|(_, &qty)| qty).sum();
360        let ask_depth: f64 = self.asks.range(&mid..=&upper).map(|(_, &qty)| qty).sum();
361
362        Some(DepthLevel {
363            bid: bid_depth * mid_price,
364            ask: ask_depth * mid_price,
365            percentage,
366        })
367    }
368
369    fn truncate_to_depth(&mut self) {
370        while self.bids.len() > self.depth {
371            self.bids.pop_first();
372        }
373
374        while self.asks.len() > self.depth {
375            self.asks.pop_last();
376        }
377    }
378}
379
380// Unit tests
381#[cfg(test)]
382mod tests {
383    use super::*;
384
385    #[test]
386    fn test_new_orderbook() {
387        let ob = Orderbook::new(10);
388        assert_eq!(ob.last_update_id(), 0);
389        assert!(!ob.snapshot_was_applied());
390        assert!(ob.pending_updates.is_none());
391        assert!(ob.mid_price().is_none());
392    }
393
394    #[test]
395    fn test_apply_update_before_snapshot() {
396        let mut ob = Orderbook::new(10);
397        ob.apply_update(vec![(100.0, 10.0), (99.0, 5.0)], vec![(101.0, 15.0)], 1)
398            .unwrap();
399
400        let pending = ob.pending_updates.as_ref().unwrap();
401        assert_eq!(pending.bids.len(), 2);
402        assert_eq!(pending.asks.len(), 1);
403        assert_eq!(pending.latest_update_id, 1);
404        assert_eq!(pending.bids.get(&BigDecimal::from(100)), Some(&(10.0, 1)));
405        assert_eq!(pending.bids.get(&BigDecimal::from(99)), Some(&(5.0, 1)));
406        assert_eq!(pending.asks.get(&BigDecimal::from(101)), Some(&(15.0, 1)));
407        assert_eq!(ob.orderbook.bids.len(), 0); // Not applied yet
408    }
409
410    #[test]
411    fn test_merge_updates_before_snapshot() {
412        let mut ob = Orderbook::new(10);
413        ob.apply_update(vec![(100.0, 10.0)], vec![(101.0, 15.0)], 1)
414            .unwrap();
415        ob.apply_update(
416            vec![(100.0, 0.0), (99.0, 5.0)], // Remove 100.0, add 99.0
417            vec![(101.0, 20.0)],             // Update 101.0
418            2,
419        )
420        .unwrap();
421
422        let pending = ob.pending_updates.as_ref().unwrap();
423        assert_eq!(pending.bids.len(), 1); // 100.0 removed
424        assert_eq!(pending.asks.len(), 1);
425        assert_eq!(pending.latest_update_id, 2);
426        assert_eq!(pending.bids.get(&BigDecimal::from(99)), Some(&(5.0, 2)));
427        assert_eq!(pending.asks.get(&BigDecimal::from(101)), Some(&(20.0, 2)));
428    }
429
430    #[test]
431    fn test_apply_snapshot_no_pending() {
432        let mut ob = Orderbook::new(10);
433        ob.apply_snapshot(vec![(100.0, 10.0), (99.0, 5.0)], vec![(101.0, 15.0)], 5)
434            .unwrap();
435
436        assert!(ob.snapshot_was_applied());
437        assert_eq!(ob.last_update_id(), 5);
438        assert!(ob.pending_updates.is_none());
439        assert_eq!(ob.orderbook.bids.len(), 2);
440        assert_eq!(ob.orderbook.asks.len(), 1);
441        assert_eq!(ob.orderbook.bids.get(&BigDecimal::from(100)), Some(&10.0));
442        assert_eq!(ob.mid_price(), Some(BigDecimal::from_f64(100.5).unwrap()));
443    }
444
445    #[test]
446    fn test_apply_snapshot_with_pending() {
447        let mut ob = Orderbook::new(10);
448        ob.apply_update(vec![(100.0, 10.0)], vec![(102.0, 20.0)], 1)
449            .unwrap();
450        ob.apply_update(vec![(99.0, 5.0)], vec![(103.0, 15.0)], 3)
451            .unwrap();
452        ob.apply_snapshot(vec![(100.0, 8.0)], vec![(101.0, 12.0)], 2)
453            .unwrap();
454
455        assert_eq!(ob.last_update_id(), 3);
456        assert_eq!(ob.orderbook.bids.len(), 2); // 100.0 from snapshot, 99.0 from update 3
457        assert_eq!(ob.orderbook.asks.len(), 2); // 101.0 from snapshot, 103.0 from update 3
458        assert_eq!(ob.orderbook.bids.get(&BigDecimal::from(100)), Some(&8.0));
459        assert_eq!(ob.orderbook.bids.get(&BigDecimal::from(99)), Some(&5.0));
460        assert_eq!(ob.orderbook.asks.get(&BigDecimal::from(102)), None); // Update 1 ignored
461    }
462
463    #[test]
464    fn test_update_after_snapshot() {
465        let mut ob = Orderbook::new(10);
466        ob.apply_snapshot(vec![(100.0, 10.0)], vec![(101.0, 15.0)], 5)
467            .unwrap();
468        ob.apply_update(vec![(100.0, 0.0), (99.0, 5.0)], vec![(102.0, 20.0)], 6)
469            .unwrap();
470
471        assert_eq!(ob.last_update_id(), 6);
472        assert_eq!(ob.orderbook.bids.len(), 1); // 100.0 removed
473        assert_eq!(ob.orderbook.asks.len(), 2);
474        assert_eq!(ob.orderbook.bids.get(&BigDecimal::from(99)), Some(&5.0));
475        assert_eq!(ob.orderbook.asks.get(&BigDecimal::from(102)), Some(&20.0));
476    }
477
478    #[test]
479    fn test_old_update_after_snapshot() {
480        let mut ob = Orderbook::new(10);
481        ob.apply_snapshot(vec![(100.0, 10.0)], vec![(101.0, 15.0)], 5)
482            .unwrap();
483        ob.apply_update(vec![(100.0, 5.0)], vec![(101.0, 10.0)], 4)
484            .unwrap();
485
486        assert_eq!(ob.last_update_id(), 5); // Update 4 ignored
487        assert_eq!(ob.orderbook.bids.get(&BigDecimal::from(100)), Some(&10.0));
488        assert_eq!(ob.orderbook.asks.get(&BigDecimal::from(101)), Some(&15.0));
489    }
490
491    #[test]
492    fn test_invalid_price() {
493        let mut ob = Orderbook::new(10);
494        let result = ob.apply_update(vec![(f64::NAN, 10.0)], vec![(101.0, 15.0)], 1);
495        assert!(result.is_err());
496    }
497
498    #[test]
499    #[allow(clippy::float_cmp)]
500    fn test_mid_price_and_depth() {
501        let mut ob = Orderbook::new(10);
502        ob.apply_snapshot(
503            vec![(99.0, 5.0), (100.0, 10.0)],
504            vec![(101.0, 15.0), (102.0, 20.0)],
505            1,
506        )
507        .unwrap();
508
509        let mid = ob.mid_price().unwrap();
510        assert_eq!(mid, BigDecimal::from_f64(100.5).unwrap());
511    }
512
513    #[test]
514    fn test_zero_quantity_in_snapshot() {
515        let mut ob = Orderbook::new(10);
516        ob.apply_snapshot(vec![(100.0, 0.0), (99.0, 5.0)], vec![(101.0, 0.0)], 1)
517            .unwrap();
518
519        assert_eq!(ob.orderbook.bids.len(), 1); // 100.0 ignored
520        assert_eq!(ob.orderbook.asks.len(), 0); // 101.0 ignored
521        assert_eq!(ob.orderbook.bids.get(&BigDecimal::from(99)), Some(&5.0));
522    }
523}
pragma_common/orderbook/mod.rs

pragma_common/orderbook/
mod.rs
