Simplified Orderbook implementation in Rust with unsafe (fixed version, no UB)

blog-rust-orderbook-unsafe-fixed.rs

#![allow(dead_code)]
// See https://blog.gaborkozar.me/2024/07/01/order-book-in.html
//
// Notice how `Orderbook` is no longer public because it cannot be safely exposed.
// `OrderbookSide` can also no longer be public because it references `Orderbook` in `get_side()`. That in turn causes `OrderKey<Side>` to remain private.

use core::cmp::Ordering;
use std::{
    collections::{BTreeMap, HashMap},
    fmt::{Display, Formatter, Result},
    marker::PhantomData,
};

trait OrderbookSide: Copy + Eq + Sized + 'static {
    type Opposite: OrderbookSide;

    fn is_price_better_than(left: f64, right: f64) -> bool;
    unsafe fn get_orders(orderbook: *mut Orderbook) -> *mut BTreeMap<OrderKey<Self>, Order>;
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct Bid;

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct Ask;

impl OrderbookSide for Bid {
    type Opposite = Ask;

    fn is_price_better_than(left: f64, right: f64) -> bool {
        left > right
    }

    unsafe fn get_orders(orderbook: *mut Orderbook) -> *mut BTreeMap<OrderKey<Self>, Order> {
        &mut (*orderbook).bids
    }
}

impl OrderbookSide for Ask {
    type Opposite = Bid;

    fn is_price_better_than(left: f64, right: f64) -> bool {
        left < right
    }

    unsafe fn get_orders(orderbook: *mut Orderbook) -> *mut BTreeMap<OrderKey<Self>, Order> {
        &mut (*orderbook).asks
    }
}

#[derive(Debug, PartialEq, Clone, Copy)]
struct OrderKey<Side: OrderbookSide> {
    pub price: f64, // assume not-NaN
    pub seq_num: u64,
    _side: PhantomData<Side>,
}

impl<Side: OrderbookSide> Eq for OrderKey<Side> {}

impl<Side: OrderbookSide> Ord for OrderKey<Side> {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.price == other.price {
            self.seq_num.cmp(&other.seq_num)
        } else if Side::is_price_better_than(self.price, other.price) {
            Ordering::Less
        } else {
            Ordering::Greater
        }
    }
}

impl<Side: OrderbookSide> PartialOrd for OrderKey<Side> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum Side {
    Bid,
    Ask,
}

impl Display for Side {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        match *self {
            Side::Bid => write!(f, "Bid"),
            Side::Ask => write!(f, "Ask"),
        }
    }
}

#[derive(Hash, PartialEq, Eq, Debug, Clone)]
pub struct Symbol(String);

impl Display for Symbol {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        self.0.fmt(f)
    }
}

#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
pub struct OrderId(u32);

#[derive(Debug, Clone)]
pub struct Order {
    pub seq_num: u64,
    pub id: OrderId,
    pub side: Side,
    pub volume: i32,
    pub price: f64,
}

impl Display for Order {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        write!(
            f,
            "[{id} #{seq_num}] {side} {volume:3} @ {price:.02}",
            id = self.id.0,
            seq_num = self.seq_num,
            side = self.side,
            price = self.price,
            volume = self.volume,
        )
    }
}

pub struct OrderHandle {
    remove_fn: Box<dyn FnOnce()>,
}

#[derive(Default)]
struct Orderbook {
    bids: BTreeMap<OrderKey<Bid>, Order>,
    asks: BTreeMap<OrderKey<Ask>, Order>,
}

impl Orderbook {
    fn new() -> Orderbook {
        Orderbook {
            bids: BTreeMap::new(),
            asks: BTreeMap::new(),
        }
    }

    unsafe fn insert<Side: OrderbookSide>(
        this: *mut Self,
        order: Order,
    ) -> Option<OrderHandle> {
        let order_key = OrderKey::<Side> {
            price: order.price,
            seq_num: order.seq_num,
            _side: PhantomData,
        };

        // ...

        if order.volume <= 0 {
            return None;
        }

        let same_side = Side::get_orders(this);
        let old_order = (*same_side).insert(order_key, order);
        assert!(old_order.is_none());

        let order_handle = {
            let orderbook_ptr = this;
            OrderHandle {
                remove_fn: Box::new(move || unsafe {
                    let same_side = Side::get_orders(orderbook_ptr);
                    (*same_side).remove_entry(&order_key);
                }),
            }
        };
        Some(order_handle)
    }

    fn print(&self) {
        for (_, ask) in self.asks.iter().rev() {
            println!("{ask}");
        }
        for (_, bid) in self.bids.iter() {
            println!("{bid}");
        }
    }
}

#[derive(Default)]
pub struct Exchange {
    order_books: HashMap<Symbol, *mut Orderbook>,
    order_handles: HashMap<OrderId, OrderHandle>,
}

impl Exchange {
    pub fn new() -> Exchange {
        Exchange {
            order_books: HashMap::new(),
            order_handles: HashMap::new(),
        }
    }

    pub fn insert(&mut self, symbol: Symbol, order: Order) {
        let order_id = order.id;
        let orderbook = *self
            .order_books
            .entry(symbol)
            .or_insert_with(|| Box::into_raw(Box::new(Orderbook::new())));

        let new_handle = unsafe {
            match order.side {
                Side::Bid => Orderbook::insert::<Bid>(orderbook, order),
                Side::Ask => Orderbook::insert::<Ask>(orderbook, order),
            }
        };

        if let Some(new_handle) = new_handle {
            self.order_handles.insert(order_id, new_handle);
        }
    }

    pub fn cancel(&mut self, order_id: OrderId) {
        if let Some(order_handle) = self.order_handles.remove(&order_id) {
            (order_handle.remove_fn)();
        }
    }

    pub fn print(&self, symbol: Symbol) {
        match self.order_books.get(&symbol) {
            Some(orderbook) => {
                println!("{symbol}");
                unsafe {
                    (**orderbook).print();
                }
            }
            None => println!("{symbol}: no such orderbook"),
        };
    }

    pub fn print_all(&self) {
        for (symbol, orderbook) in self.order_books.iter() {
            println!("{symbol}");
            unsafe { (**orderbook).print() };
        }
    }
}

impl Drop for Exchange {
    fn drop(&mut self) {
        for (_, orderbook) in self.order_books.iter() {
            drop(unsafe { Box::from_raw(*orderbook) });
        }
    }
}

fn main() {
    let mut exchange = Exchange::new();

    exchange.insert(
        Symbol("AAPL".to_owned()),
        Order {
            seq_num: 1,
            id: OrderId(1),
            side: Side::Bid,
            volume: 2,
            price: 12.50,
        },
    );
    exchange.print_all();

    exchange.insert(
        Symbol("AAPL".to_owned()),
        Order {
            seq_num: 2,
            id: OrderId(2),
            side: Side::Bid,
            volume: 10,
            price: 13.00,
        },
    );
    exchange.print_all();

    exchange.cancel(OrderId(1));
    exchange.print_all();

    exchange.cancel(OrderId(2));
    exchange.print_all();

    for i in 3..11u32 {
        exchange.insert(
            Symbol("AAPL".to_owned()),
            Order {
                seq_num: i as u64,
                id: OrderId(i),
                side: Side::Bid,
                volume: 2i32 + i as i32,
                price: 12.50 + (i as f64) / 100.0,
            },
        );
    }
    exchange.print_all();

    exchange.cancel(OrderId(3));
    exchange.print_all();
}