use std::collections::BTreeMap;
use rust_decimal::Decimal;
use digdigdig3::core::types::{OrderBook, OrderbookDelta};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct DecimalKey(Decimal);
impl PartialOrd for DecimalKey {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for DecimalKey {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.0.cmp(&other.0)
}
}
#[derive(Debug, thiserror::Error)]
pub enum OrderBookError {
#[error("sequence gap: book last_update_id={last}, delta prev_update_id={got}")]
SequenceGap { last: u64, got: u64 },
#[error("delta applied before snapshot — call apply_snapshot first")]
NoSnapshot,
#[error("symbol mismatch: tracker={tracker}, incoming={incoming}")]
SymbolMismatch { tracker: String, incoming: String },
#[error("invalid price/size value: {value}")]
InvalidDecimal { value: f64 },
}
#[derive(Debug, Clone)]
pub struct OrderBookTracker {
pub symbol: String,
bids: BTreeMap<DecimalKey, Decimal>,
asks: BTreeMap<DecimalKey, Decimal>,
last_update_id: Option<u64>,
last_timestamp_ms: i64,
has_snapshot: bool,
}
impl OrderBookTracker {
pub fn new(symbol: impl Into<String>) -> Self {
Self {
symbol: symbol.into(),
bids: BTreeMap::new(),
asks: BTreeMap::new(),
last_update_id: None,
last_timestamp_ms: 0,
has_snapshot: false,
}
}
pub fn apply_snapshot(&mut self, snapshot: &OrderBook) -> Result<(), OrderBookError> {
self.bids.clear();
self.asks.clear();
for level in &snapshot.bids {
let price = to_decimal(level.price)?;
let qty = to_decimal(level.size)?;
if qty > Decimal::ZERO {
self.bids.insert(DecimalKey(price), qty);
}
}
for level in &snapshot.asks {
let price = to_decimal(level.price)?;
let qty = to_decimal(level.size)?;
if qty > Decimal::ZERO {
self.asks.insert(DecimalKey(price), qty);
}
}
self.last_update_id = snapshot.last_update_id;
self.last_timestamp_ms = snapshot.timestamp;
self.has_snapshot = true;
Ok(())
}
pub fn apply_delta(&mut self, delta: &OrderbookDelta) -> Result<(), OrderBookError> {
if !self.has_snapshot {
return Err(OrderBookError::NoSnapshot);
}
if let (Some(last), Some(prev)) = (self.last_update_id, delta.prev_update_id) {
if prev != last {
return Err(OrderBookError::SequenceGap { last, got: prev });
}
}
for level in &delta.bids {
let price = to_decimal(level.price)?;
let qty = to_decimal(level.size)?;
if qty == Decimal::ZERO {
self.bids.remove(&DecimalKey(price));
} else {
self.bids.insert(DecimalKey(price), qty);
}
}
for level in &delta.asks {
let price = to_decimal(level.price)?;
let qty = to_decimal(level.size)?;
if qty == Decimal::ZERO {
self.asks.remove(&DecimalKey(price));
} else {
self.asks.insert(DecimalKey(price), qty);
}
}
if let Some(uid) = delta.last_update_id {
self.last_update_id = Some(uid);
}
self.last_timestamp_ms = delta.timestamp;
Ok(())
}
pub fn top_bids(&self, n: usize) -> Vec<(Decimal, Decimal)> {
self.bids.iter().rev().take(n).map(|(k, v)| (k.0, *v)).collect()
}
pub fn top_asks(&self, n: usize) -> Vec<(Decimal, Decimal)> {
self.asks.iter().take(n).map(|(k, v)| (k.0, *v)).collect()
}
pub fn bbo(&self) -> Option<(Decimal, Decimal)> {
let bid = self.bids.iter().rev().next()?.0 .0;
let ask = self.asks.iter().next()?.0 .0;
Some((bid, ask))
}
pub fn mid(&self) -> Option<Decimal> {
let (b, a) = self.bbo()?;
Some((b + a) / Decimal::new(2, 0))
}
pub fn spread(&self) -> Option<Decimal> {
let (b, a) = self.bbo()?;
Some(a - b)
}
pub fn total_bid_volume(&self) -> Decimal {
self.bids.values().copied().sum()
}
pub fn total_ask_volume(&self) -> Decimal {
self.asks.values().copied().sum()
}
pub fn depth(&self) -> (usize, usize) {
(self.bids.len(), self.asks.len())
}
pub fn last_update_id(&self) -> Option<u64> {
self.last_update_id
}
pub fn last_timestamp_ms(&self) -> i64 {
self.last_timestamp_ms
}
pub fn has_snapshot(&self) -> bool {
self.has_snapshot
}
pub fn reset(&mut self) {
self.bids.clear();
self.asks.clear();
self.last_update_id = None;
self.last_timestamp_ms = 0;
self.has_snapshot = false;
}
}
fn to_decimal(v: f64) -> Result<Decimal, OrderBookError> {
Decimal::try_from(v).map_err(|_| OrderBookError::InvalidDecimal { value: v })
}