diff --git a/contracts/sumtree-orderbook/Cargo.toml b/contracts/sumtree-orderbook/Cargo.toml index 1ea3ab7..cfe0bc9 100644 --- a/contracts/sumtree-orderbook/Cargo.toml +++ b/contracts/sumtree-orderbook/Cargo.toml @@ -54,3 +54,4 @@ thiserror = { version = "1.0.49" } [dev-dependencies] cw-multi-test = "0.18.0" +rand = "0.8.4" diff --git a/contracts/sumtree-orderbook/src/order.rs b/contracts/sumtree-orderbook/src/order.rs index 426ac2e..ffd4418 100644 --- a/contracts/sumtree-orderbook/src/order.rs +++ b/contracts/sumtree-orderbook/src/order.rs @@ -4,7 +4,7 @@ use crate::constants::{MAX_TICK, MIN_TICK}; use crate::error::ContractError; use crate::state::{new_order_id, orders, ORDERBOOKS, TICK_STATE}; use crate::sumtree::node::{generate_node_id, NodeType, TreeNode}; -use crate::sumtree::tree::TREE; +use crate::sumtree::tree::{get_root_node, TREE}; use crate::types::{LimitOrder, OrderDirection, REPLY_ID_REFUND}; use cosmwasm_std::{ coin, ensure, ensure_eq, BankMsg, Decimal256, DepsMut, Env, MessageInfo, Response, SubMsg, @@ -152,14 +152,18 @@ pub fn cancel_limit( ); // Fetch the sumtree from storage, or create one if it does not exist - let mut tree = TREE - .load(deps.storage, &(order.book_id, order.tick_id)) - .unwrap_or(TreeNode::new( + let mut tree = if let Ok(tree) = get_root_node(deps.storage, book_id, tick_id) { + tree + } else { + let new_root = TreeNode::new( order.book_id, order.tick_id, - generate_node_id(deps.storage, order.book_id, order.tick_id)?, + generate_node_id(deps.storage, book_id, tick_id)?, NodeType::default(), - )); + ); + TREE.save(deps.storage, &(book_id, tick_id), &new_root.key)?; + new_root + }; // Generate info for new node to insert to sumtree let node_id = generate_node_id(deps.storage, order.book_id, order.tick_id)?; @@ -216,7 +220,6 @@ pub fn cancel_limit( )?; tree.save(deps.storage)?; - TREE.save(deps.storage, &(book_id, tick_id), &tree)?; Ok(Response::new() .add_attribute("method", "cancelLimit") diff --git a/contracts/sumtree-orderbook/src/sumtree/node.rs b/contracts/sumtree-orderbook/src/sumtree/node.rs index 5ccccb5..af645d4 100644 --- a/contracts/sumtree-orderbook/src/sumtree/node.rs +++ b/contracts/sumtree-orderbook/src/sumtree/node.rs @@ -9,7 +9,7 @@ use cosmwasm_schema::cw_serde; use cosmwasm_std::{ensure, Storage, Uint128}; use cw_storage_plus::Map; -use crate::{error::ContractResult, ContractError}; +use crate::{error::ContractResult, sumtree::tree::TREE, ContractError}; pub const NODES: Map<&(u64, i64, u64), TreeNode> = Map::new("nodes"); pub const NODE_ID_COUNTER: Map<&(u64, i64), u64> = Map::new("node_id"); @@ -40,6 +40,8 @@ pub enum NodeType { accumulator: Uint128, // Range from min ETAS to max ETAS + value of max ETAS range: (Uint128, Uint128), + // Depth of subtree below node + weight: u64, }, } @@ -58,6 +60,7 @@ impl NodeType { Self::Internal { range: (range.0.into(), range.1.into()), accumulator: accumulator.into(), + weight: 0, } } } @@ -67,6 +70,7 @@ impl Default for NodeType { Self::Internal { accumulator: Uint128::zero(), range: (Uint128::MAX, Uint128::MIN), + weight: 0, } } } @@ -134,6 +138,12 @@ impl TreeNode { Ok(NODES.save(storage, &(self.book_id, self.tick_id, self.key), self)?) } + /// Resyncs a node with values stored in CosmWasm Storage + pub fn sync(&mut self, storage: &dyn Storage) -> ContractResult<()> { + *self = NODES.load(storage, &(self.book_id, self.tick_id, self.key))?; + Ok(()) + } + /// Returns the maximum range value of a node. /// /// For `Internal` nodes, this is the maximum value of the associated range. @@ -186,6 +196,13 @@ impl TreeNode { } } + pub fn get_weight(&self) -> u64 { + match self.node_type { + NodeType::Internal { weight, .. } => weight, + NodeType::Leaf { .. } => 1, + } + } + /// Adds a given value to an internal node's accumulator /// /// Errors if given node is not internal @@ -193,6 +210,37 @@ impl TreeNode { self.set_value(self.get_value().checked_add(value)?) } + pub fn set_weight(&mut self, new_weight: u64) -> ContractResult<()> { + match &mut self.node_type { + NodeType::Internal { weight, .. } => { + *weight = new_weight; + Ok(()) + } + NodeType::Leaf { .. } => Err(ContractError::InvalidNodeType), + } + } + + /// Gets the value for a given node. + /// + /// For `Leaf` nodes this is the `value`. + /// + /// For `Internal` nodes this is the `accumulator`. + pub fn get_value(&self) -> Uint128 { + match self.node_type { + NodeType::Leaf { value, .. } => value, + NodeType::Internal { accumulator, .. } => accumulator, + } + } + + /// Synchronizes the range and value of the current node and recursively updates its ancestors. + pub fn sync_range_and_value_up(&mut self, storage: &mut dyn Storage) -> ContractResult<()> { + self.sync_range_and_value(storage)?; + if let Some(mut parent) = self.get_parent(storage)? { + parent.sync_range_and_value_up(storage)?; + } + Ok(()) + } + /// Recalculates the range and accumulated value for a node and propagates it up the tree /// /// Must be an internal node @@ -205,9 +253,10 @@ impl TreeNode { let right_exists = maybe_right.is_some(); if !self.has_child() { - return Err(ContractError::ChildlessInternalNode); + return Ok(()); } + // Calculate new range let (min, max) = if left_exists && !right_exists { let left = maybe_left.clone().unwrap(); (left.get_min_range(), left.get_max_range()) @@ -223,38 +272,35 @@ impl TreeNode { left.get_max_range().max(right.get_max_range()), ) }; - self.set_min_range(min)?; self.set_max_range(max)?; + // Calculate new value let value = maybe_left + .clone() .map(|n| n.get_value()) .unwrap_or_default() - .checked_add(maybe_right.map(|n| n.get_value()).unwrap_or_default())?; + .checked_add( + maybe_right + .clone() + .map(|n| n.get_value()) + .unwrap_or_default(), + )?; self.set_value(value)?; + // Calculate new weight + let weight = maybe_left + .map(|n| n.get_weight()) + .unwrap_or_default() + .max(maybe_right.map(|n| n.get_weight()).unwrap_or_default()); + self.set_weight(weight + 1)?; + // Must save before propagating as parent will read this node self.save(storage)?; - if let Some(mut parent) = self.get_parent(storage)? { - parent.sync_range_and_value(storage)?; - } - Ok(()) } - /// Gets the value for a given node. - /// - /// For `Leaf` nodes this is the `value`. - /// - /// For `Internal` nodes this is the `accumulator`. - pub fn get_value(&self) -> Uint128 { - match self.node_type { - NodeType::Leaf { value, .. } => value, - NodeType::Internal { accumulator, .. } => accumulator, - } - } - /// Inserts a given node in to the tree /// /// If the node is internal an error is returned. @@ -286,28 +332,40 @@ impl TreeNode { let maybe_left = self.get_left(storage)?; let maybe_right = self.get_right(storage)?; + let is_left_internal = maybe_left.clone().map_or(false, |l| l.is_internal()); + let is_right_internal = maybe_right.clone().map_or(false, |r| r.is_internal()); let is_in_left_range = maybe_left.clone().map_or(false, |left| { - left.is_internal() && new_node.get_min_range() < left.get_max_range() + new_node.get_min_range() <= left.get_max_range() }); let is_in_right_range = maybe_right.clone().map_or(false, |right| { - right.is_internal() && new_node.get_min_range() > right.get_min_range() + new_node.get_min_range() >= right.get_min_range() }); // Case 1 Left - if is_in_left_range { + if is_left_internal && is_in_left_range { + self.save(storage)?; // Can unwrap as node must exist let mut left = maybe_left.unwrap(); - left.insert(storage, new_node)?; - self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } + // Case 1 Right - if is_in_right_range { + if is_right_internal && is_in_right_range { + self.save(storage)?; // Can unwrap as node must exist let mut right = maybe_right.unwrap(); right.insert(storage, new_node)?; + self.rebalance(storage)?; + return Ok(()); + } + + if is_right_internal && is_left_internal { self.save(storage)?; + let mut left = maybe_left.unwrap(); + left.insert(storage, new_node)?; + self.rebalance(storage)?; return Ok(()); } @@ -317,6 +375,7 @@ impl TreeNode { new_node.parent = Some(self.key); new_node.save(storage)?; self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } @@ -330,6 +389,7 @@ impl TreeNode { new_node.parent = Some(self.key); new_node.save(storage)?; self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } @@ -339,6 +399,7 @@ impl TreeNode { new_node.parent = Some(self.key); new_node.save(storage)?; self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } @@ -346,27 +407,29 @@ impl TreeNode { let right_is_leaf = maybe_right .clone() .map_or(false, |right| !right.is_internal()); + let is_higher_than_right_leaf = maybe_right.clone().map_or(false, |r| { + !r.is_internal() && new_node.get_min_range() >= r.get_max_range() + }); // Case 4 - if left_is_leaf { + if left_is_leaf && !is_higher_than_right_leaf { let mut left = maybe_left.unwrap(); let new_left = left.split(storage, new_node)?; self.left = Some(new_left); self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } // Case 5: Reordering // TODO: Add edge case test for this - let is_higher_than_right_leaf = maybe_right.clone().map_or(false, |r| { - !r.is_internal() && new_node.get_min_range() >= r.get_max_range() - }); if is_higher_than_right_leaf && maybe_left.is_none() { self.left = self.right; self.right = Some(new_node.key); new_node.parent = Some(self.key); new_node.save(storage)?; self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } @@ -376,6 +439,7 @@ impl TreeNode { let new_right = right.split(storage, new_node)?; self.right = Some(new_right); self.save(storage)?; + self.rebalance(storage)?; return Ok(()); } @@ -413,6 +477,7 @@ impl TreeNode { id, NodeType::internal(accumulator, (new_min, new_max)), ); + new_parent.set_weight(2)?; // Save new key references new_parent.parent = self.parent; @@ -420,7 +485,6 @@ impl TreeNode { new_parent.right = Some(new_right); self.parent = Some(id); new_node.parent = Some(id); - new_parent.save(storage)?; self.save(storage)?; new_node.save(storage)?; @@ -446,7 +510,7 @@ impl TreeNode { parent.delete(storage)?; } else { // Update parents values after removing node - parent.sync_range_and_value(storage)?; + parent.sync_range_and_value_up(storage)?; } } @@ -455,6 +519,217 @@ impl TreeNode { Ok(()) } + /// Returns the balance factor of a node: `(left weight - right weight)` + /// + /// Empty nodes return a weight of 0, so a childless node/leaf will return a balance factor of 0 + pub fn get_balance_factor(&self, storage: &dyn Storage) -> ContractResult { + let left_weight = self.get_left(storage)?.map_or(0, |n| n.get_weight()); + let right_weight = self.get_right(storage)?.map_or(0, |n| n.get_weight()); + Ok(right_weight as i32 - left_weight as i32) + } + + /// Rebalances the tree starting from the current node. + /// + /// This method ensures that the AVL tree properties are maintained after insertions or deletions + /// have been performed. It checks the balance factor of the current node and performs rotations + /// as necessary to bring the tree back into balance. + pub fn rebalance(&mut self, storage: &mut dyn Storage) -> ContractResult<()> { + ensure!(self.is_internal(), ContractError::InvalidNodeType); + ensure!(self.has_child(), ContractError::ChildlessInternalNode); + + // Synchronize the current node's state with storage before rebalancing. + self.sync(storage)?; + + // Calculate the balance factor to determine if rebalancing is needed. + let balance_factor = self.get_balance_factor(storage)?; + // Early return if the tree is already balanced. + if balance_factor.abs() <= 1 { + self.sync_range_and_value(storage)?; + return Ok(()); + } + + // Retrieve optional references to left and right children. + let maybe_left = self.get_left(storage)?; + let maybe_right = self.get_right(storage)?; + + // Determine the direction of imbalance. + let is_right_leaning = balance_factor > 0; + let is_left_leaning = balance_factor < 0; + + // Calculate balance factors for child nodes to determine rotation type. + let right_balance_factor = maybe_right + .as_ref() + .map_or(0, |n| n.get_balance_factor(storage).unwrap_or(0)); + let left_balance_factor = maybe_left + .as_ref() + .map_or(0, |n| n.get_balance_factor(storage).unwrap_or(0)); + + // Perform rotations based on the type of imbalance detected. + // Case 1: Right-Right (Right rotation needed) + if is_right_leaning && right_balance_factor >= 0 { + self.rotate_left(storage)?; + } + // Case 2: Left-Left (Right rotation needed) + else if is_left_leaning && left_balance_factor <= 0 { + self.rotate_right(storage)?; + } + // Case 3: Right-Left (Right rotation on right child followed by Left rotation on self) + else if is_right_leaning && right_balance_factor < 0 { + maybe_right.unwrap().rotate_right(storage)?; + self.sync(storage)?; + self.rotate_left(storage)?; + } + // Case 4: Left-Right (Left rotation on left child followed by Right rotation on self) + else if is_left_leaning && left_balance_factor > 0 { + maybe_left.unwrap().rotate_left(storage)?; + self.sync(storage)?; + self.rotate_right(storage)?; + } + + Ok(()) + } + + /// Performs a right rotation on the current node. **Called by the root of the subtree to be rotated.** + /// + /// This operation is used to rebalance the tree when the left subtree + /// has a greater height than the right subtree. It adjusts the pointers + /// accordingly to ensure the tree remains a valid binary search tree. + pub fn rotate_right(&mut self, storage: &mut dyn Storage) -> ContractResult<()> { + // Retrieve the parent node, if any. + let maybe_parent = self.get_parent(storage)?; + // Determine if the current node is a left or right child of its parent. + let is_left_child = maybe_parent + .clone() + .map_or(false, |p| p.left == Some(self.key)); + let is_right_child = maybe_parent + .clone() + .map_or(false, |p| p.right == Some(self.key)); + + // Ensure the current node has a left child to rotate. + let maybe_left = self.get_left(storage)?; + ensure!(maybe_left.is_some(), ContractError::InvalidNodeType); + + // Perform the rotation. + let mut left = maybe_left.unwrap(); + left.parent = self.parent; + self.parent = Some(left.key); + self.left = left.right; + + // Update the parent of the new left child, if it exists. + if let Some(mut new_left) = self.get_left(storage)? { + new_left.parent = Some(self.key); + new_left.save(storage)?; + } + + // Complete the rotation by setting the right child of the left node to the current node. + left.right = Some(self.key); + // Save the changes to both nodes. + left.save(storage)?; + self.save(storage)?; + + // Synchronize the range and value of the current node. + self.sync_range_and_value(storage)?; + left.sync_range_and_value(storage)?; + + // If the left node has no parent, it becomes the new root. + if left.parent.is_none() { + TREE.save(storage, &(left.book_id, left.tick_id), &left.key)?; + } + + // Update the parent's child pointers. + if is_left_child { + let mut parent = maybe_parent.clone().unwrap(); + parent.left = Some(left.key); + parent.save(storage)?; + } + if is_right_child { + let mut parent = maybe_parent.unwrap(); + parent.right = Some(left.key); + parent.save(storage)?; + } + + Ok(()) + } + + /// Performs a left rotation on the current node within the binary tree. **Called by the root of the subtree to be rotated.** + /// + /// This operation is used to rebalance the tree when the right subtree + /// has a greater height than the left subtree. It adjusts the pointers + /// accordingly to ensure the tree remains a valid binary search tree. + pub fn rotate_left(&mut self, storage: &mut dyn Storage) -> ContractResult<()> { + // Retrieve the parent node, if any, to determine the current node's relationship. + let maybe_parent = self.get_parent(storage)?; + let is_left_child = maybe_parent + .clone() + .map_or(false, |p| p.left == Some(self.key)); + let is_right_child = maybe_parent + .clone() + .map_or(false, |p| p.right == Some(self.key)); + + // Ensure the current node has a right child to perform the rotation. + let maybe_right = self.get_right(storage)?; + ensure!(maybe_right.is_some(), ContractError::InvalidNodeType); + + // Perform the rotation by reassigning parent and child references. + let mut right = maybe_right.unwrap(); + right.parent = self.parent; + self.parent = Some(right.key); + self.right = right.left; + + // Update the parent reference of the new right child, if it exists. + if let Some(mut new_right) = self.get_left(storage)? { + new_right.parent = Some(self.key); + new_right.save(storage)?; + } + + // Complete the rotation by setting the left child of the right node to the current node. + right.left = Some(self.key); + + // Persist the changes to both nodes. + right.save(storage)?; + self.save(storage)?; + + // Synchronize the range and value of the current node. + self.sync_range_and_value(storage)?; + right.sync_range_and_value(storage)?; + + // If the right node has no parent after the rotation, it becomes the new root. + if right.parent.is_none() { + TREE.save(storage, &(right.book_id, right.tick_id), &right.key)?; + } + + // Update the child references of the parent node to reflect the rotation. + if is_left_child { + let mut parent = maybe_parent.clone().unwrap(); + parent.left = Some(right.key); + parent.save(storage)?; + } + if is_right_child { + let mut parent = maybe_parent.unwrap(); + parent.right = Some(right.key); + parent.save(storage)?; + } + + Ok(()) + } + + #[cfg(test)] + pub fn with_children(self, left: Option, right: Option) -> Self { + Self { + left, + right, + ..self + } + } + + #[cfg(test)] + pub fn with_parent(self, parent: u64) -> Self { + Self { + parent: Some(parent), + ..self + } + } + #[cfg(test)] /// Depth first search traversal of tree pub fn traverse(&self, storage: &dyn Storage) -> ContractResult> { @@ -473,7 +748,7 @@ impl TreeNode { } #[cfg(test)] - pub fn get_height(&self, storage: &dyn Storage) -> ContractResult { + pub fn get_height(&self, storage: &dyn Storage) -> ContractResult { let mut height = 0; if let Some(left) = self.get_left(storage)? { height = height.max(left.get_height(storage)?); @@ -515,7 +790,9 @@ impl Display for NodeType { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { NodeType::Leaf { value, etas } => write!(f, "{etas} {value}"), - NodeType::Internal { accumulator, range } => { + NodeType::Internal { + accumulator, range, .. + } => { write!(f, "{} {}-{}", accumulator, range.0, range.1) } } diff --git a/contracts/sumtree-orderbook/src/sumtree/test/test_node.rs b/contracts/sumtree-orderbook/src/sumtree/test/test_node.rs index a5a09c0..498221f 100644 --- a/contracts/sumtree-orderbook/src/sumtree/test/test_node.rs +++ b/contracts/sumtree-orderbook/src/sumtree/test/test_node.rs @@ -1,6 +1,12 @@ -use cosmwasm_std::{testing::mock_dependencies, Uint128}; +use cosmwasm_std::{testing::mock_dependencies, Deps, Storage, Uint128}; -use crate::sumtree::node::{generate_node_id, NodeType, TreeNode, NODES}; +use crate::{ + sumtree::{ + node::{generate_node_id, NodeType, TreeNode, NODES}, + tree::{get_root_node, TREE}, + }, + ContractError, +}; struct TestNodeInsertCase { name: &'static str, @@ -11,6 +17,70 @@ struct TestNodeInsertCase { print: bool, } +// Asserts all values of internal nodes are as expected +fn assert_internal_values( + test_name: &'static str, + deps: Deps, + internals: Vec<&TreeNode>, + should_be_balanced: bool, +) { + for internal_node in internals { + let left_node = internal_node.get_left(deps.storage).unwrap(); + let right_node = internal_node.get_right(deps.storage).unwrap(); + + let accumulated_value = left_node + .clone() + .map_or(Uint128::zero(), |x| x.get_value()) + .checked_add( + right_node + .clone() + .map_or(Uint128::zero(), |x| x.get_value()), + ) + .unwrap(); + assert_eq!(internal_node.get_value(), accumulated_value); + + let min = left_node + .clone() + .map_or(Uint128::MAX, |n| n.get_min_range()) + .min( + right_node + .clone() + .map_or(Uint128::MAX, |n| n.get_min_range()), + ); + let max = left_node + .clone() + .map_or(Uint128::MIN, |n| n.get_max_range()) + .max( + right_node + .clone() + .map_or(Uint128::MIN, |n| n.get_max_range()), + ); + assert_eq!(internal_node.get_min_range(), min); + assert_eq!(internal_node.get_max_range(), max); + + let balance_factor = right_node + .map_or(0, |n| n.get_height(deps.storage).unwrap()) + .abs_diff(left_node.map_or(0, |n| n.get_height(deps.storage).unwrap())); + + assert_eq!( + internal_node.get_weight(), + internal_node.get_height(deps.storage).unwrap(), + "{}: Internal weight incorrect for node {}", + test_name, + internal_node.key + ); + + if should_be_balanced { + assert!( + balance_factor <= 1, + "{}: Balance factor greater than 1 for node {}", + test_name, + internal_node.key + ); + } + } +} + #[test] fn test_node_insert_valid() { let book_id = 1; @@ -203,26 +273,12 @@ fn test_node_insert_valid() { // Print tree at second last node to see pre-insert if test.nodes.len() >= 2 && idx == test.nodes.len() - 2 && test.print { - println!("Pre Insert Tree: {}", test.name); - println!("--------------------------"); - - let nodes = tree.traverse_bfs(deps.as_ref().storage).unwrap(); - for (idx, row) in nodes.iter().enumerate() { - print_tree_row(row.clone(), idx == 0, (nodes.len() - idx - 1) as u32); - } - println!(); + print_tree("Pre-Insert Tree", test.name, &tree, &deps.as_ref()); } } if test.print { - println!("Post Insert Tree: {}", test.name); - println!("--------------------------"); - - let nodes = tree.traverse_bfs(deps.as_ref().storage).unwrap(); - for (idx, row) in nodes.iter().enumerate() { - print_tree_row(row.clone(), idx == 0, (nodes.len() - idx - 1) as u32); - } - println!(); + print_tree("Post-Insert Tree", test.name, &tree, &deps.as_ref()); } // Return tree in vector form from Depth First Search @@ -240,116 +296,8 @@ fn test_node_insert_valid() { // Ensure all internal nodes are correctly summed and contain correct ranges let internals: Vec<&TreeNode> = result.iter().filter(|x| x.is_internal()).collect(); - for internal_node in internals { - let left_node = internal_node.get_left(deps.as_ref().storage).unwrap(); - let right_node = internal_node.get_right(deps.as_ref().storage).unwrap(); - - let accumulated_value = left_node - .clone() - .map_or(Uint128::zero(), |x| x.get_value()) - .checked_add( - right_node - .clone() - .map_or(Uint128::zero(), |x| x.get_value()), - ) - .unwrap(); - assert_eq!(internal_node.get_value(), accumulated_value); - - let min = left_node - .clone() - .map_or(Uint128::MAX, |n| n.get_min_range()) - .min( - right_node - .clone() - .map_or(Uint128::MAX, |n| n.get_min_range()), - ); - let max = left_node - .map_or(Uint128::MIN, |n| n.get_max_range()) - .max(right_node.map_or(Uint128::MIN, |n| n.get_max_range())); - assert_eq!(internal_node.get_min_range(), min); - assert_eq!(internal_node.get_max_range(), max); - } - } -} - -const SPACING: u32 = 2u32; -const RIGHT_CORNER: &str = "┐"; -const LEFT_CORNER: &str = "┌"; -const STRAIGHT: &str = "─"; - -pub fn spacing(len: u32) -> String { - let mut s = "".to_string(); - for _ in 0..len { - s.push(' '); + assert_internal_values(test.name, deps.as_ref(), internals, true); } - s -} - -pub fn print_tree_row(row: Vec<(Option, Option)>, top: bool, height: u32) { - let blank_spacing_length = 2u32.pow(height + 1) * SPACING; - let blank_spacing = spacing(blank_spacing_length); - - let mut node_spacing = "".to_string(); - for _ in 0..blank_spacing_length { - node_spacing.push_str(STRAIGHT); - } - - if !top { - let mut line = "".to_string(); - for (left, right) in row.clone() { - let print_left_top = if left.is_some() { - format!("{blank_spacing}{LEFT_CORNER}{node_spacing}") - } else { - spacing(blank_spacing_length * 2) - }; - let print_right_top = if right.is_some() { - format!("{node_spacing}{RIGHT_CORNER}{blank_spacing}") - } else { - spacing(blank_spacing_length * 2) - }; - line.push_str(format!("{print_left_top}{print_right_top}").as_str()) - } - println!("{line}") - } - - let mut line = "".to_string(); - for (left, right) in row { - let left_node_length = if let Some(left) = left.clone() { - left.to_string().len() - } else { - 0 - }; - let print_left_top = if let Some(left) = left { - // Shift spacing to adjust for length of node string - let left_space = - spacing(blank_spacing_length - (left_node_length as f32 / 2.0).ceil() as u32); - - format!("{left_space}{left}{blank_spacing}") - } else { - spacing(blank_spacing_length * 2) - }; - let right_node_length = if let Some(right) = right.clone() { - right.to_string().len() - } else { - 0 - }; - let print_right_top = if let Some(right) = right { - // Shift spacing to adjust for length of left and right node string - let right_space = spacing( - blank_spacing_length - - (right_node_length as f32 / 2.0).ceil() as u32 - - (left_node_length as f32 / 2.0).floor() as u32, - ); - format!("{right_space}{right}{blank_spacing}") - } else if !top { - spacing(blank_spacing_length * 2) - } else { - // Prevents root from going on to new line - "".to_string() - }; - line.push_str(format!("{print_left_top}{print_right_top}").as_str()) - } - println!("{line}") } struct NodeDeletionTestCase { @@ -504,13 +452,7 @@ fn test_node_deletion_valid() { } if test.print { - println!("Pre-Deletion Tree: {}", test.name); - println!("--------------------------"); - let nodes = tree.traverse_bfs(deps.as_ref().storage).unwrap(); - for (idx, row) in nodes.iter().enumerate() { - print_tree_row(row.clone(), idx == 0, (nodes.len() - idx - 1) as u32); - } - println!(); + print_tree("Pre-Deletion Tree", test.name, &tree, &deps.as_ref()); } for key in test.delete.clone() { @@ -531,13 +473,7 @@ fn test_node_deletion_valid() { .unwrap(); if test.print { - println!("Post-Deletion Tree: {}", test.name); - println!("--------------------------"); - let nodes = tree.traverse_bfs(deps.as_ref().storage).unwrap(); - for (idx, row) in nodes.iter().enumerate() { - print_tree_row(row.clone(), idx == 0, (nodes.len() - idx - 1) as u32); - } - println!(); + print_tree("Post-Deletion Tree", test.name, &tree, &deps.as_ref()); } let result = tree.traverse(deps.as_ref().storage).unwrap(); @@ -560,43 +496,1461 @@ fn test_node_deletion_valid() { } let internals: Vec<&TreeNode> = result.iter().filter(|x| x.is_internal()).collect(); - for internal_node in internals { - let left_node = internal_node.get_left(deps.as_ref().storage).unwrap(); - let right_node = internal_node.get_right(deps.as_ref().storage).unwrap(); - - let accumulated_value = left_node - .clone() - .map(|x| x.get_value()) - .unwrap_or_default() - .checked_add( - right_node - .clone() - .map(|x| x.get_value()) - .unwrap_or_default(), + assert_internal_values(test.name, deps.as_ref(), internals, true); + } +} + +struct RotateRightTestCase { + name: &'static str, + nodes: Vec, + expected: Vec, + expected_error: Option, + print: bool, +} + +#[test] +fn test_rotate_right() { + let book_id = 1; + let tick_id = 1; + let test_cases: Vec = vec![ + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌──────────────── + // 2: 2 1-3 + // ┌────────────────┐ + // 3: 1 1 4: 2 1 + // + // Post-rotation + // -------------------------- + // 2: 2 1-3 + // ┌────────────────────────────────┐ + // 3: 1 1 1: 1 2-3 + // ┌──────── + // 4: 2 1 + RotateRightTestCase { + name: "Left internal right empty", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), None), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), Some(4)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Left-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + ], + expected: vec![2, 3, 1, 4], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1-3 4: 2 1 + // ┌──────── + // 3: 1 1 + // + // Post-rotation + // -------------------------- + // 2: 2 1-3 + // ┌────────────────────────────────┐ + // 3: 1 1 1: 1 2-3 + // ────────┐ + // 4: 2 1 + RotateRightTestCase { + name: "Left internal (no left-right ancestor) right leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(4)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), None) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![2, 3, 1, 4], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1-3 4: 2 1 + // ────────┐ + // 3: 1 1 + // + // Post-rotation + // -------------------------- + // 2: 2 1-3 + // ────────────────┐ + // 1: 2 1-3 + // ┌────────────────┐ + // 3: 1 1 4: 2 1 + RotateRightTestCase { + name: "Left internal (no left-left ancestor) right leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(4)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(None, Some(3)) + .with_parent(1), + // Left-Right + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![2, 1, 3, 4], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1-3 5: 3 1 + // ┌────────────────┐ + // 3: 1 1 4: 2 1 + // + // Post-rotation + // -------------------------- + // 2: 3 1-4 + // ┌────────────────────────────────┐ + // 3: 1 1 1: 2 2-4 + // ┌────────────────┐ + // 4: 2 1 5: 3 1 + RotateRightTestCase { + name: "Left internal right leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(5)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), Some(4)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Left-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(3u32, 1u32)).with_parent(1), + ], + expected: vec![2, 3, 1, 4, 5], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1-3 5: 2 3-5 + // ┌────────────────┐ ┌────────────────┐ + // 3: 1 1 4: 2 1 6: 3 1 7: 4 1 + // + // Post-rotation + // -------------------------- + // 2: 4 1-5 + // ┌────────────────────────────────────────────────────────────────┐ + // 3: 1 1 1: 3 2-5 + // ┌────────────────────────────────┐ + // 4: 2 1 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateRightTestCase { + name: "Left internal right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(5)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), Some(4)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Left-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![2, 3, 1, 4, 5, 6, 7], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1-3 5: 2 3-5 + // ────────┐ ┌────────────────┐ + // 4: 2 1 6: 3 1 7: 4 1 + // + // Post-rotation + // -------------------------- + // 2: 3 2-5 + // ────────────────┐ + // 1: 3 2-5 + // ┌────────────────────────────────┐ + // 4: 2 1 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateRightTestCase { + name: "Left internal (no left-left) right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(5)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(None, Some(4)) + .with_parent(1), + // Left-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![2, 1, 4, 5, 6, 7], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1-3 5: 2 3-5 + // ┌──────── ┌────────────────┐ + // 4: 2 1 6: 3 1 7: 4 1 + // + // Post-rotation + // -------------------------- + // 2: 3 2-5 + // ┌────────────────────────────────────────────────────────────────┐ + // 4: 2 1 1: 2 3-5 + // ────────────────┐ + // 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateRightTestCase { + name: "Left internal (no left-right) right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(5)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(4), None) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![2, 4, 1, 5, 6, 7], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ─────────────────┐ + // 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateRightTestCase { + name: "Left empty right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), ) + .with_children(None, Some(5)), + // Right + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![], + expected_error: Some(ContractError::InvalidNodeType), + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateRightTestCase { + name: "Left leaf right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(5)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(2u32, 1u32)).with_parent(1), + // Right + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![], + expected_error: Some(ContractError::InvalidNodeType), + print: true, + }, + ]; + + for mut test in test_cases { + let mut deps = mock_dependencies(); + // Save nodes in storage + for (idx, node) in test.nodes.iter_mut().enumerate() { + // Save root node + if idx == 0 { + TREE.save(deps.as_mut().storage, &(book_id, tick_id), &node.key) + .unwrap(); + } + NODES + .save(deps.as_mut().storage, &(book_id, tick_id, node.key), node) .unwrap(); - assert_eq!(internal_node.get_value(), accumulated_value); - - let min = left_node - .clone() - .map(|n| n.get_min_range()) - .unwrap_or(Uint128::MAX) - .min( - right_node - .clone() - .map(|n| n.get_min_range()) - .unwrap_or(Uint128::MAX), - ); - let max = left_node - .map(|n| n.get_max_range()) - .unwrap_or(Uint128::MIN) - .max( - right_node - .map(|n| n.get_max_range()) - .unwrap_or(Uint128::MIN), - ); - assert_eq!(internal_node.get_min_range(), min); - assert_eq!(internal_node.get_max_range(), max); } + + // Sync weights post node storage + for mut node in test.nodes { + if node.is_internal() || node.parent.is_none() { + continue; + } + node.sync(deps.as_ref().storage).unwrap(); + let mut parent = node.get_parent(deps.as_ref().storage).unwrap().unwrap(); + parent + .sync_range_and_value_up(deps.as_mut().storage) + .unwrap(); + } + + let mut tree = get_root_node(deps.as_ref().storage, book_id, tick_id).unwrap(); + if test.print { + print_tree("Pre-rotation", test.name, &tree, &deps.as_ref()); + } + + let res = tree.rotate_right(deps.as_mut().storage); + + if let Some(err) = test.expected_error { + assert_eq!(res, Err(err), "{}", test.name); + continue; + } + + // Get new root node, it may have changed due to rotationss + let tree = get_root_node(deps.as_ref().storage, book_id, tick_id).unwrap(); + if test.print { + print_tree("Post-rotation", test.name, &tree, &deps.as_ref()); + } + + let nodes = tree.traverse(deps.as_ref().storage).unwrap(); + let res: Vec = nodes.iter().map(|n| n.key).collect(); + assert_eq!(res, test.expected, "{}", test.name); + + let internals = nodes.iter().filter(|n| n.is_internal()).collect(); + assert_internal_values(test.name, deps.as_ref(), internals, false); } } + +struct RotateLeftTestCase { + name: &'static str, + nodes: Vec, + expected: Vec, + expected_error: Option, + print: bool, +} + +#[test] +fn test_rotate_left() { + let book_id = 1; + let tick_id = 1; + let test_cases: Vec = vec![ + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ────────────────┐ + // 2: 2 1-3 + // ┌────────────────┐ + // 3: 1 1 4: 2 1 + // + // Post-rotation + // -------------------------- + // 2: 2 1-3 + // ┌────────────────────────────────┐ + // 1: 1 1-2 4: 2 1 + // ────────┐ + // 3: 1 1 + RotateLeftTestCase { + name: "Right internal left empty", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(None, Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), Some(4)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Right-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + ], + expected: vec![2, 1, 3, 4], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 4: 2 1 2: 2 1-3 + // ┌──────── + // 3: 1 1 + // + // Post-rotation + // -------------------------- + // 2: 2 1-3 + // ┌──────────────── + // 1: 2 1-3 + // ┌────────────────┐ + // 4: 2 1 3: 1 1 + RotateLeftTestCase { + name: "Right internal (no right-right ancestor) left leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(4), Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), None) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![2, 1, 4, 3], + expected_error: None, + print: true, + }, + // Pre-rotation: + // -------------------------- + // 1: 2 1-3 + // ┌────────────────────────────────┐ + // 4: 2 1 2: 1 1-2 + // ────────┐ + // 3: 1 1 + // + // Post-rotation: + // -------------------------- + // 2: 2 1-3 + // ┌────────────────────────────────┐ + // 1: 1 2-3 3: 1 1 + // ┌──────── + // 4: 2 1 + RotateLeftTestCase { + name: "Right internal (no right-left ancestor) left leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(4), Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(None, Some(3)) + .with_parent(1), + // Right-Right + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![2, 1, 4, 3], + expected_error: None, + print: true, + }, + // Pre-rotation: + // -------------------------- + // 1: 3 1-4 + // ┌────────────────────────────────┐ + // 5: 3 1 2: 2 1-3 + // ┌────────────────┐ + // 3: 1 1 4: 2 1 + // + // Post-rotation: + // -------------------------- + // 2: 3 1-4 + // ┌────────────────────────────────┐ + // 1: 2 1-4 4: 2 1 + // ┌────────────────┐ + // 5: 3 1 3: 1 1 + RotateLeftTestCase { + name: "Right internal left leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(5), Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(3), Some(4)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(2), + // Right-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Left + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(3u32, 1u32)).with_parent(1), + ], + expected: vec![2, 1, 5, 3, 4], + expected_error: None, + print: true, + }, + // Pre-rotation: + // -------------------------- + // 1: 3 2-5 + // ┌────────────────────────────────┐ + // 5: 2 3-5 2: 1 2-3 + // ┌────────────────┐ ────────┐ + // 6: 3 1 7: 4 1 4: 2 1 + // + // Post-rotation: + // -------------------------- + // 2: 3 2-5 + // ┌────────────────────────────────────────────────────────────────┐ + // 1: 2 3-5 4: 2 1 + // ┌──────────────── + // 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateLeftTestCase { + name: "Right internal (no right-left) left internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(5), Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(None, Some(4)) + .with_parent(1), + // Right-Right + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Left + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![2, 1, 5, 6, 7, 4], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 3 2-5 + // ┌────────────────────────────────┐ + // 5: 2 3-5 2: 1 2-3 + // ┌────────────────┐ ┌──────── + // 6: 3 1 7: 4 1 4: 2 1 + // + // Post-rotation + // -------------------------- + // 2: 3 2-5 + // ┌──────────────────────────────── + // 1: 3 2-5 + // ┌────────────────────────────────┐ + // 5: 2 3-5 4: 2 1 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateLeftTestCase { + name: "Right internal (no right-right) left internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(5), Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(4), None) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Left + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![2, 1, 5, 6, 7, 4], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ─────────────────┐ + // 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateLeftTestCase { + name: "Right empty left internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(5), None), + // Left + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![], + expected_error: Some(ContractError::InvalidNodeType), + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────┐ + // 2: 2 1 5: 2 3-5 + // ┌────────────────┐ + // 6: 3 1 7: 4 1 + RotateLeftTestCase { + name: "Right leaf left internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(5), Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(2u32, 1u32)).with_parent(1), + // Left + TreeNode::new(book_id, tick_id, 5, NodeType::internal(2u32, (3u32, 5u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(3u32, 1u32)).with_parent(5), + // Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(4u32, 1u32)).with_parent(5), + ], + expected: vec![], + expected_error: Some(ContractError::InvalidNodeType), + print: true, + }, + ]; + + for mut test in test_cases { + let mut deps = mock_dependencies(); + // Save nodes in storage + for (idx, node) in test.nodes.iter_mut().enumerate() { + // Save root node + if idx == 0 { + TREE.save(deps.as_mut().storage, &(book_id, tick_id), &node.key) + .unwrap(); + } + NODES + .save(deps.as_mut().storage, &(book_id, tick_id, node.key), node) + .unwrap(); + } + + // Sync weights post node storage + for mut node in test.nodes { + if node.is_internal() || node.parent.is_none() { + continue; + } + node.sync(deps.as_ref().storage).unwrap(); + let mut parent = node.get_parent(deps.as_ref().storage).unwrap().unwrap(); + parent + .sync_range_and_value_up(deps.as_mut().storage) + .unwrap(); + } + + let mut tree = get_root_node(deps.as_ref().storage, book_id, tick_id).unwrap(); + if test.print { + print_tree("Pre-rotation", test.name, &tree, &deps.as_ref()); + } + + let res = tree.rotate_left(deps.as_mut().storage); + + if let Some(err) = test.expected_error { + assert_eq!(res, Err(err), "{}", test.name); + continue; + } + + // Get new root node, it may have changed due to rotations + let tree = get_root_node(deps.as_ref().storage, book_id, tick_id).unwrap(); + if test.print { + print_tree("Post-rotation", test.name, &tree, &deps.as_ref()); + } + + let nodes = tree.traverse(deps.as_ref().storage).unwrap(); + let res: Vec = nodes.iter().map(|n| n.key).collect(); + assert_eq!(res, test.expected, "{}", test.name); + + let internals = nodes.iter().filter(|n| n.is_internal()).collect(); + assert_internal_values(test.name, deps.as_ref(), internals, false); + } +} + +struct RebalanceTestCase { + name: &'static str, + nodes: Vec, + expected: Vec, + expected_error: Option, + print: bool, +} + +#[test] +fn test_rebalance() { + let book_id = 1; + let tick_id = 1; + let test_cases: Vec = vec![ + // Pre-rotation: Case 1: Right Right + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────────────────────────────────────┐ + // 2: 1 1 3: 2 1-3 + // ┌────────────────────────────────┐ + // 4: 4 1 5: 0 4294967295-0 + // ┌────────────────┐ + // 6: 2 1 7: 3 1 + // + // Post-rotation: Case 1: Right Right + // -------------------------- + // 3: 2 1-5 + // ┌────────────────────────────────┐ + // 1: 2 1-5 5: 0 4294967295-0 + // ┌────────────────┐ ┌────────────────┐ + // 2: 1 1 4: 4 1 6: 2 1 7: 3 1 + RebalanceTestCase { + name: "Case 1: Right Right", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(1u32, 1u32)).with_parent(1), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(4u32, 1u32)).with_parent(3), + // Right-Right + TreeNode::new( + book_id, + tick_id, + 5, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(6), Some(7)) + .with_parent(3), + // Right-Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(2u32, 1u32)).with_parent(5), + // Right-Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(3u32, 1u32)).with_parent(5), + ], + expected: vec![3, 1, 2, 4, 5, 6, 7], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 4 1-5 + // ┌────────────────────────────────────────────────────────────────┐ + // 2: 3 2-5 3: 1 1 + // ┌────────────────────────────────┐ + // 4: 2 2-4 5: 4 1 + // ┌────────────────┐ + // 6: 2 1 7: 3 1 + // + // Post-rotation + // -------------------------- + // 2: 4 1-5 + // ┌────────────────────────────────┐ + // 4: 2 2-4 1: 2 1-5 + // ┌────────────────┐ ┌────────────────┐ + // 6: 2 1 7: 3 1 5: 4 1 3: 1 1 + RebalanceTestCase { + name: "Case 2: Left Left", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Left-Left + TreeNode::new( + book_id, + tick_id, + 4, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(6), Some(7)) + .with_parent(2), + // Left-Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(2u32, 1u32)).with_parent(4), + // Left-Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(3u32, 1u32)).with_parent(4), + // Left-Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(4u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(1), + ], + expected: vec![2, 4, 6, 7, 1, 5, 3], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 0 4294967295-0 + // ┌────────────────────────────────────────────────────────────────┐ + // 2: 1 1 3: 2 1-3 + // ┌────────────────────────────────┐ + // 4: 0 4294967295-0 5: 4 1 + // ┌────────────────┐ + // 6: 2 1 7: 3 1 + // + // Post-rotation + // -------------------------- + // 4: 4 1-5 + // ┌────────────────────────────────┐ + // 1: 2 1-3 3: 2 3-5 + // ┌────────────────┐ ┌────────────────┐ + // 2: 1 1 6: 2 1 7: 3 1 5: 4 1 + RebalanceTestCase { + name: "Case 3: Right Left", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(1u32, 1u32)).with_parent(1), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Right-Left + TreeNode::new( + book_id, + tick_id, + 4, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(6), Some(7)) + .with_parent(3), + // Right-Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(2u32, 1u32)).with_parent(4), + // Right-Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(3u32, 1u32)).with_parent(4), + // Right-Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(4u32, 1u32)).with_parent(3), + ], + expected: vec![4, 1, 2, 6, 3, 7, 5], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 4 1-5 + // ┌────────────────────────────────────────────────────────────────┐ + // 2: 3 2-5 3: 1 1 + // ┌────────────────────────────────┐ + // 4: 2 2-4 5: 4 1 + // ┌────────────────┐ + // 6: 2 1 7: 3 1 + // + // Post-rotation + // -------------------------- + // 2: 4 1-5 + // ┌────────────────────────────────┐ + // 4: 2 2-4 1: 2 1-5 + // ┌────────────────┐ ┌────────────────┐ + // 6: 2 1 7: 3 1 5: 4 1 3: 1 1 + RebalanceTestCase { + name: "Case 4: Left Right", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (1u32, 3u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Left-Left + TreeNode::new( + book_id, + tick_id, + 4, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(6), Some(7)) + .with_parent(2), + // Left-Left-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Left-Left-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(3u32, 1u32)).with_parent(2), + // Left-Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(4u32, 1u32)).with_parent(4), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(1u32, 1u32)).with_parent(1), + ], + expected: vec![2, 4, 6, 7, 1, 5, 3], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 1 2-3 + // ┌──────── + // 2: 2 1 + // + // Post-rotation + // -------------------------- + // 1: 1 2-3 + // ┌──────── + // 2: 2 1 + RebalanceTestCase { + name: "Pre-balanced: left leaf right empty", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), None), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![1, 2], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 1 2-3 + // ────────┐ + // 2: 2 1 + // + // Post-rotation + // -------------------------- + // 1: 1 2-3 + // ────────┐ + // 2: 2 1 + RebalanceTestCase { + name: "Pre-balanced: left empty right leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(None, Some(2)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![1, 2], + expected_error: None, + print: true, + }, + // Pre-rotation: + // -------------------------- + // 1: 3 2-4 + // ┌────────────────────────────────┐ + // 2: 2 2-4 3: 2 1 + // ┌────────────────┐ + // 4: 2 1 5: 3 1 + // + // Post-rotation: + // -------------------------- + // 1: 3 2-4 + // ┌────────────────────────────────┐ + // 2: 2 2-4 3: 2 1 + // ┌────────────────┐ + // 4: 2 1 5: 3 1 + RebalanceTestCase { + name: "Pre-balanced: left internal right leaf", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (2u32, 4u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Left-Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Left-Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(3u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::leaf(2u32, 1u32)).with_parent(1), + ], + expected: vec![1, 2, 4, 5, 3], + expected_error: None, + print: true, + }, + // Pre-rotation + // -------------------------- + // 1: 3 2-4 + // ┌────────────────────────────────┐ + // 2: 2 1 3: 2 2-4 + // ┌────────────────┐ + // 4: 2 1 5: 3 1 + // + // Post-rotation + // -------------------------- + // 1: 3 2-4 + // ┌────────────────────────────────┐ + // 2: 2 1 3: 2 2-4 + // ┌────────────────┐ + // 4: 2 1 5: 3 1 + RebalanceTestCase { + name: "Pre-balanced: left leaf right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Left + TreeNode::new(book_id, tick_id, 2, NodeType::leaf(2u32, 1u32)).with_parent(1), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::internal(2u32, (2u32, 4u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(3), + // Right-Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(3u32, 1u32)).with_parent(3), + ], + expected: vec![1, 2, 3, 4, 5], + expected_error: None, + print: true, + }, + // Pre-rotation: + // -------------------------- + // 1: 4 2-6 + // ┌────────────────────────────────┐ + // 2: 2 2-4 3: 2 4-6 + // ┌────────────────┐ ┌────────────────┐ + // 4: 2 1 5: 3 1 6: 4 1 7: 5 1 + // + // Post-rotation: + // -------------------------- + // 1: 4 2-6 + // ┌────────────────────────────────┐ + // 2: 2 2-4 3: 2 4-6 + // ┌────────────────┐ ┌────────────────┐ + // 4: 2 1 5: 3 1 6: 4 1 7: 5 1 + RebalanceTestCase { + name: "Pre-balanced: left internal right internal", + nodes: vec![ + // Root + TreeNode::new( + book_id, + tick_id, + 1, + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ) + .with_children(Some(2), Some(3)), + // Right + TreeNode::new(book_id, tick_id, 2, NodeType::internal(2u32, (2u32, 4u32))) + .with_children(Some(4), Some(5)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 4, NodeType::leaf(2u32, 1u32)).with_parent(2), + // Right-Right + TreeNode::new(book_id, tick_id, 5, NodeType::leaf(3u32, 1u32)).with_parent(2), + // Right + TreeNode::new(book_id, tick_id, 3, NodeType::internal(2u32, (2u32, 4u32))) + .with_children(Some(6), Some(7)) + .with_parent(1), + // Right-Left + TreeNode::new(book_id, tick_id, 6, NodeType::leaf(4u32, 1u32)).with_parent(3), + // Right-Right + TreeNode::new(book_id, tick_id, 7, NodeType::leaf(5u32, 1u32)).with_parent(3), + ], + expected: vec![1, 2, 4, 5, 3, 6, 7], + expected_error: None, + print: true, + }, + RebalanceTestCase { + name: "invalid node type", + nodes: vec![ + // Root + TreeNode::new(book_id, tick_id, 1, NodeType::leaf(1u32, 1u32)), + ], + expected: vec![], + expected_error: Some(ContractError::InvalidNodeType), + print: true, + }, + RebalanceTestCase { + name: "childless internal node", + nodes: vec![ + // Root + TreeNode::new(book_id, tick_id, 1, NodeType::internal(1u32, (1u32, 2u32))), + ], + expected: vec![], + expected_error: Some(ContractError::ChildlessInternalNode), + print: true, + }, + ]; + + for mut test in test_cases { + let mut deps = mock_dependencies(); + // Save nodes in storage + for (idx, node) in test.nodes.iter_mut().enumerate() { + // Save root node + if idx == 0 { + TREE.save(deps.as_mut().storage, &(book_id, tick_id), &node.key) + .unwrap(); + } + NODES + .save(deps.as_mut().storage, &(book_id, tick_id, node.key), node) + .unwrap(); + } + + // Sync weights post node storage + for mut node in test.nodes { + if node.is_internal() || node.parent.is_none() { + continue; + } + node.sync(deps.as_ref().storage).unwrap(); + let mut parent = node.get_parent(deps.as_ref().storage).unwrap().unwrap(); + parent + .sync_range_and_value_up(deps.as_mut().storage) + .unwrap(); + } + + let mut tree = get_root_node(deps.as_ref().storage, book_id, tick_id).unwrap(); + if test.print { + print_tree("Pre-rotation", test.name, &tree, &deps.as_ref()); + } + + let res = tree.rebalance(deps.as_mut().storage); + + if let Some(err) = test.expected_error { + assert_eq!(res, Err(err), "{}", test.name); + continue; + } + + // Get new root node, it may have changed due to rotations + let tree = get_root_node(deps.as_ref().storage, book_id, tick_id).unwrap(); + if test.print { + print_tree("Post-rotation", test.name, &tree, &deps.as_ref()); + } + + let nodes = tree.traverse(deps.as_ref().storage).unwrap(); + let res: Vec = nodes.iter().map(|n| n.key).collect(); + assert_eq!(res, test.expected, "{}", test.name); + + let internals = nodes.iter().filter(|n| n.is_internal()).collect(); + assert_internal_values(test.name, deps.as_ref(), internals, true); + } +} + +fn generate_nodes( + storage: &mut dyn Storage, + book_id: u64, + tick_id: i64, + quantity: u32, +) -> Vec { + use rand::seq::SliceRandom; + use rand::thread_rng; + + let mut range: Vec = (0..quantity).collect(); + range.shuffle(&mut thread_rng()); + + let mut nodes = vec![]; + for val in range { + let id = generate_node_id(storage, book_id, tick_id).unwrap(); + nodes.push(TreeNode::new( + book_id, + tick_id, + id, + NodeType::leaf(val, 1u32), + )); + } + nodes +} + +#[test] +fn test_node_insert_large_quantity() { + let book_id = 1; + let tick_id = 1; + + let mut deps = mock_dependencies(); + // Create tree root + let mut tree = TreeNode::new( + book_id, + tick_id, + generate_node_id(deps.as_mut().storage, book_id, tick_id).unwrap(), + NodeType::internal(Uint128::zero(), (u32::MAX, u32::MIN)), + ); + + let nodes = generate_nodes(deps.as_mut().storage, book_id, tick_id, 1000); + + // Insert nodes into tree + for mut node in nodes { + NODES + .save(deps.as_mut().storage, &(book_id, tick_id, node.key), &node) + .unwrap(); + tree.insert(deps.as_mut().storage, &mut node).unwrap(); + } + + // Return tree in vector form from Depth First Search + let result = tree.traverse(deps.as_ref().storage).unwrap(); + + // Ensure all internal nodes are correctly summed and contain correct ranges + let internals: Vec<&TreeNode> = result.iter().filter(|x| x.is_internal()).collect(); + assert_internal_values("Large amount of nodes", deps.as_ref(), internals, true); +} + +const SPACING: u32 = 2u32; +const RIGHT_CORNER: &str = "┐"; +const LEFT_CORNER: &str = "┌"; +const STRAIGHT: &str = "─"; + +pub fn spacing(len: u32) -> String { + let mut s = "".to_string(); + for _ in 0..len { + s.push(' '); + } + s +} + +pub fn print_tree(title: &'static str, test_name: &'static str, root: &TreeNode, deps: &Deps) { + println!("{}: {}", title, test_name); + println!("--------------------------"); + let nodes = root.traverse_bfs(deps.storage).unwrap(); + for (idx, row) in nodes.iter().enumerate() { + print_tree_row(row.clone(), idx == 0, (nodes.len() - idx - 1) as u32); + } + println!(); +} + +pub fn print_tree_row(row: Vec<(Option, Option)>, top: bool, height: u32) { + let blank_spacing_length = 2u32.pow(height + 1) * SPACING; + let blank_spacing = spacing(blank_spacing_length); + + let mut node_spacing = "".to_string(); + for _ in 0..blank_spacing_length { + node_spacing.push_str(STRAIGHT); + } + + if !top { + let mut line = "".to_string(); + for (left, right) in row.clone() { + let print_left_top = if left.is_some() { + format!("{blank_spacing}{LEFT_CORNER}{node_spacing}") + } else { + spacing(blank_spacing_length * 2) + }; + let print_right_top = if right.is_some() { + format!("{node_spacing}{RIGHT_CORNER}{blank_spacing}") + } else { + spacing(blank_spacing_length * 2) + }; + line.push_str(format!("{print_left_top}{print_right_top}").as_str()) + } + println!("{line}") + } + + let mut line = "".to_string(); + for (left, right) in row { + let left_node_length = if let Some(left) = left.clone() { + left.to_string().len() + } else { + 0 + }; + let print_left_top = if let Some(left) = left { + // Shift spacing to adjust for length of node string + let left_space = + spacing(blank_spacing_length - (left_node_length as f32 / 2.0).ceil() as u32); + + format!("{left_space}{left}{blank_spacing}") + } else { + spacing(blank_spacing_length * 2) + }; + let right_node_length = if let Some(right) = right.clone() { + right.to_string().len() + } else { + 0 + }; + let print_right_top = if let Some(right) = right { + // Shift spacing to adjust for length of left and right node string + let right_space = spacing( + blank_spacing_length + - (right_node_length as f32 / 2.0).ceil() as u32 + - (left_node_length as f32 / 2.0).floor() as u32, + ); + format!("{right_space}{right}{blank_spacing}") + } else if !top { + spacing(blank_spacing_length * 2) + } else { + // Prevents root from going on to new line + "".to_string() + }; + line.push_str(format!("{print_left_top}{print_right_top}").as_str()) + } + println!("{line}") +} diff --git a/contracts/sumtree-orderbook/src/sumtree/tree.rs b/contracts/sumtree-orderbook/src/sumtree/tree.rs index 3d81696..230fd6b 100644 --- a/contracts/sumtree-orderbook/src/sumtree/tree.rs +++ b/contracts/sumtree-orderbook/src/sumtree/tree.rs @@ -1,7 +1,19 @@ +use cosmwasm_std::Storage; use cw_storage_plus::Map; -use super::node::TreeNode; +use crate::error::ContractResult; + +use super::node::{TreeNode, NODES}; + +pub const TREE: Map<&(u64, i64), u64> = Map::new("tree"); -// TODO: REMOVE #[allow(dead_code)] -pub const TREE: Map<&(u64, i64), TreeNode> = Map::new("tree"); +/// Retrieves the root node of a specific book and tick from storage. +pub fn get_root_node( + storage: &dyn Storage, + book_id: u64, + tick_id: i64, +) -> ContractResult { + let root_id = TREE.load(storage, &(book_id, tick_id))?; + Ok(NODES.load(storage, &(book_id, tick_id, root_id))?) +} diff --git a/contracts/sumtree-orderbook/src/tests/test_order.rs b/contracts/sumtree-orderbook/src/tests/test_order.rs index 1bbe602..2626d97 100644 --- a/contracts/sumtree-orderbook/src/tests/test_order.rs +++ b/contracts/sumtree-orderbook/src/tests/test_order.rs @@ -6,7 +6,7 @@ use crate::{ state::*, sumtree::{ node::{NodeType, TreeNode}, - tree::TREE, + tree::get_root_node, }, types::{OrderDirection, REPLY_ID_REFUND}, }; @@ -578,9 +578,7 @@ fn test_cancel_limit() { // -- Sumtree -- // Ensure tree is saved correctly - let tree = TREE - .load(deps.as_ref().storage, &(valid_book_id, test.tick_id)) - .unwrap(); + let tree = get_root_node(deps.as_ref().storage, valid_book_id, test.tick_id).unwrap(); // Traverse the tree to check its form let res = tree.traverse(deps.as_ref().storage).unwrap(); @@ -590,6 +588,7 @@ fn test_cancel_limit() { 1, NodeType::internal(test.quantity, (0u128, test.quantity)), ); + root_node.set_weight(2).unwrap(); let mut cancelled_node = TreeNode::new( valid_book_id, test.tick_id,