opt: add binary search in branch updater

nomt/src/beatree/branch/node.rs

@@ -5,6 +5,7 @@ use bitvec::prelude::*;
 use super::BRANCH_NODE_SIZE;
 use crate::{
     beatree::{
+        allocator::PageNumber,
         ops::{bit_ops::bitwise_memcpy, get_key},
         Key,
     },
@@ -360,6 +361,11 @@ impl BranchNodeBuilder {
         }
     }
 
+    // returns the number of separtors already pushed
+    pub fn n_pushed(&self) -> usize {
+        self.index
+    }
+
     pub fn push(&mut self, key: Key, mut separator_len: usize, pn: u32) {
         assert!(self.index < self.branch.n() as usize);
 
@@ -390,8 +396,17 @@ impl BranchNodeBuilder {
     }
 
     // Copy the given chunk of separators from the provided base to the new node.
-    // Only compressed separators are expected in the specified range.
-    pub fn push_chunk(&mut self, base: &BranchNode, from: usize, to: usize) {
+    // Only compressed separators are expected in the specified range `from..to`.
+    //
+    // `updated` represents a list of new page numbers alongside the index
+    // of the separator to update within the specified range
+    pub fn push_chunk(
+        &mut self,
+        base: &BranchNode,
+        from: usize,
+        to: usize,
+        updated: impl Iterator<Item = (usize, PageNumber)>,
+    ) {
         let n_items = to - from;
         assert!(self.index + n_items <= self.branch.prefix_compressed() as usize);
 
@@ -439,6 +454,11 @@ impl BranchNodeBuilder {
         self.branch.node_pointers_mut()[self.index..self.index + n_items]
             .copy_from_slice(base_node_pointers);
 
+        // update page numbers of modified separators
+        for (i, new_pn) in updated {
+            self.branch.set_node_pointer(self.index + i, new_pn.0)
+        }
+
         // 3. copy and shift separators
         if bit_prefix_len_difference == 0 {
             // fast path, copy and shift all compressed separators at once
@@ -625,7 +645,7 @@ mod test {
             4, /*prefix_compressed*/
             9, /*prefix_len*/
         );
-        builder.push_chunk(&base_branch, 1, 4);
+        builder.push_chunk(&base_branch, 1, 4, [].into_iter());
         let mut key255 = [0; 32];
         key255[0] = 0x11;
         key255[1] = 255;
@@ -648,7 +668,7 @@ mod test {
             4, /*prefix_compressed*/
             7, /*prefix_len*/
         );
-        builder.push_chunk(&base_branch, 1, 4);
+        builder.push_chunk(&base_branch, 1, 4, [].into_iter());
         builder.push(key255, separator_len(&key255), 10);
         // prefix: 0001000
         // key 1:         1_10000000_00000001 // cell_poiter: 17

nomt/src/beatree/ops/mod.rs

@@ -50,38 +50,50 @@ pub fn lookup(key: Key, bbn_index: &Index, leaf_store: &StoreReader) -> Result<O
 /// Binary search a branch node for the child node containing the key. This returns the last child
 /// node pointer whose separator is less than or equal to the given key.
 fn search_branch(branch: &BranchNode, key: Key) -> Option<(usize, PageNumber)> {
+    let (found, pos) = find_key_pos(branch, &key, None);
+
+    if found {
+        return Some((pos, branch.node_pointer(pos).into()));
+    } else if pos == 0 {
+        return None;
+    } else {
+        // first key greater than the one we are looking for has been returned,
+        // thus the correct child is the previous one
+        return Some((pos - 1, branch.node_pointer(pos - 1).into()));
+    }
+}
+
+// Binary search for a key within a branch node.
+// Accept a field to override the starting point of the binary search.
+// It returns true and the index of the specified key,
+// or false and the index containing the first key greater than the specified one.
+pub fn find_key_pos(branch: &BranchNode, key: &Key, low: Option<usize>) -> (bool, usize) {
     let prefix = branch.prefix();
     let n = branch.n() as usize;
     let prefix_compressed = branch.prefix_compressed() as usize;
 
     match key.view_bits::<Msb0>()[..prefix.len()].cmp(prefix) {
-        Ordering::Less => return None,
-        Ordering::Greater if n == prefix_compressed => {
-            let i = n - 1;
-            return Some((i, branch.node_pointer(i).into()));
-        }
+        Ordering::Less => return (false, 0),
+        Ordering::Greater if n == prefix_compressed => return (false, n),
         Ordering::Equal | Ordering::Greater => {}
     }
 
-    let mut low = 0;
+    let mut low = low.unwrap_or(0);
     let mut high = branch.n() as usize;
 
     while low < high {
         let mid = low + (high - low) / 2;
 
-        if key < get_key(branch, mid) {
-            high = mid;
-        } else {
-            low = mid + 1;
+        match key.cmp(&get_key(branch, mid)) {
+            Ordering::Equal => {
+                return (true, mid);
+            }
+            Ordering::Less => high = mid,
+            Ordering::Greater => low = mid + 1,
         }
     }
 
-    // sanity: this only happens if `key` is less than separator 0.
-    if high == 0 {
-        return None;
-    }
-    let node_pointer = branch.node_pointer(high - 1);
-    Some((high - 1, node_pointer.into()))
+    (false, high)
 }
 
 // Extract the key at a given index from a BranchNode, taking into account prefix compression.

nomt/src/beatree/ops/update/branch_stage.rs

@@ -231,7 +231,7 @@ fn reset_branch_base(
     }
 
     if let Some((_, node, next_separator)) = branches_tracker.pending_base.take() {
-        let base = BaseBranch { node, iter_pos: 0 };
+        let base = BaseBranch::new(node);
         branch_updater.reset_base(Some(base), next_separator);
     } else {
         if let Some(separator) = branches_tracker
@@ -266,10 +266,7 @@ fn reset_branch_base_fresh(
 
     branches_tracker.delete(separator, branch.bbn_pn().into(), cutoff);
 
-    let base = BaseBranch {
-        node: branch,
-        iter_pos: 0,
-    };
+    let base = BaseBranch::new(branch);
     branch_updater.reset_base(Some(base), cutoff);
 }