Skip to content

Commit

Permalink
linear code
Browse files Browse the repository at this point in the history
  • Loading branch information
@rnbguy
rnbguy committed May 3, 2024
1 parent 19249e2 commit a42b75c
Showing 1 changed file with 81 additions and 90 deletions.
171 changes: 81 additions & 90 deletions basecoin/store/src/avl/tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,18 +76,14 @@ impl<K: Ord + AsBytes, V: Borrow<[u8]>> AvlTree<K, V> {

/// Remove a value from the tree.
fn remove_rec(node_ref: &mut NodeRef<K, V>, key: K) -> Option<V> {
let removed_value = if let Some(node) = node_ref {
match node.key.cmp(&key) {
Ordering::Greater => AvlTree::remove_rec(&mut node.left, key),
Ordering::Less => AvlTree::remove_rec(&mut node.right, key),
Ordering::Equal => AvlTree::remove_top(node_ref).map(|node| node.value),
}
} else {
None
let node = node_ref.as_deref_mut()?;

let removed_value = match node.key.cmp(&key) {
Ordering::Greater => AvlTree::remove_rec(&mut node.left, key),
Ordering::Less => AvlTree::remove_rec(&mut node.right, key),
Ordering::Equal => AvlTree::remove_top(node_ref).map(|node| node.value),
};

// this should be in the `if let Some(node) = node_ref` branch,
// just like `insert_rec`. But it is here to avoid borrowing issues.
if let Some(node) = node_ref {
node.update();
AvlTree::balance_node(node_ref);
Expand All @@ -98,59 +94,58 @@ impl<K: Ord + AsBytes, V: Borrow<[u8]>> AvlTree<K, V> {

/// Removes the top node in the tree, if it exists.
fn remove_top(node_ref: &mut NodeRef<K, V>) -> NodeRef<K, V> {
let removed_node = if let Some(node) = node_ref {
let substitute_node_ref = if node.right.is_none() {
// there is no right node, replace the current node with the left node.
// substitute_node_ref <- node.left
node.left.take()
} else if node.left.is_none() {
// there is no left node, replace the current node with the right node.
// substitute_node_ref <- node.right
node.right.take()
} else if node.balance_factor() <= 0 {
// both left and right nodes exist.
// left.height <= right.height; right skewed.
// removing from right subtree is better.

// Remove the leftmost node in the right subtree and replace the current.
let mut leftmost_node_ref = AvlTree::remove_leftmost(&mut node.right);
// leftmost_node_ref.right <- node_ref.right
// leftmost_node_ref.left <- node_ref.left
if let Some(leftmost_node) = leftmost_node_ref.as_mut() {
// removed leftmost node must be a leaf; not asserting, as it is an invariant.
// assert!(leftmost_node.right.is_none() && leftmost_node.left.is_none());

leftmost_node.right = node.right.take();
leftmost_node.left = node.left.take();
}
// substitute_node_ref <- leftmost_node_ref
leftmost_node_ref
} else {
// node.balance_factor() > 0
// both left and right nodes exist.
// left.height > right.height; left skewed.
// removing from left subtree is better.

// Remove the rightmost node in the left subtree and replace the current.
let mut rightmost_node_ref = AvlTree::remove_rightmost(&mut node.left);
// rightmost_node_ref.right <- node_ref.right
// rightmost_node_ref.left <- node_ref.left
if let Some(rightmost_node) = rightmost_node_ref.as_mut() {
// removed rightmost node must be a leaf; not asserting, as it is an invariant.
// assert!(rightmost_node.right.is_none() && rightmost_node.left.is_none());

rightmost_node.right = node.right.take();
rightmost_node.left = node.left.take();
}
// substitute_node_ref <- rightmost_node_ref
rightmost_node_ref
};
// removed_node <- node_ref <- substitute_node_ref
std::mem::replace(node_ref, substitute_node_ref)
let node = node_ref.as_deref_mut()?;

let substitute_node_ref = if node.right.is_none() {
// there is no right node, replace the current node with the left node.
// substitute_node_ref <- node.left
node.left.take()
} else if node.left.is_none() {
// there is no left node, replace the current node with the right node.
// substitute_node_ref <- node.right
node.right.take()
} else if node.balance_factor() <= 0 {
// both left and right nodes exist.
// left.height <= right.height; right skewed.
// removing from right subtree is better.

// Remove the leftmost node in the right subtree and replace the current.
let mut leftmost_node_ref = AvlTree::remove_leftmost(&mut node.right);
// leftmost_node_ref.right <- node_ref.right
// leftmost_node_ref.left <- node_ref.left
if let Some(leftmost_node) = leftmost_node_ref.as_mut() {
// removed leftmost node must be a leaf; not asserting, as it is an invariant.
// assert!(leftmost_node.right.is_none() && leftmost_node.left.is_none());

leftmost_node.right = node.right.take();
leftmost_node.left = node.left.take();
}
// substitute_node_ref <- leftmost_node_ref
leftmost_node_ref
} else {
None
// node.balance_factor() > 0
// both left and right nodes exist.
// left.height > right.height; left skewed.
// removing from left subtree is better.

// Remove the rightmost node in the left subtree and replace the current.
let mut rightmost_node_ref = AvlTree::remove_rightmost(&mut node.left);
// rightmost_node_ref.right <- node_ref.right
// rightmost_node_ref.left <- node_ref.left
if let Some(rightmost_node) = rightmost_node_ref.as_mut() {
// removed rightmost node must be a leaf; not asserting, as it is an invariant.
// assert!(rightmost_node.right.is_none() && rightmost_node.left.is_none());

rightmost_node.right = node.right.take();
rightmost_node.left = node.left.take();
}
// substitute_node_ref <- rightmost_node_ref
rightmost_node_ref
};

// removed_node <- node_ref <- substitute_node_ref
let removed_node = std::mem::replace(node_ref, substitute_node_ref);

if let Some(node) = node_ref {
// need to update, as top node is replaced
node.update();
Expand All @@ -162,47 +157,43 @@ impl<K: Ord + AsBytes, V: Borrow<[u8]>> AvlTree<K, V> {

/// Removes the leftmost key in the tree, if it exists.
fn remove_leftmost(node_ref: &mut NodeRef<K, V>) -> NodeRef<K, V> {
if let Some(node) = node_ref {
if node.left.is_none() {
let right_node = node.right.take();
// removed_node <- node_ref <- right_node
std::mem::replace(node_ref, right_node)

// no need to update, as current node (right_node) is already updated
} else {
let removed_node = AvlTree::remove_leftmost(&mut node.left);
let node = node_ref.as_deref_mut()?;

// need to update, as left node is updated
node.update();
AvlTree::balance_node(node_ref);
if node.left.is_none() {
let right_node = node.right.take();
// removed_node <- node_ref <- right_node
std::mem::replace(node_ref, right_node)

removed_node
}
// no need to update, as current node (right_node) is already updated
} else {
None
let removed_node = AvlTree::remove_leftmost(&mut node.left);

// need to update, as left node is updated
node.update();
AvlTree::balance_node(node_ref);

removed_node
}
}

/// Removes the rightmost key in the tree, if it exists.
fn remove_rightmost(node_ref: &mut NodeRef<K, V>) -> NodeRef<K, V> {
if let Some(node) = node_ref {
if node.right.is_none() {
let left_node = node.left.take();
// removed_node <- node_ref <- left_node
std::mem::replace(node_ref, left_node)

// no need to update, as current node (left_node) is already updated
} else {
let removed_node = AvlTree::remove_rightmost(&mut node.right);
let node = node_ref.as_deref_mut()?;

// need to update, as right node is updated
node.update();
AvlTree::balance_node(node_ref);
if node.right.is_none() {
let left_node = node.left.take();
// removed_node <- node_ref <- left_node
std::mem::replace(node_ref, left_node)

removed_node
}
// no need to update, as current node (left_node) is already updated
} else {
None
let removed_node = AvlTree::remove_rightmost(&mut node.right);

// need to update, as right node is updated
node.update();
AvlTree::balance_node(node_ref);

removed_node
}
}

Expand Down

0 comments on commit a42b75c

Please sign in to comment.