feat(store): implement AvlTree::remove (#141)

* impl remove for avltree

* impl Store::delete for InMemoryStore

* use InMemoryStore with delete support

* add test for AvlTree::remove

* mv test to tests file

* add shuffle test for get

* refactor insert and remove methods with return values

* rm redundant node updates

* fix rand dependency

* simplify mem::replace use

* assert balance factor

* insert 256 items

* test avltree thoroughly

* optimize remove_top algo

* comment

* comment for inconsistency

* optimize node removal

* rm old comment

* reuse mem::replace

* linear code

* update comments

* update comment

* minor refactor

basecoin/src/runner.rs

@@ -6,7 +6,7 @@ use basecoin_modules::gov::Governance;
 use basecoin_modules::ibc::Ibc;
 use basecoin_modules::staking::Staking;
 use basecoin_modules::upgrade::Upgrade;
-use basecoin_store::impls::{GrowingStore, InMemoryStore};
+use basecoin_store::impls::InMemoryStore;
 use ibc_proto::cosmos::base::tendermint::v1beta1::service_server::ServiceServer as HealthServer;
 use ibc_proto::cosmos::tx::v1beta1::service_server::ServiceServer as TxServer;
 #[cfg(all(feature = "v0_38", not(feature = "v0_37")))]
@@ -18,7 +18,7 @@ use crate::config::ServerConfig;
 
 pub async fn default_app_runner(server_cfg: ServerConfig) {
     // instantiate the application with a KV store implementation of choice
-    let app_builder = Builder::new(GrowingStore::<InMemoryStore>::default());
+    let app_builder = Builder::new(InMemoryStore::default());
 
     // instantiate modules and setup inter-module communication (if required)
     let auth = Auth::new(app_builder.module_store(&prefix::Auth {}.identifier()));

basecoin/store/src/avl/tests.rs

@@ -37,6 +37,55 @@ fn get() {
     assert_eq!(tree.get(&[4]), None);
 }
 
+#[test]
+fn test_shuffle_insert_get_remove() {
+    let mut tree = AvlTree::new();
+    let mut std_tree = std::collections::BTreeMap::new();
+
+    let mut keys: Vec<u8> = (0..=255).collect();
+
+    keys.shuffle(&mut thread_rng());
+    for &i in keys.iter() {
+        tree.insert([i], vec![i]);
+        std_tree.insert([i], vec![i]);
+
+        // keys from BTreeMap must be present in AvlTree.
+        for key in std_tree.keys() {
+            assert_eq!(tree.get(key), std_tree.get(key));
+        }
+
+        // keys from AvlTree must be present in BTreeMap.
+        for key in tree.get_keys() {
+            assert_eq!(tree.get(key), std_tree.get(key));
+        }
+
+        // AvlTree must stay balanced.
+        assert!(tree.root.as_ref().unwrap().balance_factor().abs() <= 1);
+    }
+
+    keys.shuffle(&mut thread_rng());
+    for &i in keys.iter() {
+        // keys from BTreeMap must be present in AvlTree.
+        for key in std_tree.keys() {
+            assert_eq!(tree.get(key), std_tree.get(key));
+        }
+
+        // keys from AvlTree must be present in BTreeMap.
+        for key in tree.get_keys() {
+            assert_eq!(tree.get(key), std_tree.get(key));
+        }
+
+        // AvlTree must stay balanced.
+        assert!(tree.root.as_ref().unwrap().balance_factor().abs() <= 1);
+
+        assert_eq!(tree.remove([i]), Some(vec![i]));
+        assert_eq!(std_tree.remove(&[i]), Some(vec![i]));
+    }
+
+    assert!(std_tree.is_empty());
+    assert!(tree.root.is_none());
+}
+
 #[test]
 fn rotate_right() {
     let mut before = AvlTree {

basecoin/store/src/avl/tree.rs

@@ -48,23 +48,172 @@ impl<K: Ord + AsBytes, V: Borrow<[u8]>> AvlTree<K, V> {
     /// Insert a value into the AVL tree, this operation runs in amortized O(log(n)).
     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
         let node_ref = &mut self.root;
-        let mut old_value = None;
-        AvlTree::insert_rec(node_ref, key, value, &mut old_value);
-        old_value
+        AvlTree::insert_rec(node_ref, key, value)
     }
 
     /// Insert a value in the tree.
-    fn insert_rec(node_ref: &mut NodeRef<K, V>, key: K, value: V, old_value: &mut Option<V>) {
+    fn insert_rec(node_ref: &mut NodeRef<K, V>, key: K, value: V) -> Option<V> {
         if let Some(node) = node_ref {
-            match node.key.cmp(&key) {
-                Ordering::Greater => AvlTree::insert_rec(&mut node.left, key, value, old_value),
-                Ordering::Less => AvlTree::insert_rec(&mut node.right, key, value, old_value),
-                Ordering::Equal => *old_value = Some(node.set_value(value)),
-            }
+            let old_value = match node.key.cmp(&key) {
+                Ordering::Greater => AvlTree::insert_rec(&mut node.left, key, value),
+                Ordering::Less => AvlTree::insert_rec(&mut node.right, key, value),
+                Ordering::Equal => Some(node.set_value(value)),
+            };
             node.update();
+            // Note: when old_value is None, balancing not necessary.
+            // But we do it anyway as general rule.
             AvlTree::balance_node(node_ref);
+            old_value
         } else {
             *node_ref = as_node_ref(key, value);
+            None
+        }
+    }
+
+    /// Remove a value from the AVL tree, this operation runs in amortized O(log(n)).
+    pub fn remove(&mut self, key: K) -> Option<V> {
+        let node_ref = &mut self.root;
+        AvlTree::remove_rec(node_ref, key).map(|node| node.value)
+    }
+
+    /// Remove a value from the tree.
+    ///
+    /// Find the sub-tree whose root to be removed. Then, use [`Self::remove_root`].
+    fn remove_rec(node_ref: &mut NodeRef<K, V>, key: K) -> NodeRef<K, V> {
+        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_root(node_ref),
+        };
+
+        if let Some(node) = node_ref {
+            if removed_value.is_some() {
+                // need to update, as root node is updated
+                // Note: if removed_value is None, nothing is removed.
+                // So no need to update and balance.
+                node.update();
+                AvlTree::balance_node(node_ref);
+            }
+        }
+
+        removed_value
+    }
+
+    /// Removes the root node in the tree, if it exists.
+    ///
+    /// Since we are removing the root node, we need to replace it with a new node.
+    /// The new node is chosen as follows:
+    /// - If the root node has no children, the new node is `None`.
+    /// - If the root node has only one child, the new node is the child.
+    /// - If the root node has both children.
+    ///   - If left child is shorter: the new node is the leftmost node in the right subtree.
+    ///     Also, the root node's children are set to the new node's children.
+    ///   - If right child is shorter, vice versa.
+    ///
+    /// Never called on an empty tree.
+    fn remove_root(node_ref: &mut NodeRef<K, V>) -> NodeRef<K, V> {
+        let node = node_ref.as_deref_mut()?;
+
+        let substitute_node_ref = if node.right.is_none() {
+            // there is no right node, replace the root 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 root 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; 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; 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();
+            AvlTree::balance_node(node_ref);
+        }
+
+        removed_node
+    }
+
+    /// Removes the leftmost key in the tree, if it exists.
+    fn remove_leftmost(node_ref: &mut NodeRef<K, V>) -> NodeRef<K, V> {
+        let node = node_ref.as_deref_mut()?;
+
+        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 the new node (right_node) is already updated
+        } else {
+            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> {
+        let node = node_ref.as_deref_mut()?;
+
+        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 the new node (left_node) is already updated
+        } else {
+            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
         }
     }
 

basecoin/store/src/impls/in_memory.rs

@@ -142,8 +142,8 @@ impl Store for InMemoryStore {
         self.get_state(height).and_then(|v| v.get(path).cloned())
     }
 
-    fn delete(&mut self, _path: &Path) {
-        todo!()
+    fn delete(&mut self, path: &Path) {
+        self.pending.remove(path.clone());
     }
 
     fn commit(&mut self) -> Result<Vec<u8>, Self::Error> {