update order dlob

src/dlob/dlob.rs

@@ -1,5 +1,5 @@
 use dashmap::{DashMap, DashSet};
-use drift::state::user::{Order, OrderStatus};
+use drift::state::user::{MarketType, Order, OrderStatus};
 use solana_sdk::pubkey::Pubkey;
 
 use crate::dlob::node_list::NodeList;
@@ -30,7 +30,7 @@ impl DLOB {
         DLOB {
             exchange,
             open_orders,
-            initialized: false,
+            initialized: true,
             max_slot_for_resting_limit_orders
         }
     }
@@ -82,12 +82,92 @@ impl DLOB {
 
         None
     }
+
+    pub fn delete_order(&self, order_id: u32, user_account: Pubkey) {
+        let order_signature = get_order_signature(order_id, user_account);
+        for mut node_list in get_node_lists(&self.exchange) {
+            if let Some(_) = node_list.get_node(&order_signature) {
+                node_list.remove(&order_signature);
+            }
+        }
+    }
+
+    fn update_resting_limit_orders_for_market_type(&self, market_type: &MarketType, slot: u64) {
+        let mut ask_order_sigs_to_update: Vec<String> = vec![];
+        let mut bid_order_sigs_to_update: Vec<String> = vec![];
+        let market_type = market_type_to_string(market_type);
+        if let Some(market) = self.exchange.get_mut(&market_type) {
+            for mut market_ref in market.iter_mut() {
+                let market = market_ref.value_mut();
+                for node in market.taking_limit_orders.taking_limit_bids.iter() {
+                    if !node.get_order().is_resting_limit_order(slot).unwrap() {
+                        continue;
+                    }
+                    bid_order_sigs_to_update.push(get_order_signature(node.get_order().order_id, node.get_user_account()));
+                }
+                for node in market.taking_limit_orders.taking_limit_asks.iter() {
+                    if !node.get_order().is_resting_limit_order(slot).unwrap() {
+                        continue;
+                    }
+                    ask_order_sigs_to_update.push(get_order_signature(node.get_order().order_id, node.get_user_account()));
+                }
+
+                for order_sig in ask_order_sigs_to_update.iter() {
+                    if let Some(mut node) = market.taking_limit_orders.taking_limit_asks.remove(order_sig) {
+                        node.set_node_type(NodeType::RestingLimit);
+                        market.resting_limit_orders.resting_limit_asks.insert(node);
+                    }
+                }
+
+                for order_sig in bid_order_sigs_to_update.iter() {
+                    if let Some(mut node) = market.taking_limit_orders.taking_limit_bids.remove(order_sig) {
+                        node.set_node_type(NodeType::RestingLimit);
+                        market.resting_limit_orders.resting_limit_bids.insert(node);
+                    }
+                }
+            }
+        }
+    }
+    pub fn update_resting_limit_orders(&mut self, slot: u64) {
+        if slot <= self.max_slot_for_resting_limit_orders {
+            return
+        }
+
+        self.max_slot_for_resting_limit_orders = slot;
+
+        self.update_resting_limit_orders_for_market_type(&MarketType::Perp, slot);
+        self.update_resting_limit_orders_for_market_type(&MarketType::Spot, slot);
+
+    }
+
+    pub fn update_order(&mut self, order: Order, user_account: Pubkey, slot: u64, cumulative_base_asset_amount_filled: u64) {
+        self.update_resting_limit_orders(slot);
+
+        if order.base_asset_amount == cumulative_base_asset_amount_filled {
+            self.delete_order(order.order_id, user_account);
+            return
+        }
+
+        if order.base_asset_amount_filled == cumulative_base_asset_amount_filled {
+            return
+        }
+
+        let mut new_order = order.clone();
+
+        new_order.base_asset_amount_filled = cumulative_base_asset_amount_filled;
+
+        if let Some(mut market) = self.exchange.get_mut(&market_type_to_string(&order.market_type)).unwrap().get_mut(&order.market_index) {
+            if let Some(node_list) = market.get_list_for_order(&order, slot) {
+                node_list.update_order(&get_order_signature(order.order_id, user_account), new_order);
+            }
+        }
+    }
 }
 
 
 #[cfg(test)]
 mod tests {
-    use drift::state::user::{MarketType, OrderType};
+    use drift::{controller::position::PositionDirection, state::user::{MarketType, OrderType}};
 
     use super::*;
 
@@ -108,5 +188,146 @@ mod tests {
 
 
         assert!(dlob.get_order(order.order_id, user_account).is_some());
+        assert!(dlob.get_order(2, user_account).is_none());
+    }
+
+    #[test]
+    fn test_delete_order() {
+        let dlob = DLOB::new(0);
+        let order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            ..Order::default()
+        };
+        let user_account = Pubkey::new_unique();
+        let slot = 0;
+        dlob.insert_order(order, user_account, slot);
+
+        assert!(dlob.get_order(order.order_id, user_account).is_some());
+
+        dlob.delete_order(order.order_id, user_account);
+
+        assert!(dlob.get_order(order.order_id, user_account).is_none());
+    }
+
+    #[test]
+    fn test_update_order_cumulative_baa_filled_below_order_baa_filled() {
+
+        let mut dlob = DLOB::new(0);
+        let order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 0,
+            ..Order::default()
+        };
+        let user_account = Pubkey::new_unique();
+        let slot = 0;
+        dlob.insert_order(order, user_account, slot);
+
+        assert!(dlob.get_order(order.order_id, user_account).unwrap().slot == 0);
+
+        let updated_order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 1,
+            base_asset_amount_filled: 1_000_000,
+            ..Order::default()
+        };
+        dlob.update_order(updated_order, user_account, slot, 500_000);
+
+        assert!(dlob.get_order(order.order_id, user_account).unwrap().slot == 1);
+    }
+
+    #[test]
+    fn test_update_order_cumulative_baa_equal_order_baa_filled() {    
+        let mut dlob = DLOB::new(0);
+        let order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 0,
+            ..Order::default()
+        };
+        let user_account = Pubkey::new_unique();
+        let slot = 0;
+        dlob.insert_order(order, user_account, slot);
+
+        assert!(dlob.get_order(order.order_id, user_account).unwrap().slot == 0);
+
+        let updated_order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 1,
+            base_asset_amount: 1_000_000,
+            ..Order::default()
+        };
+        dlob.update_order(updated_order, user_account, slot, 1_000_000);
+
+        assert!(dlob.get_order(order.order_id, user_account).is_none());
+    }
+
+    #[test]
+    fn test_update_resting_limit_nodes() {
+        let mut dlob = DLOB::new(0);
+        let order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 0,
+            ..Order::default()
+        };
+        let order_2 = Order {
+            order_id: 2,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 0,
+            auction_duration: 100,
+            ..Order::default()
+        };
+        let user_account = Pubkey::new_unique();
+        let slot = 1;
+        dlob.insert_order(order, user_account, slot);
+        dlob.insert_order(order_2, user_account, slot);
+
+        let updated_order = Order {
+            order_id: 1,
+            market_type: MarketType::Spot,
+            market_index: 1,
+            order_type: OrderType::Limit,
+            status: OrderStatus::Open,
+            slot: 1,
+            base_asset_amount: 1_000_000,
+            direction: PositionDirection::Long,
+            ..Order::default()
+        };
+
+        let update_slot = 105;
+        dlob.update_order(updated_order, user_account, update_slot, 1_000_000);
+
+        assert!(dlob.get_order(order.order_id, user_account).is_none());
+
+        let side = &dlob.exchange.get("spot").unwrap();
+        let node_list = &side.get(&1).unwrap().resting_limit_orders.resting_limit_bids;
+        let order_sig = get_order_signature(order_2.order_id, user_account);
+
+        assert!(node_list.contains(&order_sig));
     }
 }

src/dlob/dlob_node.rs

@@ -121,6 +121,14 @@ impl DLOBNode for Node {
             Node::VAMMNode(_) => NodeType::VAMM,
         }
     }
+
+    fn set_node_type(&mut self, node_type: NodeType) {
+        match self {
+            Node::OrderNode(order_node) => order_node.set_node_type(node_type),
+            Node::VAMMNode(_) => unimplemented!(),
+        }
+
+    }
 }
 
 #[derive(Clone, Copy)]
@@ -202,6 +210,10 @@ impl DLOBNode for OrderNode {
     fn get_node_type(&self) -> NodeType {
         self.node_type
     }
+
+    fn set_node_type(&mut self, node_type: NodeType) {
+        self.node_type = node_type;
+    }
 }
 
 pub(crate) fn create_node<'a>(arena: &'a Arena<Node>, node_type: NodeType, order: Order, user_account: Pubkey) -> &'a mut Node {
@@ -236,6 +248,8 @@ pub(crate) trait DLOBNode {
     fn set_order(&mut self, order: Order);
 
     fn get_node_type(&self) -> NodeType;
+
+    fn set_node_type(&mut self, node_type: NodeType);
 }
 
 pub(crate) trait DLOBNodePointerExt {
@@ -302,6 +316,10 @@ impl DLOBNode for VAMMNode {
     fn get_node_type(&self) -> NodeType {
         NodeType::VAMM
     }
+
+    fn set_node_type(&mut self, node_type: NodeType) {
+        unimplemented!()
+    }
 }
 
 impl VAMMNode {

src/dlob/node_list.rs

@@ -3,6 +3,7 @@ use crate::dlob::dlob_node::{DLOBNode, SortDirection, NodeType, get_order_signat
 use dashmap::DashMap;
 use drift::state::user::Order;
 use std::sync::Arc;
+use std::panic::Location;
 
 #[derive(Clone)]
 pub(crate) struct NodeList {
@@ -24,9 +25,11 @@ impl NodeList {
         }
     }
 
+    #[track_caller]
     pub(crate) fn insert(&mut self, node: Node) {
         if node.get_node_type() != self.node_type {
-            panic!("{}", format!("Node type mismatch. Expected: {:?}, Got: {:?}", self.node_type, node.get_node_type()));
+            let caller = Location::caller();
+            panic!("{}", format!("Node type mismatch. Expected: {:?}, Got: {:?} at {:?}", self.node_type, node.get_node_type(), caller));
         }
         let sort_value = node.get_sort_value(node.get_order()).unwrap();
         let order_signature = get_order_signature(node.get_order().order_id, node.get_user_account());
@@ -98,13 +101,13 @@ impl NodeList {
         }
     }
 
-    pub(crate) fn remove(&mut self, order_signature: &str) -> Option<Box<dyn DLOBNode>> {
+    pub(crate) fn remove(&mut self, order_signature: &str) -> Option< Node> {
         if let Some(node_ptr) = self.node_map.remove(order_signature) {
             let node_ptr = node_ptr.1;
             unsafe {
                 let prev_ptr = (&*node_ptr).get_prev_ptr();
                 let next_ptr = (&*node_ptr).get_next_ptr();
-
+    
                 if let Some(prev_ptr) = prev_ptr {
                     (&mut *prev_ptr).set_next(next_ptr);
                 } else {
@@ -114,9 +117,11 @@ impl NodeList {
                 if let Some(next_ptr) = next_ptr {
                     (&mut *next_ptr).set_prev(prev_ptr);
                 }
+
+                let node = *Box::from_raw(node_ptr as *mut Node);
+                self.length -= 1;
+                Some(node)
             }
-            self.length -= 1;
-            Some(unsafe { Box::from_raw(node_ptr) })
         } else {
             None
         }