diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 87f2f5dc0..54fff8b3b 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -10,7 +10,7 @@ use crate::{ TableInteriorCell, TableLeafCell, CELL_PTR_SIZE_BYTES, INTERIOR_PAGE_HEADER_SIZE_BYTES, LEAF_PAGE_HEADER_SIZE_BYTES, LEFT_CHILD_PTR_SIZE_BYTES, }, - state_machines::{EmptyTableState, MoveToRightState, SeekToLastState}, + state_machines::{EmptyTableState, MoveToRightState, RewindState, SeekToLastState}, }, translate::plan::IterationDirection, turso_assert, @@ -576,7 +576,10 @@ pub struct BTreeCursor { /// State machine for [BTreeCursor::move_to_rightmost] and, optionally, the id of the rightmost page in the btree. /// If we know the rightmost page id and are already on that page, we can skip a seek. move_to_right_state: (MoveToRightState, Option), + /// State machine for [BTreeCursor::seek_to_last] seek_to_last_state: SeekToLastState, + /// State machine for [BTreeCursor::rewind] + rewind_state: RewindState, } /// We store the cell index and cell count for each page in the stack. @@ -634,6 +637,7 @@ impl BTreeCursor { is_empty_table_state: RefCell::new(EmptyTableState::Start), move_to_right_state: (MoveToRightState::Start, None), seek_to_last_state: SeekToLastState::Start, + rewind_state: RewindState::Start, } } @@ -4241,22 +4245,27 @@ impl BTreeCursor { #[instrument(skip_all, level = Level::DEBUG)] pub fn rewind(&mut self) -> Result> { - if let Some(mv_cursor) = &self.mv_cursor { - { - let mut mv_cursor = mv_cursor.borrow_mut(); - mv_cursor.rewind(); + loop { + match self.rewind_state { + RewindState::Start => { + self.rewind_state = RewindState::NextRecord; + if let Some(mv_cursor) = &self.mv_cursor { + let mut mv_cursor = mv_cursor.borrow_mut(); + mv_cursor.rewind(); + } else { + let _c = self.move_to_root()?; + return Ok(IOResult::IO); + } + } + RewindState::NextRecord => { + let cursor_has_record = return_if_io!(self.get_next_record()); + self.invalidate_record(); + self.has_record.replace(cursor_has_record); + self.rewind_state = RewindState::Start; + return Ok(IOResult::Done(())); + } } - let cursor_has_record = return_if_io!(self.get_next_record()); - self.invalidate_record(); - self.has_record.replace(cursor_has_record); - } else { - let _c = self.move_to_root()?; - - let cursor_has_record = return_if_io!(self.get_next_record()); - self.invalidate_record(); - self.has_record.replace(cursor_has_record); } - Ok(IOResult::Done(())) } #[instrument(skip_all, level = Level::DEBUG)] diff --git a/core/storage/state_machines.rs b/core/storage/state_machines.rs index ffce33498..53d83157d 100644 --- a/core/storage/state_machines.rs +++ b/core/storage/state_machines.rs @@ -17,3 +17,9 @@ pub enum SeekToLastState { Start, IsEmpty, } + +#[derive(Debug, Clone, Copy)] +pub enum RewindState { + Start, + NextRecord, +} diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 56bb30e23..d6cb09e99 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -6388,10 +6388,21 @@ pub fn op_noop( Ok(InsnFunctionStepResult::Step) } +#[derive(Default)] pub enum OpOpenEphemeralState { + #[default] Start, - StartingTxn { pager: Rc }, - CreateBtree { pager: Rc }, + StartingTxn { + pager: Rc, + }, + CreateBtree { + pager: Rc, + }, + // clippy complains this variant is too big when compared to the rest of the variants + // so it says we need to box it here + Rewind { + cursor: Box, + }, } pub fn op_open_ephemeral( program: &Program, @@ -6408,7 +6419,7 @@ pub fn op_open_ephemeral( Insn::OpenAutoindex { cursor_id } => (*cursor_id, false), _ => unreachable!("unexpected Insn {:?}", insn), }; - match &state.op_open_ephemeral_state { + match &mut state.op_open_ephemeral_state { OpOpenEphemeralState::Start => { tracing::trace!("Start"); let conn = program.connection.clone(); @@ -6491,7 +6502,7 @@ pub fn op_open_ephemeral( _ => unreachable!("This should not have happened"), }; - let mut cursor = if let CursorType::BTreeIndex(index) = cursor_type { + let cursor = if let CursorType::BTreeIndex(index) = cursor_type { BTreeCursor::new_index( mv_cursor, pager.clone(), @@ -6502,24 +6513,37 @@ pub fn op_open_ephemeral( } else { BTreeCursor::new_table(mv_cursor, pager.clone(), root_page as usize, num_columns) }; - let res = cursor.rewind()?; // Will never return io + state.op_open_ephemeral_state = OpOpenEphemeralState::Rewind { + cursor: Box::new(cursor), + }; + } + OpOpenEphemeralState::Rewind { cursor } => { + return_if_io!(cursor.rewind()); let mut cursors: std::cell::RefMut<'_, Vec>> = state.cursors.borrow_mut(); + let (_, cursor_type) = program.cursor_ref.get(cursor_id).unwrap(); + + let OpOpenEphemeralState::Rewind { cursor } = + std::mem::take(&mut state.op_open_ephemeral_state) + else { + unreachable!() + }; + // Table content is erased if the cursor already exists match cursor_type { CursorType::BTreeTable(_) => { cursors .get_mut(cursor_id) .unwrap() - .replace(Cursor::new_btree(cursor)); + .replace(Cursor::new_btree(*cursor)); } CursorType::BTreeIndex(_) => { cursors .get_mut(cursor_id) .unwrap() - .replace(Cursor::new_btree(cursor)); + .replace(Cursor::new_btree(*cursor)); } CursorType::Pseudo(_) => { panic!("OpenEphemeral on pseudo cursor");