From 6b7575bf3f2b01852ddae77dcce034e1d39df01a Mon Sep 17 00:00:00 2001
From: Pere Diaz Bou <pere-altea@hotmail.com>
Date: Wed, 9 Apr 2025 15:03:23 +0200
Subject: [PATCH] fix tree traversal assumptions on traversal

---
 core/storage/btree.rs | 46 ++++++++++++++++++++++++++++++++++++-------
 1 file changed, 39 insertions(+), 7 deletions(-)

diff --git a/core/storage/btree.rs b/core/storage/btree.rs
index 9e9d6800c..95dcad520 100644
--- a/core/storage/btree.rs
+++ b/core/storage/btree.rs
@@ -1630,7 +1630,12 @@ impl BTreeCursor {
                     let write_info = self.state.mut_write_info().unwrap();
                     write_info.state = WriteState::BalanceNonRoot;
                     self.stack.pop();
-                    self.stack.retreat();
+                    // with `move_to` we advance the current cell idx of TableInterior once we move to left subtree.
+                    // On the other hand, with IndexInterior, we do not because we tranver in-order. In the latter case
+                    // since we haven't consumed the cell we can avoid retreating the current cell index.
+                    if matches!(current_page.get_contents().page_type(), PageType::TableLeaf) {
+                        self.stack.retreat();
+                    }
                     return_if_io!(self.balance_non_root());
                 }
                 WriteState::BalanceNonRoot | WriteState::BalanceNonRootWaitLoadPages => {
@@ -1682,7 +1687,12 @@ impl BTreeCursor {
                     parent_contents.overflow_cells.is_empty(),
                     "balancing child page with overflowed parent not yet implemented"
                 );
-                assert!(page_to_balance_idx <= parent_contents.cell_count());
+                assert!(
+                    page_to_balance_idx <= parent_contents.cell_count(),
+                    "page_to_balance_idx={} is out of bounds for parent cell count {}",
+                    page_to_balance_idx,
+                    number_of_cells_in_parent
+                );
                 // As there will be at maximum 3 pages used to balance:
                 // sibling_pointer is the index represeneting one of those 3 pages, and we initialize it to the last possible page.
                 // next_divider is the first divider that contains the first page of the 3 pages.
@@ -1813,6 +1823,7 @@ impl BTreeCursor {
                 // Now do real balancing
                 let parent_page = self.stack.top();
                 let parent_contents = parent_page.get_contents();
+
                 assert!(
                     parent_contents.overflow_cells.is_empty(),
                     "overflow parent not yet implemented"
@@ -2259,6 +2270,7 @@ impl BTreeCursor {
                         write_varint_to_vec(rowid, &mut new_divider_cell);
                     } else {
                         // Leaf index
+                        new_divider_cell.extend_from_slice(&(page.get().id as u32).to_be_bytes());
                         new_divider_cell.extend_from_slice(divider_cell);
                     }
 
@@ -2855,7 +2867,10 @@ impl BTreeCursor {
         child_buf[0..root_contents.header_size()]
             .copy_from_slice(&root_buf[offset..offset + root_contents.header_size()]);
         // Copy overflow cells
-        child_contents.overflow_cells = root_contents.overflow_cells.clone();
+        std::mem::swap(
+            &mut child_contents.overflow_cells,
+            &mut root_contents.overflow_cells,
+        );
 
         // 2. Modify root
         let new_root_page_type = match root_contents.page_type() {
@@ -2878,7 +2893,9 @@ impl BTreeCursor {
         self.root_page = root.get().id;
         self.stack.clear();
         self.stack.push(root.clone());
-        self.stack.advance();
+        if matches!(root_contents.page_type(), PageType::TableInterior) {
+            self.stack.advance();
+        }
         self.stack.push(child.clone());
     }
 
@@ -2933,6 +2950,7 @@ impl BTreeCursor {
             }
             cell_idx += 1;
         }
+        assert!(cell_idx <= cell_count);
         cell_idx
     }
 
@@ -4449,10 +4467,21 @@ fn debug_validate_cells_core(page: &PageContent, usable_space: u16) {
             payload_overflow_threshold_min(page.page_type(), usable_space),
             usable_space as usize,
         );
+        let buf = &page.as_ptr()[offset..offset + size];
+        assert!(
+            size >= 4,
+            "cell size should be at least 4 bytes idx={}, cell={:?}, offset={}",
+            i,
+            buf,
+            offset
+        );
         if page.is_leaf() {
             assert!(page.as_ptr()[offset] != 0);
         }
-        assert!(size >= 4, "cell size should be at least 4 bytes idx={}", i);
+        assert!(
+            offset + size <= usable_space as usize,
+            "cell spans out of usable space"
+        );
     }
 }
 
@@ -4467,6 +4496,7 @@ fn insert_into_cell(
     cell_idx: usize,
     usable_space: u16,
 ) -> Result<()> {
+    debug_validate_cells!(page, usable_space);
     assert!(
         cell_idx <= page.cell_count() + page.overflow_cells.len(),
         "attempting to add cell to an incorrect place cell_idx={} cell_count={}",
@@ -4487,10 +4517,12 @@ fn insert_into_cell(
 
     let new_cell_data_pointer = allocate_cell_space(page, payload.len() as u16, usable_space)?;
     tracing::debug!(
-        "insert_into_cell(idx={}, pc={})",
+        "insert_into_cell(idx={}, pc={}, size={})",
         cell_idx,
-        new_cell_data_pointer
+        new_cell_data_pointer,
+        payload.len()
     );
+    assert!(new_cell_data_pointer + payload.len() as u16 <= usable_space);
     let buf = page.as_ptr();
 
     // copy data