add the ability to remove keys

src/lib.rs

@@ -2,6 +2,7 @@
 
 extern crate alloc;
 
+use alloc::collections::VecDeque;
 use alloc::vec;
 use alloc::vec::Vec;
 
@@ -150,6 +151,18 @@ impl<T: SizeOf + Clone + Default + ToBytes + FromBytes> BTree<T> {
         }
     }
 
+    /// deletes a node from the internal node list, and updates all references to the following nodes so that they point to the correct node
+    pub fn delete_from_internal_list(&mut self, n: usize) {
+        for node in &mut self.nodes {
+            for child in &mut node.children {
+                if *child >= n && *child > 0 {
+                    *child -= 1;
+                }
+            }
+        }
+        self.nodes.remove(n);
+    }
+
     /// finds the node with the given key, or None if not found
     pub fn search(&mut self, key: u32) -> Option<&mut BTreeNode<T>> {
         let mut index = self.root?;
@@ -277,6 +290,291 @@ impl<T: SizeOf + Clone + Default + ToBytes + FromBytes> BTree<T> {
         }
     }
 
+    /// removes the given key from this b-tree (THIS IS THE FUNCTION YOU LIKELY WANT)
+    pub fn remove(&mut self, key: u32) {
+        if let Some(mut root) = self.root {
+            let mut res = self.remove_from_nodetree(root, key);
+            let mut root_removed = false;
+            while let Some(node) = res.pop_front() {
+                if root >= node {
+                    root -= 1;
+                }
+                if node == root {
+                    root_removed = true;
+                }
+            }
+            self.root = Some(root);
+            // if root has no keys, make the first child the new root (if it has a child), else set root to None
+            if self.nodes[root].n == 0 || root_removed {
+                if self.nodes[root].leaf {
+                    self.delete_from_internal_list(root);
+                    self.root = None;
+                } else {
+                    self.root = Some(self.nodes[root].children[0]);
+                    self.delete_from_internal_list(root);
+                }
+            }
+        }
+    }
+
+    /// removes the given key from the tree starting from a given node, returns VecDeque<usize> where Vec<usize> is a vector of nodes that were removed (so that other references can be updated)
+    pub fn remove_from_nodetree(&mut self, mut node: usize, key: u32) -> VecDeque<usize> {
+        let index = self.find_gte_key_on_node(node, key);
+
+        let mut result = VecDeque::new();
+
+        if index < self.nodes[node].n as usize && self.nodes[node].keys[index].0 == key {
+            // if this is a leaf node, call remove_assuming_leaf
+            // vice versa
+            if self.nodes[node].leaf {
+                self.remove_assuming_leaf(node, index);
+            } else {
+                let res = self.remove_assuming_non_leaf(node, index);
+                result.append(&mut res.clone());
+            }
+
+            result
+        } else {
+            // if this is a leaf node, key not found
+            if self.nodes[node].leaf {
+                return result;
+            }
+
+            // saved for later
+            let flag = index == self.nodes[node].n as usize;
+
+            // if the child that should contain the key has less than m keys, steal from it
+            if self.nodes[self.nodes[node].children[index]].n < self.m {
+                let shift = self.steal_give(node, index);
+                if let Some(shift) = shift {
+                    // if node is greater than shift, decrease them by 1
+                    if node >= shift {
+                        node -= 1;
+                    }
+                    result.push_front(shift);
+                }
+            }
+
+            // if the last child was merged, recurse on the (index-1)th child
+            // vice versa
+            if flag && index > self.nodes[node].n as usize {
+                let res = self.remove_from_nodetree(self.nodes[node].children[index - 1], key);
+                result.append(&mut res.clone());
+            } else {
+                let res = self.remove_from_nodetree(self.nodes[node].children[index], key);
+                result.append(&mut res.clone());
+            }
+
+            result
+        }
+    }
+
+    /// finds the first key on a node that is greater than or equal to the given key
+    pub fn find_gte_key_on_node(&mut self, node: usize, key: u32) -> usize {
+        let mut i = 0;
+        while i < self.nodes[node].n && self.nodes[node].keys[i as usize].0 < key {
+            i += 1;
+        }
+        i as usize
+    }
+
+    /// steals keys from a node's children to fill child n, returns the position of the node that was removed so that other references can be updated, or None if no node was removed
+    pub fn steal_give(&mut self, node: usize, n: usize) -> Option<usize> {
+        // if the previous child has more than m-1 keys, steal from it
+        // else if the next child has more than m-1 keys, steal from it
+        // else, merge with one of the children
+        if n != 0 && self.nodes[self.nodes[node].children[n-1]].n >= self.m {
+            self.cfptn(node, n);
+            None
+        } else if n != self.nodes[node].n as usize && self.nodes[self.nodes[node].children[n+1]].n >= self.m {
+            self.cfftn(node, n);
+            None
+        } else if n != self.nodes[node].n as usize {
+            Some(self.merge_children(node, n).unwrap())
+        } else {
+            Some(self.merge_children(node, n-1).unwrap())
+        }
+    }
+
+    /// copies the last key from the (n+1)th child and gives it to the nth child
+    pub fn cfftn(&mut self, node: usize, n: usize) {
+        let child = self.nodes[node].children[n];
+        let sibling = self.nodes[node].children[n+1];
+
+        // give child the nth key from the current node
+        let child_n = self.nodes[child].n;
+        self.nodes[child].keys[child_n as usize] = self.nodes[node].keys[n].clone();
+
+        // if child isn't leaf, move sibling's first child as child's last child
+        if !self.nodes[child].leaf {
+            self.nodes[child].children[child_n as usize + 1] = self.nodes[sibling].children[0];
+        }
+
+        // move the first key from sibling to the current node
+        self.nodes[node].keys[n] = self.nodes[sibling].keys[0].clone();
+
+        // shift sibling keys to the left
+        for i in 1..self.nodes[sibling].n {
+            self.nodes[sibling].keys[i as usize - 1] = self.nodes[sibling].keys[i as usize].clone();
+        }
+
+        // if sibling isn't leaf, shift sibling children to the left
+        if !self.nodes[sibling].leaf {
+            for i in 1..self.nodes[sibling].n+1 {
+                self.nodes[sibling].children[i as usize - 1] = self.nodes[sibling].children[i as usize];
+            }
+        }
+
+        // decrease number of keys in sibling
+        self.nodes[sibling].n -= 1;
+        // increase number of keys in child
+        self.nodes[child].n += 1;
+    }
+
+    /// copies the last key from the (n-1)th child and gives it to the nth child
+    pub fn cfptn(&mut self, node: usize, n: usize) {
+        let child = self.nodes[node].children[n];
+        let sibling = self.nodes[node].children[n-1];
+
+        // move all keys in child one to the right
+        for i in (0..=self.nodes[child].n-1).rev() {
+            self.nodes[child].keys[i as usize + 1] = self.nodes[child].keys[i as usize].clone();
+        }
+
+        // if child isn't leaf, move all child pointers one step ahead
+        if !self.nodes[child].leaf {
+            for i in (0..=self.nodes[child].n).rev() {
+                self.nodes[child].children[i as usize + 1] = self.nodes[child].children[i as usize];
+            }
+        }
+
+        // give child the (n-1)th key from the current node
+        self.nodes[child].keys[0] = self.nodes[node].keys[n-1].clone();
+
+        // if child isn't leaf, move sibling's last child as child's first child
+        if !self.nodes[child].leaf {
+            self.nodes[child].children[0] = self.nodes[sibling].children[self.nodes[sibling].n as usize];
+        }
+
+        // move the (n-1)th key from sibling to the current node
+        self.nodes[node].keys[n-1] = self.nodes[sibling].keys[self.nodes[sibling].n as usize - 1].clone();
+
+        // decrease number of keys in sibling
+        self.nodes[sibling].n -= 1;
+        // increase number of keys in child
+        self.nodes[child].n += 1;
+    }
+
+    /// merges the given node's nth child with its (n+1)th child, returns Some(n) if the merge was successful, None if the merge was not successful
+    /// where `n` is the index of the node that was removed so that other references can be updated
+    pub fn merge_children(&mut self, node: usize, n: usize) -> Option<usize> {
+        let child = self.nodes[node].children[n];
+        let sibling = self.nodes[node].children[n+1];
+
+        // pop key from current node and insert it into child
+        self.nodes[child].keys[self.m as usize - 1] = self.nodes[node].keys[n].clone();
+
+        // copy keys from sibling into child
+        for i in 0..self.nodes[sibling].n {
+            self.nodes[child].keys[(self.m + i) as usize] = self.nodes[sibling].keys[i as usize].clone();
+        }
+
+        // if child isn't leaf, copy children from sibling into child
+        if !self.nodes[child].leaf {
+            for i in 0..self.nodes[sibling].n {
+                self.nodes[child].children[(self.m + i) as usize] = self.nodes[sibling].children[i as usize];
+            }
+        }
+
+        // move keys after n one to the left, filling the gap
+        for i in n+1..self.nodes[node].n as usize {
+            self.nodes[node].keys[i-1] = self.nodes[node].keys[i].clone();
+        }
+
+        // move children after n+1 one to the left, filling the gap
+        for i in n+2..self.nodes[node].n as usize + 1 {
+            self.nodes[node].children[i-1] = self.nodes[node].children[i];
+        }
+        // decrease number of keys in node
+        self.nodes[child].n += self.nodes[sibling].n + 1;
+        self.nodes[node].n -= 1;
+
+        // delete sibling
+        self.delete_from_internal_list(sibling);
+
+        Some(sibling)
+    }
+
+    /// removes the nth key from the given node assuming it is a non-leaf node, returns VecDeque<usize> where Vec<usize> is a vector of nodes that were removed (so that other references can be updated)
+    pub fn remove_assuming_non_leaf(&mut self, mut node: usize, n: usize) -> VecDeque<usize> {
+        let key = self.nodes[node].keys[n].0;
+
+        let mut result = VecDeque::new();
+
+        let child_n = self.nodes[self.nodes[node].children[n]].n;
+
+        // if child before k (children[n]) has at least m keys, find predecessor, replace key with predecessor, delete predecessor
+        // else if nth child has less than m keys, and (n+1)th child has at least m keys, replace key with successor, delete successor
+        // else merge nth child with (n+1)th child, recurse on nth child
+        if child_n >= self.m {
+            let pre_node = self.fkfl(node, n).unwrap().clone();
+            self.nodes[node].keys[n] = pre_node.clone();
+            let res = self.remove_from_nodetree(self.nodes[node].children[n], pre_node.0);
+            result.append(&mut res.clone());
+        } else if self.nodes[self.nodes[node].children[n+1]].n >= self.m {
+            let suc_node = self.fkfr(node, n).unwrap().clone();
+            self.nodes[node].keys[n] = suc_node.clone();
+            let res = self.remove_from_nodetree(self.nodes[node].children[n+1], suc_node.0);
+            result.append(&mut res.clone());
+        } else {
+            let res = self.merge_children(node, n).unwrap();
+            if node >= res {
+                node -= 1;
+                result.push_front(res);
+            }
+            let res = self.remove_from_nodetree(self.nodes[node].children[n], key);
+            result.append(&mut res.clone());
+        }
+
+        result
+    }
+
+    /// removes the nth key from the given node assuming it is a leaf node
+    pub fn remove_assuming_leaf(&mut self, node: usize, n: usize) {
+        // move all keys after n one to the left
+        for i in n+1..self.nodes[node].n as usize {
+            self.nodes[node].keys[i-1] = self.nodes[node].keys[i].clone();
+        }
+        self.nodes[node].n -= 1;
+    }
+
+    /// returns the last key of the last leaf node starting from the right of the children of the given node
+    pub fn fkfr(&self, node: usize, i: usize) -> Option<&(u32, T)> {
+        let children = &self.nodes[node].children;
+        if children.is_empty() {
+            return None;
+        }
+        let mut node = children[i];
+        while !self.nodes[node].leaf {
+            node = self.nodes[node].children[self.nodes[node].n as usize];
+        }
+        Some(&self.nodes[node].keys[self.nodes[node].n as usize - 1])
+    }
+
+    /// returns the first key of the first leaf node starting from the left of the children of the given node
+    pub fn fkfl(&self, node: usize, i: usize) -> Option<&(u32, T)> {
+        let children = &self.nodes[node].children;
+        if children.is_empty() {
+            return None;
+        }
+        let mut node = children[i];
+        while !self.nodes[node].leaf {
+            node = self.nodes[node].children[0];
+        }
+        Some(&self.nodes[node].keys[0])
+
+    }
+
     /// traverses the tree in order and returns a vector of all keys
     pub fn traverse_in_order(&self, node: usize) -> Vec<u32> {
         let mut result = vec![];
@@ -293,6 +591,23 @@ impl<T: SizeOf + Clone + Default + ToBytes + FromBytes> BTree<T> {
 
         result
     }
+
+    /// traverses the tree in order and returns a vector of all values {
+    pub fn traverse_in_order_values(&self, node: usize) -> Vec<T> {
+        let mut result = vec![];
+        for i in 0..self.nodes[node].n {
+            if !self.nodes[node].leaf {
+                result.append(&mut self.traverse_in_order_values(self.nodes[node].children[i as usize]));
+            }
+            result.push(self.nodes[node].keys[i as usize].1.clone());
+        }
+
+        if !self.nodes[node].leaf {
+            result.append(&mut self.traverse_in_order_values(self.nodes[node].children[self.nodes[node].n as usize]));
+        }
+
+        result
+    }
 }
 
 impl<T: SizeOf + Clone + Default + ToBytes + FromBytes> ToBytes for BTree<T> {
@@ -368,6 +683,11 @@ mod tests {
         tree.insert(3, 0);
         tree.insert(2, 0);
         tree.insert(1, 0);
+        tree.insert(24, 1);
+        tree.insert(25, 0);
+
+        tree.remove(25);
+        tree.remove(24);
 
         fn tree_good(mut tree: BTree<i32>) {
             // print traversed tree
@@ -385,5 +705,26 @@ mod tests {
         let tree2 = BTree::<i32>::from_bytes(&bytes);
 
         tree_good(tree2.clone());
+
+        // remove 10 and 20, then check if tree is still good
+        tree.remove(10);
+        assert_eq!(tree.traverse_in_order(tree.root.unwrap()), vec![1, 2, 3, 5, 6, 7, 11, 12, 17, 20, 30]);
+        tree.remove(20);
+        assert_eq!(tree.traverse_in_order(tree.root.unwrap()), vec![1, 2, 3, 5, 6, 7, 11, 12, 17, 30]);
+
+        // remove all keys but 30
+        let keys_remaining = vec![1, 2, 3, 5, 6, 7, 11, 12, 17];
+        for key in keys_remaining {
+            tree.remove(key);
+        }
+
+        // check if tree is still good
+        assert_eq!(tree.traverse_in_order(tree.root.unwrap()), vec![30]);
+
+        // remove 30
+        tree.remove(30);
+
+        assert_eq!(tree.root, None);
+        assert_eq!(tree.nodes.len(), 0);
     }
 }