From 9594058fd5ee07f72f0b56b6db21a56ee5057eaf Mon Sep 17 00:00:00 2001
From: floor-licker <kingofthecoast78@gmail.com>
Date: Fri, 30 Jan 2026 18:17:45 -0500
Subject: [PATCH] feat(book): apply websocket book updates

---
 src/book.rs                 | 69 +++++++++++++++++++++++++++++++++++++
 tests/no_alloc_hot_paths.rs | 37 ++++++++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/src/book.rs b/src/book.rs
index 244a97d..3de5c6c 100644
--- a/src/book.rs
+++ b/src/book.rs
@@ -396,6 +396,75 @@ impl OrderBook {
         Ok(())
     }
 
+    /// Apply a WebSocket `book` update for this token.
+    ///
+    /// The official Polymarket CLOB WebSocket `book` event contains batches of
+    /// price levels for both sides. Unlike `apply_delta_fast`, this method can
+    /// apply many levels that share the same message timestamp.
+    ///
+    /// Notes:
+    /// - This performs upserts (update/insert/remove) for the provided levels.
+    /// - It does **not** infer removals for levels omitted from the message.
+    /// - Insertions of *new* price levels may allocate (BTreeMap node growth).
+    pub fn apply_book_update(&mut self, update: &BookUpdate) -> Result<()> {
+        if update.asset_id != self.token_id {
+            return Err(PolyfillError::validation("Token ID mismatch"));
+        }
+
+        // Use the exchange-provided timestamp as our monotonic sequence marker.
+        // This is less strict than the REST/legacy delta sequence but works for
+        // ignoring obviously stale book snapshots.
+        if update.timestamp <= self.sequence {
+            return Ok(());
+        }
+
+        self.sequence = update.timestamp;
+        self.timestamp = chrono::DateTime::<Utc>::from_timestamp(update.timestamp as i64, 0)
+            .unwrap_or_else(Utc::now);
+
+        // Apply bids (BUY) and asks (SELL) as level upserts.
+        for level in &update.bids {
+            let price_ticks = decimal_to_price(level.price)
+                .map_err(|_| PolyfillError::validation("Invalid price"))?;
+            let size_units = decimal_to_qty(level.size)
+                .map_err(|_| PolyfillError::validation("Invalid size"))?;
+
+            if let Some(tick_size_ticks) = self.tick_size_ticks {
+                if tick_size_ticks > 0 && !price_ticks.is_multiple_of(tick_size_ticks) {
+                    return Err(PolyfillError::validation("Price not aligned to tick size"));
+                }
+            }
+
+            if size_units == 0 {
+                self.bids.remove(&price_ticks);
+            } else {
+                self.bids.insert(price_ticks, size_units);
+            }
+        }
+
+        for level in &update.asks {
+            let price_ticks = decimal_to_price(level.price)
+                .map_err(|_| PolyfillError::validation("Invalid price"))?;
+            let size_units = decimal_to_qty(level.size)
+                .map_err(|_| PolyfillError::validation("Invalid size"))?;
+
+            if let Some(tick_size_ticks) = self.tick_size_ticks {
+                if tick_size_ticks > 0 && !price_ticks.is_multiple_of(tick_size_ticks) {
+                    return Err(PolyfillError::validation("Price not aligned to tick size"));
+                }
+            }
+
+            if size_units == 0 {
+                self.asks.remove(&price_ticks);
+            } else {
+                self.asks.insert(price_ticks, size_units);
+            }
+        }
+
+        self.trim_depth();
+        Ok(())
+    }
+
     /// Apply a bid-side delta (someone wants to buy) - LEGACY VERSION
     /// If size is 0, it means "remove this price level entirely"
     /// Otherwise, set the total size at this price level
diff --git a/tests/no_alloc_hot_paths.rs b/tests/no_alloc_hot_paths.rs
index 441b1f5..6e7465e 100644
--- a/tests/no_alloc_hot_paths.rs
+++ b/tests/no_alloc_hot_paths.rs
@@ -2,9 +2,11 @@ use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::Cell;
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};
+use std::str::FromStr;
 
 use chrono::Utc;
 use polyfill_rs::{OrderBookImpl, Side};
+use rust_decimal::Decimal;
 
 thread_local! {
     static ALLOCATIONS: Cell<usize> = const { Cell::new(0) };
@@ -127,3 +129,38 @@ fn no_alloc_apply_delta_fast_existing_level_update() {
         .unwrap();
     guard.assert_no_allocations();
 }
+
+#[test]
+fn no_alloc_apply_book_update_existing_levels() {
+    let asset_id = "test_asset_id";
+    let token_hash = token_id_hash(asset_id);
+    let mut book = OrderBookImpl::new(asset_id.to_string(), 100);
+
+    // Allocate during setup: create initial price levels.
+    book.apply_delta_fast(mk_delta(token_hash, Side::BUY, 7500, 1_000_000, 1))
+        .unwrap();
+    book.apply_delta_fast(mk_delta(token_hash, Side::SELL, 7600, 1_000_000, 2))
+        .unwrap();
+
+    let update = polyfill_rs::types::BookUpdate {
+        asset_id: asset_id.to_string(),
+        market: "0xabc".to_string(),
+        timestamp: 10,
+        bids: vec![polyfill_rs::types::OrderSummary {
+            price: Decimal::from_str("0.75").unwrap(),
+            size: Decimal::from_str("200.0").unwrap(),
+        }],
+        asks: vec![polyfill_rs::types::OrderSummary {
+            price: Decimal::from_str("0.76").unwrap(),
+            size: Decimal::from_str("50.0").unwrap(),
+        }],
+        hash: None,
+    };
+
+    // Warm up TLS access before measuring (defensive).
+    let _ = allocation_count();
+
+    let guard = NoAllocGuard::new();
+    book.apply_book_update(&update).unwrap();
+    guard.assert_no_allocations();
+}