From aa056168457bb947320f44a30e1387957bfdf3ee Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Mon, 4 Aug 2025 13:01:08 -0300 Subject: [PATCH] fix tests --- core/storage/btree.rs | 158 +++++++++--------- .../query_processing/test_read_path.rs | 6 +- 2 files changed, 81 insertions(+), 83 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 265a8f091..08bab4d60 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -5398,95 +5398,93 @@ impl BTreeCursor { let mut mem_page; let mut contents; - loop { - let state = self.count_state; - match state { - CountState::Start => { - let _c = self.move_to_root()?; - self.count_state = CountState::Loop; - return Ok(IOResult::IO); + let state = self.count_state; + match state { + CountState::Start => { + let _c = self.move_to_root()?; + self.count_state = CountState::Loop; + 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 - ** this page contains countable entries. Increment the entry counter - ** accordingly. - */ - if !matches!(contents.page_type(), PageType::TableInterior) { - self.count += contents.cell_count(); + self.stack.advance(); + let cell_idx = self.stack.current_cell_index() as usize; + + // Second condition is necessary in case we return if the page is locked in the loop below + if contents.is_leaf() || cell_idx > 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 - if contents.is_leaf() || cell_idx > 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(); - 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!(), + 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!(), } } - CountState::Finish => { - return Ok(IOResult::Done(self.count)); - } + } + CountState::Finish => { + return Ok(IOResult::Done(self.count)); } } } diff --git a/tests/integration/query_processing/test_read_path.rs b/tests/integration/query_processing/test_read_path.rs index 193396362..03d8b9ceb 100644 --- a/tests/integration/query_processing/test_read_path.rs +++ b/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);