state machine for 'rewind'

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<usize>),
+    /// 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<IOResult<()>> {
-        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)]

core/storage/state_machines.rs

@@ -17,3 +17,9 @@ pub enum SeekToLastState {
     Start,
     IsEmpty,
 }
+
+#[derive(Debug, Clone, Copy)]
+pub enum RewindState {
+    Start,
+    NextRecord,
+}

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<Pager> },
-    CreateBtree { pager: Rc<Pager> },
+    StartingTxn {
+        pager: Rc<Pager>,
+    },
+    CreateBtree {
+        pager: Rc<Pager>,
+    },
+    // 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<BTreeCursor>,
+    },
 }
 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<Option<Cursor>>> =
                 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");