fix tests

core/storage/btree.rs

@@ -5398,95 +5398,93 @@ impl BTreeCursor {
         let mut mem_page;
         let mut contents;
 
-        loop {
+        let state = self.count_state;
-            let state = self.count_state;
+        match state {
-            match state {
+            CountState::Start => {
-                CountState::Start => {
+                let _c = self.move_to_root()?;
-                    let _c = self.move_to_root()?;
+                self.count_state = CountState::Loop;
-                    self.count_state = CountState::Loop;
+                return Ok(IOResult::IO);
-                    return Ok(IOResult::IO);
+            }
+            CountState::Loop => {
+                mem_page_rc = self.stack.top();
+                mem_page = mem_page_rc.get();
+                return_if_locked_maybe_load!(self.pager, mem_page_rc);
+                turso_assert!(mem_page.is_loaded(), "page should be loaded");
+                contents = mem_page.get().contents.as_ref().unwrap();
+
+                /* If this is a leaf page or the tree is not an int-key tree, then
+                 ** this page contains countable entries. Increment the entry counter
+                 ** accordingly.
+                 */
+                if !matches!(contents.page_type(), PageType::TableInterior) {
+                    self.count += contents.cell_count();
                 }
-                CountState::Loop => {
-                    mem_page_rc = self.stack.top();
-                    mem_page = mem_page_rc.get();
-                    turso_assert!(mem_page.is_loaded(), "page should be loaded");
-                    contents = mem_page.get().contents.as_ref().unwrap();
 
-                    /* If this is a leaf page or the tree is not an int-key tree, then
+                self.stack.advance();
-                     ** this page contains countable entries. Increment the entry counter
+                let cell_idx = self.stack.current_cell_index() as usize;
-                     ** accordingly.
+
-                     */
+                // Second condition is necessary in case we return if the page is locked in the loop below
-                    if !matches!(contents.page_type(), PageType::TableInterior) {
+                if contents.is_leaf() || cell_idx > contents.cell_count() {
-                        self.count += contents.cell_count();
+                    loop {
+                        if !self.stack.has_parent() {
+                            // All pages of the b-tree have been visited. Return successfully
+                            let _c = self.move_to_root()?;
+                            self.count_state = CountState::Finish;
+                            return Ok(IOResult::IO);
+                        }
+
+                        // Move to parent
+                        self.stack.pop();
+
+                        mem_page_rc = self.stack.top();
+                        mem_page = mem_page_rc.get();
+                        return_if_locked_maybe_load!(self.pager, mem_page_rc);
+                        turso_assert!(mem_page.is_loaded(), "page should be loaded");
+                        contents = mem_page.get().contents.as_ref().unwrap();
+
+                        let cell_idx = self.stack.current_cell_index() as usize;
+
+                        if cell_idx <= contents.cell_count() {
+                            break;
+                        }
                     }
+                }
 
+                let cell_idx = self.stack.current_cell_index() as usize;
+
+                assert!(cell_idx <= contents.cell_count(),);
+                assert!(!contents.is_leaf());
+
+                if cell_idx == contents.cell_count() {
+                    // Move to right child
+                    // should be safe as contents is not a leaf page
+                    let right_most_pointer = contents.rightmost_pointer().unwrap();
                     self.stack.advance();
-                    let cell_idx = self.stack.current_cell_index() as usize;
+                    let (mem_page, _c) = self.read_page(right_most_pointer as usize)?;
+                    self.stack.push(mem_page);
+                    return Ok(IOResult::IO);
+                } else {
+                    // Move to child left page
+                    let cell = contents.cell_get(cell_idx, self.usable_space())?;
 
-                    // Second condition is necessary in case we return if the page is locked in the loop below
+                    match cell {
-                    if contents.is_leaf() || cell_idx > contents.cell_count() {
+                        BTreeCell::TableInteriorCell(TableInteriorCell {
-                        loop {
+                            left_child_page, ..
-                            if !self.stack.has_parent() {
+                        })
-                                // All pages of the b-tree have been visited. Return successfully
+                        | BTreeCell::IndexInteriorCell(IndexInteriorCell {
-                                let _c = self.move_to_root()?;
+                            left_child_page, ..
-                                self.count_state = CountState::Finish;
+                        }) => {
-                                return Ok(IOResult::IO);
+                            self.stack.advance();
-                            }
+                            let (mem_page, _c) = self.read_page(left_child_page as usize)?;
-
+                            self.stack.push(mem_page);
-                            // Move to parent
+                            return Ok(IOResult::IO);
-                            self.stack.pop();
-
-                            mem_page_rc = self.stack.top();
-                            mem_page = mem_page_rc.get();
-                            turso_assert!(mem_page.is_loaded(), "page should be loaded");
-                            contents = mem_page.get().contents.as_ref().unwrap();
-
-                            let cell_idx = self.stack.current_cell_index() as usize;
-
-                            if cell_idx <= contents.cell_count() {
-                                break;
-                            }
-                        }
-                    }
-
-                    let cell_idx = self.stack.current_cell_index() as usize;
-
-                    assert!(cell_idx <= contents.cell_count(),);
-                    assert!(!contents.is_leaf());
-
-                    if cell_idx == contents.cell_count() {
-                        // Move to right child
-                        // should be safe as contents is not a leaf page
-                        let right_most_pointer = contents.rightmost_pointer().unwrap();
-                        self.stack.advance();
-                        let (mem_page, _c) = self.read_page(right_most_pointer as usize)?;
-                        self.stack.push(mem_page);
-                        return Ok(IOResult::IO);
-                    } else {
-                        // Move to child left page
-                        let cell = contents.cell_get(cell_idx, self.usable_space())?;
-
-                        match cell {
-                            BTreeCell::TableInteriorCell(TableInteriorCell {
-                                left_child_page,
-                                ..
-                            })
-                            | BTreeCell::IndexInteriorCell(IndexInteriorCell {
-                                left_child_page,
-                                ..
-                            }) => {
-                                self.stack.advance();
-                                let (mem_page, _c) = self.read_page(left_child_page as usize)?;
-                                self.stack.push(mem_page);
-                                return Ok(IOResult::IO);
-                            }
-                            _ => unreachable!(),
                         }
+                        _ => unreachable!(),
                     }
                 }
-                CountState::Finish => {
+            }
-                    return Ok(IOResult::Done(self.count));
+            CountState::Finish => {
-                }
+                return Ok(IOResult::Done(self.count));
             }
         }
     }

tests/integration/query_processing/test_read_path.rs

@@ -617,12 +617,12 @@ fn test_bind_parameters_update_rowid_alias_seek_rowid() -> anyhow::Result<()> {
                     row.get::<&Value>(2).unwrap(),
                     &Value::Integer(if i == 0 { 4 } else { 11 })
                 );
+                i += 1;
             }
             StepResult::IO => sel.run_once()?,
             StepResult::Done | StepResult::Interrupt => break,
             StepResult::Busy => panic!("database busy"),
         }
-        i += 1;
     }
     let mut ins = conn.prepare("update test set name = ? where id < ? AND age between ? and ?;")?;
     ins.bind_at(1.try_into()?, Value::build_text("updated"));
@@ -648,12 +648,12 @@ fn test_bind_parameters_update_rowid_alias_seek_rowid() -> anyhow::Result<()> {
                     row.get::<&Value>(0).unwrap(),
                     &Value::build_text(if i == 0 { "updated" } else { "test" }),
                 );
+                i += 1;
             }
             StepResult::IO => sel.run_once()?,
             StepResult::Done | StepResult::Interrupt => break,
             StepResult::Busy => panic!("database busy"),
         }
-        i += 1;
     }
 
     assert_eq!(ins.parameters().count(), 4);
@@ -692,12 +692,12 @@ fn test_bind_parameters_delete_rowid_alias_seek_out_of_order() -> anyhow::Result
             StepResult::Row => {
                 let row = sel.row().unwrap();
                 assert_eq!(row.get::<&Value>(0).unwrap(), &Value::build_text("correct"),);
+                i += 1;
             }
             StepResult::IO => sel.run_once()?,
             StepResult::Done | StepResult::Interrupt => break,
             StepResult::Busy => panic!("database busy"),
         }
-        i += 1;
     }
     assert_eq!(i, 1);
     assert_eq!(ins.parameters().count(), 4);