From 1585d5cbeedf5d6757030dcdbd31cdb158be5900 Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Sat, 2 Aug 2025 01:22:00 -0300
Subject: [PATCH] state machine for 'next' and `prev`

---
 core/storage/btree.rs          | 47 ++++++++++++++++++++++++++--------
 core/storage/state_machines.rs |  6 +++++
 2 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/core/storage/btree.rs b/core/storage/btree.rs
index 54fff8b3b..9a39669cf 100644
--- a/core/storage/btree.rs
+++ b/core/storage/btree.rs
@@ -10,7 +10,9 @@ 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, RewindState, SeekToLastState},
+        state_machines::{
+            AdvanceState, EmptyTableState, MoveToRightState, RewindState, SeekToLastState,
+        },
     },
     translate::plan::IterationDirection,
     turso_assert,
@@ -580,6 +582,8 @@ pub struct BTreeCursor {
     seek_to_last_state: SeekToLastState,
     /// State machine for [BTreeCursor::rewind]
     rewind_state: RewindState,
+    /// State machine for [BTreeCursor::next] and [BTreeCursor::prev]
+    advance_state: AdvanceState,
 }
 
 /// We store the cell index and cell count for each page in the stack.
@@ -638,6 +642,7 @@ impl BTreeCursor {
             move_to_right_state: (MoveToRightState::Start, None),
             seek_to_last_state: SeekToLastState::Start,
             rewind_state: RewindState::Start,
+            advance_state: AdvanceState::Start,
         }
     }
 
@@ -4279,11 +4284,21 @@ impl BTreeCursor {
 
     #[instrument(skip_all, level = Level::DEBUG)]
     pub fn next(&mut self) -> Result<IOResult<bool>> {
-        return_if_io!(self.restore_context());
-        let cursor_has_record = return_if_io!(self.get_next_record());
-        self.has_record.replace(cursor_has_record);
-        self.invalidate_record();
-        Ok(IOResult::Done(cursor_has_record))
+        loop {
+            match self.advance_state {
+                AdvanceState::Start => {
+                    return_if_io!(self.restore_context());
+                    self.advance_state = AdvanceState::Advance;
+                }
+                AdvanceState::Advance => {
+                    let cursor_has_record = return_if_io!(self.get_next_record());
+                    self.has_record.replace(cursor_has_record);
+                    self.invalidate_record();
+                    self.advance_state = AdvanceState::Start;
+                    return Ok(IOResult::Done(cursor_has_record));
+                }
+            }
+        }
     }
 
     fn invalidate_record(&mut self) {
@@ -4297,11 +4312,21 @@ impl BTreeCursor {
     #[instrument(skip_all, level = Level::DEBUG)]
     pub fn prev(&mut self) -> Result<IOResult<bool>> {
         assert!(self.mv_cursor.is_none());
-        return_if_io!(self.restore_context());
-        let cursor_has_record = return_if_io!(self.get_prev_record());
-        self.has_record.replace(cursor_has_record);
-        self.invalidate_record();
-        Ok(IOResult::Done(cursor_has_record))
+        loop {
+            match self.advance_state {
+                AdvanceState::Start => {
+                    return_if_io!(self.restore_context());
+                    self.advance_state = AdvanceState::Advance;
+                }
+                AdvanceState::Advance => {
+                    let cursor_has_record = return_if_io!(self.get_prev_record());
+                    self.has_record.replace(cursor_has_record);
+                    self.invalidate_record();
+                    self.advance_state = AdvanceState::Start;
+                    return Ok(IOResult::Done(cursor_has_record));
+                }
+            }
+        }
     }
 
     #[instrument(skip(self), level = Level::DEBUG)]
diff --git a/core/storage/state_machines.rs b/core/storage/state_machines.rs
index 53d83157d..1f437ebb1 100644
--- a/core/storage/state_machines.rs
+++ b/core/storage/state_machines.rs
@@ -23,3 +23,9 @@ pub enum RewindState {
     Start,
     NextRecord,
 }
+
+#[derive(Debug, Clone, Copy)]
+pub enum AdvanceState {
+    Start,
+    Advance,
+}