diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 4b05b1754..17fd94d2a 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -473,6 +473,8 @@ impl CursorSeekState { } } +type CursorMoveToState = CursorSeekState; + pub struct BTreeCursor { /// The multi-version cursor that is used to read and write to the database file. mv_cursor: Option>>, @@ -509,6 +511,7 @@ pub struct BTreeCursor { /// This Vec should be empty for Table Btree collations: Vec, seek_state: CursorSeekState, + move_to_state: CursorMoveToState, } impl BTreeCursor { @@ -539,6 +542,7 @@ impl BTreeCursor { valid_state: CursorValidState::Valid, collations, seek_state: CursorSeekState::Start, + move_to_state: CursorMoveToState::Start, } } @@ -1542,12 +1546,24 @@ impl BTreeCursor { } let cell_count = contents.cell_count(); - let mut min: isize = 0; - let mut max: isize = cell_count as isize - 1; - let mut leftmost_matching_cell = None; + { + let min: isize = 0; + let max: isize = cell_count as isize - 1; + let leftmost_matching_cell = None; + + self.move_to_state = CursorMoveToState::Seeking { + min, + max, + nearest_matching_cell: leftmost_matching_cell, + }; + } loop { + let min = self.move_to_state.get_min(); + let max = self.move_to_state.get_max(); if min > max { - if let Some(leftmost_matching_cell) = leftmost_matching_cell { + if let Some(leftmost_matching_cell) = + self.move_to_state.get_nearest_matching_cell() + { let left_child_page = contents.cell_table_interior_read_left_child_page( leftmost_matching_cell as usize, )?; @@ -1612,10 +1628,11 @@ impl BTreeCursor { SeekOp::EQ => cell_rowid >= rowid, }; if is_on_left { - leftmost_matching_cell = Some(cur_cell_idx as usize); - max = cur_cell_idx - 1; + self.move_to_state + .set_nearest_matching_cell(Some(cur_cell_idx as usize)); + self.move_to_state.set_max(cur_cell_idx - 1); } else { - min = cur_cell_idx + 1; + self.move_to_state.set_min(cur_cell_idx + 1); } } } @@ -1637,13 +1654,26 @@ impl BTreeCursor { return Ok(CursorResult::Ok(())); } - let cell_count = contents.cell_count(); - let mut min: isize = 0; - let mut max: isize = cell_count as isize - 1; - let mut leftmost_matching_cell = None; + { + let cell_count = contents.cell_count(); + let min: isize = 0; + let max: isize = cell_count as isize - 1; + let leftmost_matching_cell = None; + + self.move_to_state = CursorMoveToState::Seeking { + min, + max, + nearest_matching_cell: leftmost_matching_cell, + }; + } + loop { + let min = self.move_to_state.get_min(); + let max = self.move_to_state.get_max(); if min > max { - let Some(leftmost_matching_cell) = leftmost_matching_cell else { + let Some(leftmost_matching_cell) = + self.move_to_state.get_nearest_matching_cell() + else { self.stack.set_cell_index(contents.cell_count() as i32 + 1); match contents.rightmost_pointer() { Some(right_most_pointer) => { @@ -1724,59 +1754,62 @@ impl BTreeCursor { self.get_immutable_record_or_create().as_mut().unwrap(), )? }; - let record = self.get_immutable_record(); - let record = record.as_ref().unwrap(); - let record_slice_equal_number_of_cols = - &record.get_values().as_slice()[..index_key.get_values().len()]; - let interior_cell_vs_index_key = compare_immutable( - record_slice_equal_number_of_cols, - index_key.get_values(), - self.key_sort_order(), - &self.collations, - ); - // in sqlite btrees left child pages have <= keys. - // in general, in forwards iteration we want to find the first key that matches the seek condition. - // in backwards iteration we want to find the last key that matches the seek condition. - // - // Logic table for determining if target leaf page is in left subtree. - // For index b-trees this is a bit more complicated since the interior cells contain payloads (the key is the payload). - // and for non-unique indexes there might be several cells with the same key. - // - // Forwards iteration (looking for first match in tree): - // OP | Current Cell vs Seek Key | Action? | Explanation - // GT | > | go left | First > key could be exactly this one, or in left subtree - // GT | = or < | go right | First > key must be in right subtree - // GE | > | go left | First >= key could be exactly this one, or in left subtree - // GE | = | go left | First >= key could be exactly this one, or in left subtree - // GE | < | go right | First >= key must be in right subtree - // - // Backwards iteration (looking for last match in tree): - // OP | Current Cell vs Seek Key | Action? | Explanation - // LE | > | go left | Last <= key must be in left subtree - // LE | = | go right | Last <= key is either this one, or somewhere to the right of this one. So we need to go right to make sure - // LE | < | go right | Last <= key must be in right subtree - // LT | > | go left | Last < key must be in left subtree - // LT | = | go left | Last < key must be in left subtree since we want strictly less than - // LT | < | go right | Last < key could be exactly this one, or in right subtree - // - // No iteration (point query): - // EQ | > | go left | First = key must be in left subtree - // EQ | = | go left | First = key could be exactly this one, or in left subtree - // EQ | < | go right | First = key must be in right subtree + let target_leaf_page_is_in_left_subtree = { + let record = self.get_immutable_record(); + let record = record.as_ref().unwrap(); + let record_slice_equal_number_of_cols = + &record.get_values().as_slice()[..index_key.get_values().len()]; + let interior_cell_vs_index_key = compare_immutable( + record_slice_equal_number_of_cols, + index_key.get_values(), + self.key_sort_order(), + &self.collations, + ); + // in sqlite btrees left child pages have <= keys. + // in general, in forwards iteration we want to find the first key that matches the seek condition. + // in backwards iteration we want to find the last key that matches the seek condition. + // + // Logic table for determining if target leaf page is in left subtree. + // For index b-trees this is a bit more complicated since the interior cells contain payloads (the key is the payload). + // and for non-unique indexes there might be several cells with the same key. + // + // Forwards iteration (looking for first match in tree): + // OP | Current Cell vs Seek Key | Action? | Explanation + // GT | > | go left | First > key could be exactly this one, or in left subtree + // GT | = or < | go right | First > key must be in right subtree + // GE | > | go left | First >= key could be exactly this one, or in left subtree + // GE | = | go left | First >= key could be exactly this one, or in left subtree + // GE | < | go right | First >= key must be in right subtree + // + // Backwards iteration (looking for last match in tree): + // OP | Current Cell vs Seek Key | Action? | Explanation + // LE | > | go left | Last <= key must be in left subtree + // LE | = | go right | Last <= key is either this one, or somewhere to the right of this one. So we need to go right to make sure + // LE | < | go right | Last <= key must be in right subtree + // LT | > | go left | Last < key must be in left subtree + // LT | = | go left | Last < key must be in left subtree since we want strictly less than + // LT | < | go right | Last < key could be exactly this one, or in right subtree + // + // No iteration (point query): + // EQ | > | go left | First = key must be in left subtree + // EQ | = | go left | First = key could be exactly this one, or in left subtree + // EQ | < | go right | First = key must be in right subtree - let target_leaf_page_is_in_left_subtree = match cmp { - SeekOp::GT => interior_cell_vs_index_key.is_gt(), - SeekOp::GE => interior_cell_vs_index_key.is_ge(), - SeekOp::EQ => interior_cell_vs_index_key.is_ge(), - SeekOp::LE => interior_cell_vs_index_key.is_gt(), - SeekOp::LT => interior_cell_vs_index_key.is_ge(), + match cmp { + SeekOp::GT => interior_cell_vs_index_key.is_gt(), + SeekOp::GE => interior_cell_vs_index_key.is_ge(), + SeekOp::EQ => interior_cell_vs_index_key.is_ge(), + SeekOp::LE => interior_cell_vs_index_key.is_gt(), + SeekOp::LT => interior_cell_vs_index_key.is_ge(), + } }; if target_leaf_page_is_in_left_subtree { - leftmost_matching_cell = Some(cur_cell_idx as usize); - max = cur_cell_idx - 1; + self.move_to_state + .set_nearest_matching_cell(Some(cur_cell_idx as usize)); + self.move_to_state.set_max(cur_cell_idx - 1); } else { - min = cur_cell_idx + 1; + self.move_to_state.set_min(cur_cell_idx + 1); } } } @@ -1791,7 +1824,7 @@ impl BTreeCursor { ) -> Result> { assert!(self.mv_cursor.is_none()); self.move_to_root(); - return_if_io!(self.tablebtree_move_to(rowid, seek_op)); + return_if_io!(self.move_to(SeekKey::TableRowId(rowid), seek_op)); let page = self.stack.top(); return_if_locked!(page.get()); let page = page.get(); @@ -1954,7 +1987,7 @@ impl BTreeCursor { seek_op: SeekOp, ) -> Result> { self.move_to_root(); - return_if_io!(self.indexbtree_move_to(key, seek_op)); + return_if_io!(self.move_to(SeekKey::IndexKey(key), seek_op)); let page = self.stack.top(); return_if_locked!(page.get()); @@ -2170,14 +2203,13 @@ impl BTreeCursor { // 6. If we find the cell, we return the record. Otherwise, we return an empty result. self.move_to_root(); - match key { - SeekKey::TableRowId(rowid_key) => { - return self.tablebtree_move_to(rowid_key, cmp); - } - SeekKey::IndexKey(index_key) => { - return self.indexbtree_move_to(index_key, cmp); - } - } + let ret = match key { + SeekKey::TableRowId(rowid_key) => self.tablebtree_move_to(rowid_key, cmp), + SeekKey::IndexKey(index_key) => self.indexbtree_move_to(index_key, cmp), + }; + return_if_io!(ret); + self.move_to_state = CursorMoveToState::Start; + Ok(CursorResult::Ok(())) } /// Insert a record into the btree.