From 9909539b9deb85a8a6b04910b6517b066e062ca6 Mon Sep 17 00:00:00 2001 From: Jussi Saurio Date: Sat, 11 Jan 2025 14:04:14 +0200 Subject: [PATCH 1/4] Store cursor type (table,index,pseudo,sorter) when allocating cursor --- core/schema.rs | 8 +++++++ core/translate/group_by.rs | 12 +++++++---- core/translate/insert.rs | 42 +++++++++++++++++-------------------- core/translate/main_loop.rs | 17 +++++++++------ core/translate/mod.rs | 8 ++++--- core/translate/order_by.rs | 19 +++++++++-------- core/translate/plan.rs | 8 ++++++- core/vdbe/builder.rs | 28 ++++++++++++++++++++----- core/vdbe/explain.rs | 21 ++++++++++++++----- core/vdbe/mod.rs | 2 +- 10 files changed, 108 insertions(+), 57 deletions(-) diff --git a/core/schema.rs b/core/schema.rs index 82932058..e28a1546 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -135,6 +135,14 @@ impl Table { Self::Pseudo(_) => unimplemented!(), } } + + pub fn btree(&self) -> Option> { + match self { + Self::BTree(table) => Some(table.clone()), + Self::Index(_) => None, + Self::Pseudo(_) => None, + } + } } impl PartialEq for Table { diff --git a/core/translate/group_by.rs b/core/translate/group_by.rs index c20466f2..2b6e7afb 100644 --- a/core/translate/group_by.rs +++ b/core/translate/group_by.rs @@ -4,9 +4,13 @@ use sqlite3_parser::ast; use crate::{ function::AggFunc, - schema::{Column, PseudoTable, Table}, + schema::{Column, PseudoTable}, types::{OwnedRecord, OwnedValue}, - vdbe::{builder::ProgramBuilder, insn::Insn, BranchOffset}, + vdbe::{ + builder::{CursorType, ProgramBuilder}, + insn::Insn, + BranchOffset, + }, Result, }; @@ -50,7 +54,7 @@ pub fn init_group_by( ) -> Result<()> { let num_aggs = aggregates.len(); - let sort_cursor = program.alloc_cursor_id(None, None); + let sort_cursor = program.alloc_cursor_id(None, CursorType::Sorter); let reg_abort_flag = program.alloc_register(); let reg_group_exprs_cmp = program.alloc_registers(group_by.exprs.len()); @@ -175,7 +179,7 @@ pub fn emit_group_by<'a>( columns: pseudo_columns, }); - let pseudo_cursor = program.alloc_cursor_id(None, Some(Table::Pseudo(pseudo_table.clone()))); + let pseudo_cursor = program.alloc_cursor_id(None, CursorType::Pseudo(pseudo_table.clone())); program.emit_insn(Insn::OpenPseudo { cursor_id: pseudo_cursor, diff --git a/core/translate/insert.rs b/core/translate/insert.rs index b9f73ba8..2dec7424 100644 --- a/core/translate/insert.rs +++ b/core/translate/insert.rs @@ -6,13 +6,18 @@ use sqlite3_parser::ast::{ }; use crate::error::SQLITE_CONSTRAINT_PRIMARYKEY; +use crate::schema::BTreeTable; use crate::util::normalize_ident; use crate::vdbe::BranchOffset; use crate::{ - schema::{Column, Schema, Table}, + schema::{Column, Schema}, storage::sqlite3_ondisk::DatabaseHeader, translate::expr::translate_expr, - vdbe::{builder::ProgramBuilder, insn::Insn, Program}, + vdbe::{ + builder::{CursorType, ProgramBuilder}, + insn::Insn, + Program, + }, SymbolTable, }; use crate::{Connection, Result}; @@ -53,20 +58,15 @@ pub fn translate_insert( Some(table) => table, None => crate::bail_corrupt_error!("Parse error: no such table: {}", table_name), }; - let table = Rc::new(Table::BTree(table)); - if !table.has_rowid() { + if !table.has_rowid { crate::bail_parse_error!("INSERT into WITHOUT ROWID table is not supported"); } let cursor_id = program.alloc_cursor_id( Some(table_name.0.clone()), - Some(table.clone().deref().clone()), + CursorType::BTreeTable(table.clone()), ); - let root_page = match table.as_ref() { - Table::BTree(btree) => btree.root_page, - Table::Index(index) => index.root_page, - Table::Pseudo(_) => todo!(), - }; + let root_page = table.root_page; let values = match body { InsertBody::Select(select, None) => match &select.body.select.deref() { sqlite3_parser::ast::OneSelect::Values(values) => values, @@ -77,9 +77,9 @@ pub fn translate_insert( let column_mappings = resolve_columns_for_insert(&table, columns, values)?; // Check if rowid was provided (through INTEGER PRIMARY KEY as a rowid alias) - let rowid_alias_index = table.columns().iter().position(|c| c.is_rowid_alias); + let rowid_alias_index = table.columns.iter().position(|c| c.is_rowid_alias); let has_user_provided_rowid = { - assert!(column_mappings.len() == table.columns().len()); + assert!(column_mappings.len() == table.columns.len()); if let Some(index) = rowid_alias_index { column_mappings[index].value_index.is_some() } else { @@ -89,7 +89,7 @@ pub fn translate_insert( // allocate a register for each column in the table. if not provided by user, they will simply be set as null. // allocate an extra register for rowid regardless of whether user provided a rowid alias column. - let num_cols = table.columns().len(); + let num_cols = table.columns.len(); let rowid_reg = program.alloc_registers(num_cols + 1); let column_registers_start = rowid_reg + 1; let rowid_alias_reg = { @@ -215,14 +215,14 @@ pub fn translate_insert( target_pc: make_record_label, }); let rowid_column_name = if let Some(index) = rowid_alias_index { - table.column_index_to_name(index).unwrap() + &table.columns.get(index).unwrap().name } else { "rowid" }; program.emit_insn(Insn::Halt { err_code: SQLITE_CONSTRAINT_PRIMARYKEY, - description: format!("{}.{}", table.get_name(), rowid_column_name), + description: format!("{}.{}", table_name.0, rowid_column_name), }); program.resolve_label(make_record_label, program.offset()); @@ -293,7 +293,7 @@ struct ColumnMapping<'a> { /// - Named columns map to their corresponding value index /// - Unspecified columns map to None fn resolve_columns_for_insert<'a>( - table: &'a Table, + table: &'a BTreeTable, columns: &Option, values: &[Vec], ) -> Result>> { @@ -301,7 +301,7 @@ fn resolve_columns_for_insert<'a>( crate::bail_parse_error!("no values to insert"); } - let table_columns = table.columns(); + let table_columns = &table.columns; // Case 1: No columns specified - map values to columns in order if columns.is_none() { @@ -309,7 +309,7 @@ fn resolve_columns_for_insert<'a>( if num_values > table_columns.len() { crate::bail_parse_error!( "table {} has {} columns but {} values were supplied", - table.get_name(), + &table.name, table_columns.len(), num_values ); @@ -350,11 +350,7 @@ fn resolve_columns_for_insert<'a>( .position(|c| c.name.eq_ignore_ascii_case(&column_name)); if table_index.is_none() { - crate::bail_parse_error!( - "table {} has no column named {}", - table.get_name(), - column_name - ); + crate::bail_parse_error!("table {} has no column named {}", &table.name, column_name); } mappings[table_index.unwrap()].value_index = Some(value_index); diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index e95b3dc2..33f76a84 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -1,9 +1,12 @@ use sqlite3_parser::ast; use crate::{ - schema::Table, translate::result_row::emit_select_result, - vdbe::{builder::ProgramBuilder, insn::Insn, BranchOffset}, + vdbe::{ + builder::{CursorType, ProgramBuilder}, + insn::Insn, + BranchOffset, + }, Result, }; @@ -81,7 +84,7 @@ pub fn init_loop( } => { let cursor_id = program.alloc_cursor_id( Some(table_reference.table_identifier.clone()), - Some(table_reference.table.clone()), + CursorType::BTreeTable(table_reference.btree().unwrap().clone()), ); let root_page = table_reference.table.get_root_page(); @@ -114,7 +117,7 @@ pub fn init_loop( } => { let table_cursor_id = program.alloc_cursor_id( Some(table_reference.table_identifier.clone()), - Some(table_reference.table.clone()), + CursorType::BTreeTable(table_reference.btree().unwrap().clone()), ); match mode { @@ -138,8 +141,10 @@ pub fn init_loop( } if let Search::IndexSearch { index, .. } = search { - let index_cursor_id = program - .alloc_cursor_id(Some(index.name.clone()), Some(Table::Index(index.clone()))); + let index_cursor_id = program.alloc_cursor_id( + Some(index.name.clone()), + CursorType::BTreeIndex(index.clone()), + ); match mode { OperationMode::SELECT => { diff --git a/core/translate/mod.rs b/core/translate/mod.rs index 92b661ce..fdbbc47e 100644 --- a/core/translate/mod.rs +++ b/core/translate/mod.rs @@ -27,6 +27,7 @@ use crate::storage::pager::Pager; use crate::storage::sqlite3_ondisk::{DatabaseHeader, MIN_PAGE_CACHE_SIZE}; use crate::translate::delete::translate_delete; use crate::util::PRIMARY_KEY_AUTOMATIC_INDEX_NAME_PREFIX; +use crate::vdbe::builder::CursorType; use crate::vdbe::{builder::ProgramBuilder, insn::Insn, Program}; use crate::{bail_parse_error, Connection, LimboError, Result, SymbolTable}; use insert::translate_insert; @@ -463,9 +464,10 @@ fn translate_create_table( let table_id = "sqlite_schema".to_string(); let table = schema.get_table(&table_id).unwrap(); - let table = crate::schema::Table::BTree(table.clone()); - let sqlite_schema_cursor_id = - program.alloc_cursor_id(Some(table_id.to_owned()), Some(table.to_owned())); + let sqlite_schema_cursor_id = program.alloc_cursor_id( + Some(table_id.to_owned()), + CursorType::BTreeTable(table.clone()), + ); program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, root_page: 1, diff --git a/core/translate/order_by.rs b/core/translate/order_by.rs index b58e7ec2..1d02639d 100644 --- a/core/translate/order_by.rs +++ b/core/translate/order_by.rs @@ -3,10 +3,13 @@ use std::rc::Rc; use sqlite3_parser::ast; use crate::{ - schema::{Column, PseudoTable, Table}, + schema::{Column, PseudoTable}, types::{OwnedRecord, OwnedValue}, util::exprs_are_equivalent, - vdbe::{builder::ProgramBuilder, insn::Insn}, + vdbe::{ + builder::{CursorType, ProgramBuilder}, + insn::Insn, + }, Result, }; @@ -32,7 +35,7 @@ pub fn init_order_by( t_ctx: &mut TranslateCtx, order_by: &[(ast::Expr, Direction)], ) -> Result<()> { - let sort_cursor = program.alloc_cursor_id(None, None); + let sort_cursor = program.alloc_cursor_id(None, CursorType::Sorter); t_ctx.meta_sort = Some(SortMetadata { sort_cursor, reg_sorter_data: program.alloc_register(), @@ -93,12 +96,10 @@ pub fn emit_order_by( .map(|v| v.len()) .unwrap_or(0); - let pseudo_cursor = program.alloc_cursor_id( - None, - Some(Table::Pseudo(Rc::new(PseudoTable { - columns: pseudo_columns, - }))), - ); + let pseudo_table = Rc::new(PseudoTable { + columns: pseudo_columns, + }); + let pseudo_cursor = program.alloc_cursor_id(None, CursorType::Pseudo(pseudo_table.clone())); let SortMetadata { sort_cursor, reg_sorter_data, diff --git a/core/translate/plan.rs b/core/translate/plan.rs index b773d1ee..bffd384d 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -7,7 +7,7 @@ use std::{ use crate::{ function::AggFunc, - schema::{Column, Index, Table}, + schema::{BTreeTable, Column, Index, Table}, vdbe::BranchOffset, Result, }; @@ -255,6 +255,12 @@ pub struct TableReference { } impl TableReference { + pub fn btree(&self) -> Option> { + match self.reference_type { + TableReferenceType::BTreeTable => self.table.btree(), + TableReferenceType::Subquery { .. } => None, + } + } pub fn new_subquery(identifier: String, table_index: usize, plan: &SelectPlan) -> Self { Self { table: Table::Pseudo(Rc::new(PseudoTable::new_with_columns( diff --git a/core/vdbe/builder.rs b/core/vdbe/builder.rs index 40f936cd..7acc4be6 100644 --- a/core/vdbe/builder.rs +++ b/core/vdbe/builder.rs @@ -4,9 +4,13 @@ use std::{ rc::{Rc, Weak}, }; -use crate::{storage::sqlite3_ondisk::DatabaseHeader, Connection}; +use crate::{ + schema::{BTreeTable, Index, PseudoTable}, + storage::sqlite3_ondisk::DatabaseHeader, + Connection, +}; -use super::{BranchOffset, CursorID, Insn, InsnReference, Program, Table}; +use super::{BranchOffset, CursorID, Insn, InsnReference, Program}; #[allow(dead_code)] pub struct ProgramBuilder { @@ -18,7 +22,7 @@ pub struct ProgramBuilder { constant_insns: Vec, next_insn_label: Option, // Cursors that are referenced by the program. Indexed by CursorID. - pub cursor_ref: Vec<(Option, Option)>, + pub cursor_ref: Vec<(Option, CursorType)>, // Hashmap of label to insn reference. Resolved in build(). label_to_resolved_offset: HashMap, // Bitmask of cursors that have emitted a SeekRowid instruction. @@ -27,6 +31,20 @@ pub struct ProgramBuilder { comments: HashMap, } +#[derive(Debug, Clone)] +pub enum CursorType { + BTreeTable(Rc), + BTreeIndex(Rc), + Pseudo(Rc), + Sorter, +} + +impl CursorType { + pub fn is_index(&self) -> bool { + matches!(self, CursorType::BTreeIndex(_)) + } +} + impl ProgramBuilder { pub fn new() -> Self { Self { @@ -58,11 +76,11 @@ impl ProgramBuilder { pub fn alloc_cursor_id( &mut self, table_identifier: Option, - table: Option
, + cursor_type: CursorType, ) -> usize { let cursor = self.next_free_cursor_id; self.next_free_cursor_id += 1; - self.cursor_ref.push((table_identifier, table)); + self.cursor_ref.push((table_identifier, cursor_type)); assert!(self.cursor_ref.len() == self.next_free_cursor_id); cursor } diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index 133172b7..40fda1a2 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -1,3 +1,5 @@ +use crate::vdbe::builder::CursorType; + use super::{Insn, InsnReference, OwnedValue, Program}; use std::rc::Rc; @@ -387,7 +389,19 @@ pub fn insn_to_str( column, dest, } => { - let (table_identifier, table) = &program.cursor_ref[*cursor_id]; + let (table_identifier, cursor_type) = &program.cursor_ref[*cursor_id]; + let column_name = match cursor_type { + CursorType::BTreeTable(table) => { + Some(&table.columns.get(*column).unwrap().name) + } + CursorType::BTreeIndex(index) => { + Some(&index.columns.get(*column).unwrap().name) + } + CursorType::Pseudo(pseudo_table) => { + Some(&pseudo_table.columns.get(*column).unwrap().name) + } + CursorType::Sorter => None, + }; ( "Column", *cursor_id as i32, @@ -401,10 +415,7 @@ pub fn insn_to_str( table_identifier .as_ref() .unwrap_or(&format!("cursor {}", cursor_id)), - table - .as_ref() - .and_then(|x| x.column_index_to_name(*column)) - .unwrap_or(format!("column {}", *column).as_str()) + column_name.unwrap_or(&format!("column {}", *column)) ), ) } diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 64c41b3b..bceaeac4 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -211,7 +211,7 @@ impl ProgramState { pub struct Program { pub max_registers: usize, pub insns: Vec, - pub cursor_ref: Vec<(Option, Option
)>, + pub cursor_ref: Vec<(Option, CursorType)>, pub database_header: Rc>, pub comments: HashMap, pub connection: Weak, From 9f44d2a7ac1dfc6aa95e6ed994c7828ebb7b5696 Mon Sep 17 00:00:00 2001 From: Jussi Saurio Date: Sat, 11 Jan 2025 17:07:30 +0200 Subject: [PATCH 2/4] Remove unused Table::Index variant and unused Table methods --- core/schema.rs | 59 -------------------------------------------------- 1 file changed, 59 deletions(-) diff --git a/core/schema.rs b/core/schema.rs index e28a1546..fda6c12b 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -46,35 +46,13 @@ impl Schema { #[derive(Clone, Debug)] pub enum Table { BTree(Rc), - Index(Rc), Pseudo(Rc), } impl Table { - pub fn is_pseudo(&self) -> bool { - matches!(self, Table::Pseudo(_)) - } - - pub fn get_rowid_alias_column(&self) -> Option<(usize, &Column)> { - match self { - Self::BTree(table) => table.get_rowid_alias_column(), - Self::Index(_) => None, - Self::Pseudo(_) => None, - } - } - - pub fn column_is_rowid_alias(&self, col: &Column) -> bool { - match self { - Table::BTree(table) => table.column_is_rowid_alias(col), - Self::Index(_) => false, - Self::Pseudo(_) => false, - } - } - pub fn get_root_page(&self) -> usize { match self { Table::BTree(table) => table.root_page, - Table::Index(_) => unimplemented!(), Table::Pseudo(_) => unimplemented!(), } } @@ -82,40 +60,13 @@ impl Table { pub fn get_name(&self) -> &str { match self { Self::BTree(table) => &table.name, - Self::Index(index) => &index.name, Self::Pseudo(_) => "", } } - pub fn column_index_to_name(&self, index: usize) -> Option<&str> { - match self { - Self::BTree(table) => match table.columns.get(index) { - Some(column) => Some(&column.name), - None => None, - }, - Self::Index(i) => match i.columns.get(index) { - Some(column) => Some(&column.name), - None => None, - }, - Self::Pseudo(table) => match table.columns.get(index) { - Some(_) => None, - None => None, - }, - } - } - - pub fn get_column(&self, name: &str) -> Option<(usize, &Column)> { - match self { - Self::BTree(table) => table.get_column(name), - Self::Index(_) => unimplemented!(), - Self::Pseudo(table) => table.get_column(name), - } - } - pub fn get_column_at(&self, index: usize) -> &Column { match self { Self::BTree(table) => table.columns.get(index).unwrap(), - Self::Index(_) => unimplemented!(), Self::Pseudo(table) => table.columns.get(index).unwrap(), } } @@ -123,23 +74,13 @@ impl Table { pub fn columns(&self) -> &Vec { match self { Self::BTree(table) => &table.columns, - Self::Index(_) => unimplemented!(), Self::Pseudo(table) => &table.columns, } } - pub fn has_rowid(&self) -> bool { - match self { - Self::BTree(table) => table.has_rowid, - Self::Index(_) => unimplemented!(), - Self::Pseudo(_) => unimplemented!(), - } - } - pub fn btree(&self) -> Option> { match self { Self::BTree(table) => Some(table.clone()), - Self::Index(_) => None, Self::Pseudo(_) => None, } } From bf48c0ae72441aee9f7bb94a632b12bf7c3973a0 Mon Sep 17 00:00:00 2001 From: Jussi Saurio Date: Sat, 11 Jan 2025 17:07:51 +0200 Subject: [PATCH 3/4] Remove trait Cursor --- core/pseudo.rs | 105 +----------- core/storage/btree.rs | 288 ++++++++++++++++---------------- core/types.rs | 27 +-- core/vdbe/mod.rs | 376 ++++++++++++++++++++++++------------------ core/vdbe/sorter.rs | 93 ++--------- 5 files changed, 377 insertions(+), 512 deletions(-) diff --git a/core/pseudo.rs b/core/pseudo.rs index 45f47856..93d9d6a6 100644 --- a/core/pseudo.rs +++ b/core/pseudo.rs @@ -1,110 +1,19 @@ -use crate::{ - types::{SeekKey, SeekOp}, - Result, -}; -use std::cell::{Ref, RefCell}; - -use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue}; +use crate::types::OwnedRecord; pub struct PseudoCursor { - current: RefCell>, + current: Option, } impl PseudoCursor { pub fn new() -> Self { - Self { - current: RefCell::new(None), - } - } -} - -impl Cursor for PseudoCursor { - fn is_empty(&self) -> bool { - self.current.borrow().is_none() - } - - fn root_page(&self) -> usize { - unreachable!() - } - - fn rewind(&mut self) -> Result> { - *self.current.borrow_mut() = None; - Ok(CursorResult::Ok(())) - } - - fn next(&mut self) -> Result> { - *self.current.borrow_mut() = None; - Ok(CursorResult::Ok(())) - } - - fn wait_for_completion(&mut self) -> Result<()> { - Ok(()) - } - - fn rowid(&self) -> Result> { - let x = self - .current - .borrow() - .as_ref() - .map(|record| match record.values[0] { - OwnedValue::Integer(rowid) => rowid as u64, - ref ov => { - panic!("Expected integer value, got {:?}", ov); - } - }); - Ok(x) - } - - fn seek(&mut self, _: SeekKey<'_>, _: SeekOp) -> Result> { - unimplemented!(); - } - - fn seek_to_last(&mut self) -> Result> { - unimplemented!(); - } - - fn record(&self) -> Result>> { - Ok(self.current.borrow()) - } - - fn insert( - &mut self, - key: &OwnedValue, - record: &OwnedRecord, - moved_before: bool, - ) -> Result> { - let _ = key; - let _ = moved_before; - *self.current.borrow_mut() = Some(record.clone()); - Ok(CursorResult::Ok(())) - } - - fn delete(&mut self) -> Result> { - unimplemented!() - } - - fn get_null_flag(&self) -> bool { - false - } - - fn set_null_flag(&mut self, _null_flag: bool) { - // Do nothing - } - - fn exists(&mut self, key: &OwnedValue) -> Result> { - let _ = key; - todo!() - } - - fn btree_create(&mut self, _flags: usize) -> u32 { - unreachable!("Please don't.") + Self { current: None } } - fn last(&mut self) -> Result> { - todo!() + pub fn record(&self) -> Option<&OwnedRecord> { + self.current.as_ref() } - fn prev(&mut self) -> Result> { - todo!() + pub fn insert(&mut self, record: OwnedRecord) { + self.current = Some(record); } } diff --git a/core/storage/btree.rs b/core/storage/btree.rs index c84a0c2c..5bdcb135 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -5,7 +5,7 @@ use crate::storage::sqlite3_ondisk::{ read_btree_cell, read_varint, write_varint, BTreeCell, DatabaseHeader, PageContent, PageType, TableInteriorCell, TableLeafCell, }; -use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue, SeekKey, SeekOp}; +use crate::types::{CursorResult, OwnedRecord, OwnedValue, SeekKey, SeekOp}; use crate::Result; use std::cell::{Ref, RefCell}; @@ -419,7 +419,7 @@ impl BTreeCursor { /// This may be used to seek to a specific record in a point query (e.g. SELECT * FROM table WHERE col = 10) /// or e.g. find the first record greater than the seek key in a range query (e.g. SELECT * FROM table WHERE col > 10). /// We don't include the rowid in the comparison and that's why the last value from the record is not included. - fn seek( + fn do_seek( &mut self, key: SeekKey<'_>, op: SeekOp, @@ -1697,133 +1697,8 @@ impl BTreeCursor { } cell_idx } -} - -impl PageStack { - /// Push a new page onto the stack. - /// This effectively means traversing to a child page. - fn push(&self, page: PageRef) { - debug!( - "pagestack::push(current={}, new_page_id={})", - self.current_page.borrow(), - page.get().id - ); - *self.current_page.borrow_mut() += 1; - let current = *self.current_page.borrow(); - assert!( - current < BTCURSOR_MAX_DEPTH as i32, - "corrupted database, stack is bigger than expected" - ); - self.stack.borrow_mut()[current as usize] = Some(page); - self.cell_indices.borrow_mut()[current as usize] = 0; - } - - /// Pop a page off the stack. - /// This effectively means traversing back up to a parent page. - fn pop(&self) { - let current = *self.current_page.borrow(); - debug!("pagestack::pop(current={})", current); - self.cell_indices.borrow_mut()[current as usize] = 0; - self.stack.borrow_mut()[current as usize] = None; - *self.current_page.borrow_mut() -= 1; - } - - /// Get the top page on the stack. - /// This is the page that is currently being traversed. - fn top(&self) -> PageRef { - let current = *self.current_page.borrow(); - let page = self.stack.borrow()[current as usize] - .as_ref() - .unwrap() - .clone(); - debug!( - "pagestack::top(current={}, page_id={})", - current, - page.get().id - ); - page - } - - /// Get the parent page of the current page. - fn parent(&self) -> PageRef { - let current = *self.current_page.borrow(); - self.stack.borrow()[current as usize - 1] - .as_ref() - .unwrap() - .clone() - } - - /// Current page pointer being used - fn current(&self) -> usize { - *self.current_page.borrow() as usize - } - - /// Cell index of the current page - fn current_cell_index(&self) -> i32 { - let current = self.current(); - self.cell_indices.borrow()[current] - } - - /// Check if the current cell index is less than 0. - /// This means we have been iterating backwards and have reached the start of the page. - fn current_cell_index_less_than_min(&self) -> bool { - let cell_idx = self.current_cell_index(); - cell_idx < 0 - } - - /// Advance the current cell index of the current page to the next cell. - fn advance(&self) { - let current = self.current(); - self.cell_indices.borrow_mut()[current] += 1; - } - - fn retreat(&self) { - let current = self.current(); - self.cell_indices.borrow_mut()[current] -= 1; - } - - fn set_cell_index(&self, idx: i32) { - let current = self.current(); - self.cell_indices.borrow_mut()[current] = idx - } - - fn has_parent(&self) -> bool { - *self.current_page.borrow() > 0 - } - - fn clear(&self) { - *self.current_page.borrow_mut() = -1; - } -} - -fn find_free_cell(page_ref: &PageContent, db_header: Ref, amount: usize) -> usize { - // NOTE: freelist is in ascending order of keys and pc - // unuse_space is reserved bytes at the end of page, therefore we must substract from maxpc - let mut pc = page_ref.first_freeblock() as usize; - - let buf = page_ref.as_ptr(); - - let usable_space = (db_header.page_size - db_header.reserved_space as u16) as usize; - let maxpc = usable_space - amount; - let mut found = false; - while pc <= maxpc { - let next = u16::from_be_bytes(buf[pc..pc + 2].try_into().unwrap()); - let size = u16::from_be_bytes(buf[pc + 2..pc + 4].try_into().unwrap()); - if amount <= size as usize { - found = true; - break; - } - pc = next as usize; - } - if !found { - 0 - } else { - pc - } -} -impl Cursor for BTreeCursor { - fn seek_to_last(&mut self) -> Result> { + pub fn seek_to_last(&mut self) -> Result> { return_if_io!(self.move_to_rightmost()); let (rowid, record) = return_if_io!(self.get_next_record(None)); if rowid.is_none() { @@ -1836,15 +1711,15 @@ impl Cursor for BTreeCursor { Ok(CursorResult::Ok(())) } - fn is_empty(&self) -> bool { + pub fn is_empty(&self) -> bool { self.record.borrow().is_none() } - fn root_page(&self) -> usize { + pub fn root_page(&self) -> usize { self.root_page } - fn rewind(&mut self) -> Result> { + pub fn rewind(&mut self) -> Result> { self.move_to_root(); let (rowid, record) = return_if_io!(self.get_next_record(None)); @@ -1853,21 +1728,21 @@ impl Cursor for BTreeCursor { Ok(CursorResult::Ok(())) } - fn last(&mut self) -> Result> { + pub fn last(&mut self) -> Result> { match self.move_to_rightmost()? { CursorResult::Ok(_) => self.prev(), CursorResult::IO => Ok(CursorResult::IO), } } - fn next(&mut self) -> Result> { + pub fn next(&mut self) -> Result> { let (rowid, record) = return_if_io!(self.get_next_record(None)); self.rowid.replace(rowid); self.record.replace(record); Ok(CursorResult::Ok(())) } - fn prev(&mut self) -> Result> { + pub fn prev(&mut self) -> Result> { match self.get_prev_record()? { CursorResult::Ok((rowid, record)) => { self.rowid.replace(rowid); @@ -1878,27 +1753,27 @@ impl Cursor for BTreeCursor { } } - fn wait_for_completion(&mut self) -> Result<()> { + pub fn wait_for_completion(&mut self) -> Result<()> { // TODO: Wait for pager I/O to complete Ok(()) } - fn rowid(&self) -> Result> { + pub fn rowid(&self) -> Result> { Ok(*self.rowid.borrow()) } - fn seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result> { - let (rowid, record) = return_if_io!(self.seek(key, op)); + pub fn seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result> { + let (rowid, record) = return_if_io!(self.do_seek(key, op)); self.rowid.replace(rowid); self.record.replace(record); Ok(CursorResult::Ok(rowid.is_some())) } - fn record(&self) -> Result>> { + pub fn record(&self) -> Result>> { Ok(self.record.borrow()) } - fn insert( + pub fn insert( &mut self, key: &OwnedValue, _record: &OwnedRecord, @@ -1917,20 +1792,20 @@ impl Cursor for BTreeCursor { Ok(CursorResult::Ok(())) } - fn delete(&mut self) -> Result> { + pub fn delete(&mut self) -> Result> { println!("rowid: {:?}", self.rowid.borrow()); Ok(CursorResult::Ok(())) } - fn set_null_flag(&mut self, flag: bool) { + pub fn set_null_flag(&mut self, flag: bool) { self.null_flag = flag; } - fn get_null_flag(&self) -> bool { + pub fn get_null_flag(&self) -> bool { self.null_flag } - fn exists(&mut self, key: &OwnedValue) -> Result> { + pub fn exists(&mut self, key: &OwnedValue) -> Result> { let int_key = match key { OwnedValue::Integer(i) => i, _ => unreachable!("btree tables are indexed by integers!"), @@ -1965,7 +1840,7 @@ impl Cursor for BTreeCursor { } } - fn btree_create(&mut self, flags: usize) -> u32 { + pub fn btree_create(&mut self, flags: usize) -> u32 { let page_type = match flags { 1 => PageType::TableLeaf, 2 => PageType::IndexLeaf, @@ -1980,6 +1855,129 @@ impl Cursor for BTreeCursor { } } +impl PageStack { + /// Push a new page onto the stack. + /// This effectively means traversing to a child page. + fn push(&self, page: PageRef) { + debug!( + "pagestack::push(current={}, new_page_id={})", + self.current_page.borrow(), + page.get().id + ); + *self.current_page.borrow_mut() += 1; + let current = *self.current_page.borrow(); + assert!( + current < BTCURSOR_MAX_DEPTH as i32, + "corrupted database, stack is bigger than expected" + ); + self.stack.borrow_mut()[current as usize] = Some(page); + self.cell_indices.borrow_mut()[current as usize] = 0; + } + + /// Pop a page off the stack. + /// This effectively means traversing back up to a parent page. + fn pop(&self) { + let current = *self.current_page.borrow(); + debug!("pagestack::pop(current={})", current); + self.cell_indices.borrow_mut()[current as usize] = 0; + self.stack.borrow_mut()[current as usize] = None; + *self.current_page.borrow_mut() -= 1; + } + + /// Get the top page on the stack. + /// This is the page that is currently being traversed. + fn top(&self) -> PageRef { + let current = *self.current_page.borrow(); + let page = self.stack.borrow()[current as usize] + .as_ref() + .unwrap() + .clone(); + debug!( + "pagestack::top(current={}, page_id={})", + current, + page.get().id + ); + page + } + + /// Get the parent page of the current page. + fn parent(&self) -> PageRef { + let current = *self.current_page.borrow(); + self.stack.borrow()[current as usize - 1] + .as_ref() + .unwrap() + .clone() + } + + /// Current page pointer being used + fn current(&self) -> usize { + *self.current_page.borrow() as usize + } + + /// Cell index of the current page + fn current_cell_index(&self) -> i32 { + let current = self.current(); + self.cell_indices.borrow()[current] + } + + /// Check if the current cell index is less than 0. + /// This means we have been iterating backwards and have reached the start of the page. + fn current_cell_index_less_than_min(&self) -> bool { + let cell_idx = self.current_cell_index(); + cell_idx < 0 + } + + /// Advance the current cell index of the current page to the next cell. + fn advance(&self) { + let current = self.current(); + self.cell_indices.borrow_mut()[current] += 1; + } + + fn retreat(&self) { + let current = self.current(); + self.cell_indices.borrow_mut()[current] -= 1; + } + + fn set_cell_index(&self, idx: i32) { + let current = self.current(); + self.cell_indices.borrow_mut()[current] = idx + } + + fn has_parent(&self) -> bool { + *self.current_page.borrow() > 0 + } + + fn clear(&self) { + *self.current_page.borrow_mut() = -1; + } +} + +fn find_free_cell(page_ref: &PageContent, db_header: Ref, amount: usize) -> usize { + // NOTE: freelist is in ascending order of keys and pc + // unuse_space is reserved bytes at the end of page, therefore we must substract from maxpc + let mut pc = page_ref.first_freeblock() as usize; + + let buf = page_ref.as_ptr(); + + let usable_space = (db_header.page_size - db_header.reserved_space as u16) as usize; + let maxpc = usable_space - amount; + let mut found = false; + while pc <= maxpc { + let next = u16::from_be_bytes(buf[pc..pc + 2].try_into().unwrap()); + let size = u16::from_be_bytes(buf[pc + 2..pc + 4].try_into().unwrap()); + if amount <= size as usize { + found = true; + break; + } + pc = next as usize; + } + if !found { + 0 + } else { + pc + } +} + pub fn btree_init_page( page: &PageRef, page_type: PageType, diff --git a/core/types.rs b/core/types.rs index 7b581439..d9a496bf 100644 --- a/core/types.rs +++ b/core/types.rs @@ -1,5 +1,5 @@ use std::fmt::Display; -use std::{cell::Ref, rc::Rc}; +use std::rc::Rc; use crate::error::LimboError; use crate::Result; @@ -524,31 +524,6 @@ pub enum SeekKey<'a> { IndexKey(&'a OwnedRecord), } -pub trait Cursor { - fn is_empty(&self) -> bool; - fn root_page(&self) -> usize; - fn rewind(&mut self) -> Result>; - fn last(&mut self) -> Result>; - fn next(&mut self) -> Result>; - fn prev(&mut self) -> Result>; - fn wait_for_completion(&mut self) -> Result<()>; - fn rowid(&self) -> Result>; - fn seek(&mut self, key: SeekKey, op: SeekOp) -> Result>; - fn seek_to_last(&mut self) -> Result>; - fn record(&self) -> Result>>; - fn insert( - &mut self, - key: &OwnedValue, - record: &OwnedRecord, - moved_before: bool, /* Tells inserter that it doesn't need to traverse in order to find leaf page */ - ) -> Result>; // - fn delete(&mut self) -> Result>; - fn exists(&mut self, key: &OwnedValue) -> Result>; - fn set_null_flag(&mut self, flag: bool); - fn get_null_flag(&self) -> bool; - fn btree_create(&mut self, flags: usize) -> u32; -} - #[cfg(test)] mod tests { use super::*; diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index bceaeac4..7b62ecae 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -30,13 +30,11 @@ use crate::ext::{exec_ts_from_uuid7, exec_uuid, exec_uuidblob, exec_uuidstr, Ext use crate::function::{AggFunc, FuncCtx, MathFunc, MathFuncArity, ScalarFunc}; use crate::pseudo::PseudoCursor; use crate::result::LimboResult; -use crate::schema::Table; use crate::storage::sqlite3_ondisk::DatabaseHeader; use crate::storage::{btree::BTreeCursor, pager::Pager}; -use crate::types::{ - AggContext, Cursor, CursorResult, OwnedRecord, OwnedValue, Record, SeekKey, SeekOp, -}; +use crate::types::{AggContext, CursorResult, OwnedRecord, OwnedValue, Record, SeekKey, SeekOp}; use crate::util::parse_schema_rows; +use crate::vdbe::builder::CursorType; use crate::vdbe::insn::Insn; #[cfg(feature = "json")] use crate::{ @@ -53,6 +51,7 @@ use likeop::{construct_like_escape_arg, exec_glob, exec_like_with_escape}; use rand::distributions::{Distribution, Uniform}; use rand::{thread_rng, Rng}; use regex::{Regex, RegexBuilder}; +use sorter::Sorter; use std::borrow::{Borrow, BorrowMut}; use std::cell::RefCell; use std::collections::{BTreeMap, HashMap}; @@ -164,7 +163,10 @@ impl RegexCache { /// The program state describes the environment in which the program executes. pub struct ProgramState { pub pc: InsnReference, - cursors: RefCell>>, + btree_table_cursors: RefCell>, + btree_index_cursors: RefCell>, + pseudo_cursors: RefCell>, + sorter_cursors: RefCell>, registers: Vec, last_compare: Option, deferred_seek: Option<(CursorID, CursorID)>, @@ -175,12 +177,18 @@ pub struct ProgramState { impl ProgramState { pub fn new(max_registers: usize) -> Self { - let cursors = RefCell::new(BTreeMap::new()); + let btree_table_cursors = RefCell::new(BTreeMap::new()); + let btree_index_cursors = RefCell::new(BTreeMap::new()); + let pseudo_cursors = RefCell::new(BTreeMap::new()); + let sorter_cursors = RefCell::new(BTreeMap::new()); let mut registers = Vec::with_capacity(max_registers); registers.resize(max_registers, OwnedValue::Null); Self { pc: 0, - cursors, + btree_table_cursors, + btree_index_cursors, + pseudo_cursors, + sorter_cursors, registers, last_compare: None, deferred_seek: None, @@ -207,6 +215,19 @@ impl ProgramState { } } +macro_rules! must_be_btree_cursor { + ($cursor_id:expr, $cursor_ref:expr, $btree_table_cursors:expr, $btree_index_cursors:expr, $insn_name:expr) => {{ + let (_, cursor_type) = $cursor_ref.get($cursor_id).unwrap(); + let cursor = match cursor_type { + CursorType::BTreeTable(_) => $btree_table_cursors.get_mut(&$cursor_id).unwrap(), + CursorType::BTreeIndex(_) => $btree_index_cursors.get_mut(&$cursor_id).unwrap(), + CursorType::Pseudo(_) => panic!("{} on pseudo cursor", $insn_name), + CursorType::Sorter => panic!("{} on sorter cursor", $insn_name), + }; + cursor + }}; +} + #[derive(Debug)] pub struct Program { pub max_registers: usize, @@ -248,7 +269,10 @@ impl Program { } let insn = &self.insns[state.pc as usize]; trace_insn(self, state.pc as InsnReference, insn); - let mut cursors = state.cursors.borrow_mut(); + let mut btree_table_cursors = state.btree_table_cursors.borrow_mut(); + let mut btree_index_cursors = state.btree_index_cursors.borrow_mut(); + let mut pseudo_cursors = state.pseudo_cursors.borrow_mut(); + let mut sorter_cursors = state.sorter_cursors.borrow_mut(); match insn { Insn::Init { target_pc } => { assert!(target_pc.is_offset()); @@ -304,7 +328,13 @@ impl Program { state.pc += 1; } Insn::NullRow { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "NullRow" + ); cursor.set_null_flag(true); state.pc += 1; } @@ -569,12 +599,23 @@ impl Program { cursor_id, root_page, } => { - let cursor = Box::new(BTreeCursor::new( - pager.clone(), - *root_page, - self.database_header.clone(), - )); - cursors.insert(*cursor_id, cursor); + let (_, cursor_type) = self.cursor_ref.get(*cursor_id).unwrap(); + let cursor = + BTreeCursor::new(pager.clone(), *root_page, self.database_header.clone()); + match cursor_type { + CursorType::BTreeTable(_) => { + btree_table_cursors.insert(*cursor_id, cursor); + } + CursorType::BTreeIndex(_) => { + btree_index_cursors.insert(*cursor_id, cursor); + } + CursorType::Pseudo(_) => { + panic!("OpenReadAsync on pseudo cursor"); + } + CursorType::Sorter => { + panic!("OpenReadAsync on sorter cursor"); + } + } state.pc += 1; } Insn::OpenReadAwait => { @@ -585,17 +626,29 @@ impl Program { content_reg: _, num_fields: _, } => { - let cursor = Box::new(PseudoCursor::new()); - cursors.insert(*cursor_id, cursor); + let cursor = PseudoCursor::new(); + pseudo_cursors.insert(*cursor_id, cursor); state.pc += 1; } Insn::RewindAsync { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "RewindAsync" + ); return_if_io!(cursor.rewind()); state.pc += 1; } Insn::LastAsync { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "LastAsync" + ); return_if_io!(cursor.last()); state.pc += 1; } @@ -604,7 +657,13 @@ impl Program { pc_if_empty, } => { assert!(pc_if_empty.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "LastAwait" + ); cursor.wait_for_completion()?; if cursor.is_empty() { state.pc = pc_if_empty.to_offset_int(); @@ -617,7 +676,13 @@ impl Program { pc_if_empty, } => { assert!(pc_if_empty.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "RewindAwait" + ); cursor.wait_for_completion()?; if cursor.is_empty() { state.pc = pc_if_empty.to_offset_int(); @@ -631,9 +696,9 @@ impl Program { dest, } => { if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() { - let index_cursor = cursors.get_mut(&index_cursor_id).unwrap(); + let index_cursor = btree_index_cursors.get_mut(&index_cursor_id).unwrap(); let rowid = index_cursor.rowid()?; - let table_cursor = cursors.get_mut(&table_cursor_id).unwrap(); + let table_cursor = btree_table_cursors.get_mut(&table_cursor_id).unwrap(); match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? { CursorResult::Ok(_) => {} CursorResult::IO => { @@ -642,18 +707,45 @@ impl Program { } } } - - let cursor = cursors.get_mut(cursor_id).unwrap(); - if let Some(ref record) = *cursor.record()? { - let null_flag = cursor.get_null_flag(); - state.registers[*dest] = if null_flag { - OwnedValue::Null - } else { - record.values[*column].clone() - }; - } else { - state.registers[*dest] = OwnedValue::Null; + let (_, cursor_type) = self.cursor_ref.get(*cursor_id).unwrap(); + match cursor_type { + CursorType::BTreeTable(_) | CursorType::BTreeIndex(_) => { + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "Column" + ); + let record = cursor.record()?; + if let Some(record) = record.as_ref() { + state.registers[*dest] = if cursor.get_null_flag() { + OwnedValue::Null + } else { + record.values[*column].clone() + }; + } else { + state.registers[*dest] = OwnedValue::Null; + } + } + CursorType::Sorter => { + let cursor = sorter_cursors.get_mut(cursor_id).unwrap(); + if let Some(record) = cursor.record() { + state.registers[*dest] = record.values[*column].clone(); + } else { + state.registers[*dest] = OwnedValue::Null; + } + } + CursorType::Pseudo(_) => { + let cursor = pseudo_cursors.get_mut(cursor_id).unwrap(); + if let Some(record) = cursor.record() { + state.registers[*dest] = record.values[*column].clone(); + } else { + state.registers[*dest] = OwnedValue::Null; + } + } } + state.pc += 1; } Insn::MakeRecord { @@ -671,13 +763,25 @@ impl Program { return Ok(StepResult::Row(record)); } Insn::NextAsync { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "NextAsync" + ); cursor.set_null_flag(false); return_if_io!(cursor.next()); state.pc += 1; } Insn::PrevAsync { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "PrevAsync" + ); cursor.set_null_flag(false); return_if_io!(cursor.prev()); state.pc += 1; @@ -687,7 +791,13 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "PrevAwait" + ); cursor.wait_for_completion()?; if !cursor.is_empty() { state.pc = pc_if_next.to_offset_int(); @@ -700,7 +810,13 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "NextAwait" + ); cursor.wait_for_completion()?; if !cursor.is_empty() { state.pc = pc_if_next.to_offset_int(); @@ -818,9 +934,9 @@ impl Program { } Insn::RowId { cursor_id, dest } => { if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() { - let index_cursor = cursors.get_mut(&index_cursor_id).unwrap(); + let index_cursor = btree_index_cursors.get_mut(&index_cursor_id).unwrap(); let rowid = index_cursor.rowid()?; - let table_cursor = cursors.get_mut(&table_cursor_id).unwrap(); + let table_cursor = btree_table_cursors.get_mut(&table_cursor_id).unwrap(); match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? { CursorResult::Ok(_) => {} CursorResult::IO => { @@ -830,7 +946,7 @@ impl Program { } } - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); if let Some(ref rowid) = cursor.rowid()? { state.registers[*dest] = OwnedValue::Integer(*rowid as i64); } else { @@ -844,7 +960,7 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); let rowid = match &state.registers[*src_reg] { OwnedValue::Integer(rowid) => *rowid as u64, OwnedValue::Null => { @@ -880,7 +996,7 @@ impl Program { } => { assert!(target_pc.is_offset()); if *is_index { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_index_cursors.get_mut(cursor_id).unwrap(); let record_from_regs: OwnedRecord = make_owned_record(&state.registers, start_reg, num_regs); let found = return_if_io!( @@ -892,7 +1008,7 @@ impl Program { state.pc += 1; } } else { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); let rowid = match &state.registers[*start_reg] { OwnedValue::Null => { // All integer values are greater than null so we just rewind the cursor @@ -925,7 +1041,7 @@ impl Program { } => { assert!(target_pc.is_offset()); if *is_index { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_index_cursors.get_mut(cursor_id).unwrap(); let record_from_regs: OwnedRecord = make_owned_record(&state.registers, start_reg, num_regs); let found = return_if_io!( @@ -937,7 +1053,7 @@ impl Program { state.pc += 1; } } else { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); let rowid = match &state.registers[*start_reg] { OwnedValue::Null => { // All integer values are greater than null so we just rewind the cursor @@ -968,7 +1084,7 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_index_cursors.get_mut(cursor_id).unwrap(); let record_from_regs: OwnedRecord = make_owned_record(&state.registers, start_reg, num_regs); if let Some(ref idx_record) = *cursor.record()? { @@ -991,7 +1107,7 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_index_cursors.get_mut(cursor_id).unwrap(); let record_from_regs: OwnedRecord = make_owned_record(&state.registers, start_reg, num_regs); if let Some(ref idx_record) = *cursor.record()? { @@ -1275,47 +1391,46 @@ impl Program { _ => unreachable!(), }) .collect(); - let cursor = Box::new(sorter::Sorter::new(order)); - cursors.insert(*cursor_id, cursor); + let cursor = sorter::Sorter::new(order); + sorter_cursors.insert(*cursor_id, cursor); state.pc += 1; } Insn::SorterData { cursor_id, dest_reg, - pseudo_cursor: sorter_cursor, + pseudo_cursor, } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); - let record = match *cursor.record()? { - Some(ref record) => record.clone(), + let sorter_cursor = sorter_cursors.get_mut(cursor_id).unwrap(); + let record = match sorter_cursor.record() { + Some(record) => record.clone(), None => { state.pc += 1; continue; } }; state.registers[*dest_reg] = OwnedValue::Record(record.clone()); - let sorter_cursor = cursors.get_mut(sorter_cursor).unwrap(); - sorter_cursor.insert(&OwnedValue::Integer(0), &record, false)?; // fix key later + let pseudo_cursor = pseudo_cursors.get_mut(pseudo_cursor).unwrap(); + pseudo_cursor.insert(record); state.pc += 1; } Insn::SorterInsert { cursor_id, record_reg, } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = sorter_cursors.get_mut(cursor_id).unwrap(); let record = match &state.registers[*record_reg] { OwnedValue::Record(record) => record, _ => unreachable!("SorterInsert on non-record register"), }; - // TODO: set correct key - cursor.insert(&OwnedValue::Integer(0), record, false)?; + cursor.insert(record); state.pc += 1; } Insn::SorterSort { cursor_id, pc_if_empty, } => { - if let Some(cursor) = cursors.get_mut(cursor_id) { - cursor.rewind()?; + if let Some(cursor) = sorter_cursors.get_mut(cursor_id) { + cursor.sort(); state.pc += 1; } else { state.pc = pc_if_empty.to_offset_int(); @@ -1326,8 +1441,8 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let cursor = cursors.get_mut(cursor_id).unwrap(); - return_if_io!(cursor.next()); + let cursor = sorter_cursors.get_mut(cursor_id).unwrap(); + cursor.next(); if !cursor.is_empty() { state.pc = pc_if_next.to_offset_int(); } else { @@ -1875,7 +1990,7 @@ impl Program { record_reg, flag: _, } => { - let cursor = cursors.get_mut(cursor).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor).unwrap(); let record = match &state.registers[*record_reg] { OwnedValue::Record(r) => r, _ => unreachable!("Not a record! Cannot insert a non record value."), @@ -1885,7 +2000,7 @@ impl Program { state.pc += 1; } Insn::InsertAwait { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); cursor.wait_for_completion()?; // Only update last_insert_rowid for regular table inserts, not schema modifications if cursor.root_page() != 1 { @@ -1901,19 +2016,19 @@ impl Program { state.pc += 1; } Insn::DeleteAsync { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); return_if_io!(cursor.delete()); state.pc += 1; } Insn::DeleteAwait { cursor_id } => { - let cursor = cursors.get_mut(cursor_id).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor_id).unwrap(); cursor.wait_for_completion()?; state.pc += 1; } Insn::NewRowid { cursor, rowid_reg, .. } => { - let cursor = cursors.get_mut(cursor).unwrap(); + let cursor = btree_table_cursors.get_mut(cursor).unwrap(); // TODO: make io handle rng let rowid = return_if_io!(get_new_rowid(cursor, thread_rng())); state.registers[*rowid_reg] = OwnedValue::Integer(rowid); @@ -1939,7 +2054,13 @@ impl Program { rowid_reg, target_pc, } => { - let cursor = cursors.get_mut(cursor).unwrap(); + let cursor = must_be_btree_cursor!( + *cursor, + self.cursor_ref, + btree_table_cursors, + btree_index_cursors, + "NotExists" + ); let exists = return_if_io!(cursor.exists(&state.registers[*rowid_reg])); if exists { state.pc += 1; @@ -1954,12 +2075,15 @@ impl Program { cursor_id, root_page, } => { - let cursor = Box::new(BTreeCursor::new( - pager.clone(), - *root_page, - self.database_header.clone(), - )); - cursors.insert(*cursor_id, cursor); + let (_, cursor_type) = self.cursor_ref.get(*cursor_id).unwrap(); + let is_index = cursor_type.is_index(); + let cursor = + BTreeCursor::new(pager.clone(), *root_page, self.database_header.clone()); + if is_index { + btree_index_cursors.insert(*cursor_id, cursor); + } else { + btree_table_cursors.insert(*cursor_id, cursor); + } state.pc += 1; } Insn::OpenWriteAwait {} => { @@ -1991,7 +2115,21 @@ impl Program { state.pc += 1; } Insn::Close { cursor_id } => { - cursors.remove(cursor_id); + let (_, cursor_type) = self.cursor_ref.get(*cursor_id).unwrap(); + match cursor_type { + CursorType::BTreeTable(_) => { + let _ = btree_table_cursors.remove(cursor_id); + } + CursorType::BTreeIndex(_) => { + let _ = btree_index_cursors.remove(cursor_id); + } + CursorType::Pseudo(_) => { + let _ = pseudo_cursors.remove(cursor_id); + } + CursorType::Sorter => { + let _ = sorter_cursors.remove(cursor_id); + } + } state.pc += 1; } Insn::IsNull { src, target_pc } => { @@ -2022,7 +2160,7 @@ impl Program { } } -fn get_new_rowid(cursor: &mut Box, mut rng: R) -> Result> { +fn get_new_rowid(cursor: &mut BTreeCursor, mut rng: R) -> Result> { match cursor.seek_to_last()? { CursorResult::Ok(()) => {} CursorResult::IO => return Ok(CursorResult::IO), @@ -3076,90 +3214,6 @@ mod tests { use rand::{rngs::mock::StepRng, thread_rng}; use std::{cell::Ref, collections::HashMap, rc::Rc}; - mock! { - Cursor { - fn seek_to_last(&mut self) -> Result>; - fn seek<'a>(&mut self, key: SeekKey<'a>, op: SeekOp) -> Result>; - fn rowid(&self) -> Result>; - fn seek_rowid(&mut self, rowid: u64) -> Result>; - } - } - - impl Cursor for MockCursor { - fn root_page(&self) -> usize { - unreachable!() - } - - fn seek_to_last(&mut self) -> Result> { - self.seek_to_last() - } - - fn rowid(&self) -> Result> { - self.rowid() - } - - fn seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result> { - self.seek(key, op) - } - - fn rewind(&mut self) -> Result> { - unimplemented!() - } - - fn next(&mut self) -> Result> { - unimplemented!() - } - - fn record(&self) -> Result>> { - unimplemented!() - } - - fn is_empty(&self) -> bool { - unimplemented!() - } - - fn set_null_flag(&mut self, _flag: bool) { - unimplemented!() - } - - fn get_null_flag(&self) -> bool { - unimplemented!() - } - - fn insert( - &mut self, - _key: &OwnedValue, - _record: &OwnedRecord, - _is_leaf: bool, - ) -> Result> { - unimplemented!() - } - - fn delete(&mut self) -> Result> { - unimplemented!() - } - - fn wait_for_completion(&mut self) -> Result<()> { - unimplemented!() - } - - fn exists(&mut self, _key: &OwnedValue) -> Result> { - unimplemented!() - } - - fn btree_create(&mut self, _flags: usize) -> u32 { - unimplemented!() - } - - fn last(&mut self) -> Result> { - todo!() - } - - fn prev(&mut self) -> Result> { - todo!() - } - } - #[test] fn test_get_new_rowid() -> Result<()> { // Test case 0: Empty table diff --git a/core/vdbe/sorter.rs b/core/vdbe/sorter.rs index e3962b09..d23a3bef 100644 --- a/core/vdbe/sorter.rs +++ b/core/vdbe/sorter.rs @@ -1,13 +1,9 @@ -use crate::{ - types::{Cursor, CursorResult, OwnedRecord, OwnedValue, SeekKey, SeekOp}, - Result, -}; -use std::cell::{Ref, RefCell}; +use crate::types::OwnedRecord; use std::cmp::Ordering; pub struct Sorter { records: Vec, - current: RefCell>, + current: Option, order: Vec, } @@ -15,23 +11,15 @@ impl Sorter { pub fn new(order: Vec) -> Self { Self { records: Vec::new(), - current: RefCell::new(None), + current: None, order, } } -} - -impl Cursor for Sorter { - fn is_empty(&self) -> bool { - self.current.borrow().is_none() - } - - fn root_page(&self) -> usize { - unreachable!() + pub fn is_empty(&self) -> bool { + self.current.is_none() } - // We do the sorting here since this is what is called by the SorterSort instruction - fn rewind(&mut self) -> Result> { + pub fn sort(&mut self) { self.records.sort_by(|a, b| { let cmp_by_idx = |idx: usize, ascending: bool| { let a = &a.values[idx]; @@ -55,73 +43,14 @@ impl Cursor for Sorter { self.records.reverse(); self.next() } - - fn next(&mut self) -> Result> { - let mut c = self.current.borrow_mut(); - *c = self.records.pop(); - Ok(CursorResult::Ok(())) - } - - fn wait_for_completion(&mut self) -> Result<()> { - Ok(()) + pub fn next(&mut self) { + self.current = self.records.pop(); } - - fn rowid(&self) -> Result> { - todo!(); - } - - fn seek(&mut self, _: SeekKey<'_>, _: SeekOp) -> Result> { - unimplemented!(); - } - - fn seek_to_last(&mut self) -> Result> { - unimplemented!(); - } - - fn record(&self) -> Result>> { - let ret = self.current.borrow(); - // log::trace!("returning {:?}", ret); - Ok(ret) + pub fn record(&self) -> Option<&OwnedRecord> { + self.current.as_ref() } - fn insert( - &mut self, - key: &OwnedValue, - record: &OwnedRecord, - moved_before: bool, - ) -> Result> { - let _ = key; - let _ = moved_before; + pub fn insert(&mut self, record: &OwnedRecord) { self.records.push(OwnedRecord::new(record.values.to_vec())); - Ok(CursorResult::Ok(())) - } - - fn delete(&mut self) -> Result> { - unimplemented!() - } - - fn set_null_flag(&mut self, _flag: bool) { - todo!(); - } - - fn get_null_flag(&self) -> bool { - false - } - - fn exists(&mut self, key: &OwnedValue) -> Result> { - let _ = key; - todo!() - } - - fn btree_create(&mut self, _flags: usize) -> u32 { - unreachable!("Why did you try to build a new tree with a sorter??? Stand up, open the door and take a walk for 30 min to come back with a better plan."); - } - - fn last(&mut self) -> Result> { - todo!() - } - - fn prev(&mut self) -> Result> { - todo!() } } From 92888c54e9d8ecd0811ff08a1703d2654776c161 Mon Sep 17 00:00:00 2001 From: Jussi Saurio Date: Sat, 11 Jan 2025 17:41:10 +0200 Subject: [PATCH 4/4] remove get_new_rowid() specific tests that require mocking --- core/vdbe/mod.rs | 88 +++--------------------------------------------- 1 file changed, 4 insertions(+), 84 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 7b62ecae..bcdd6046 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -3198,96 +3198,16 @@ fn exec_math_log(arg: &OwnedValue, base: Option<&OwnedValue>) -> OwnedValue { #[cfg(test)] mod tests { - use crate::{ - types::{SeekKey, SeekOp}, - vdbe::exec_replace, - }; + use crate::vdbe::exec_replace; use super::{ exec_abs, exec_char, exec_hex, exec_if, exec_instr, exec_length, exec_like, exec_lower, exec_ltrim, exec_max, exec_min, exec_nullif, exec_quote, exec_random, exec_randomblob, exec_round, exec_rtrim, exec_sign, exec_soundex, exec_substring, exec_trim, exec_typeof, - exec_unhex, exec_unicode, exec_upper, exec_zeroblob, execute_sqlite_version, get_new_rowid, - AggContext, Cursor, CursorResult, LimboError, OwnedRecord, OwnedValue, Result, + exec_unhex, exec_unicode, exec_upper, exec_zeroblob, execute_sqlite_version, AggContext, + OwnedValue, }; - use mockall::{mock, predicate}; - use rand::{rngs::mock::StepRng, thread_rng}; - use std::{cell::Ref, collections::HashMap, rc::Rc}; - - #[test] - fn test_get_new_rowid() -> Result<()> { - // Test case 0: Empty table - let mut mock = MockCursor::new(); - mock.expect_seek_to_last() - .return_once(|| Ok(CursorResult::Ok(()))); - mock.expect_rowid().return_once(|| Ok(None)); - - let result = get_new_rowid(&mut (Box::new(mock) as Box), thread_rng())?; - assert_eq!( - result, - CursorResult::Ok(1), - "For an empty table, rowid should be 1" - ); - - // Test case 1: Normal case, rowid within i64::MAX - let mut mock = MockCursor::new(); - mock.expect_seek_to_last() - .return_once(|| Ok(CursorResult::Ok(()))); - mock.expect_rowid().return_once(|| Ok(Some(100))); - - let result = get_new_rowid(&mut (Box::new(mock) as Box), thread_rng())?; - assert_eq!(result, CursorResult::Ok(101)); - - // Test case 2: Rowid exceeds i64::MAX, need to generate random rowid - let mut mock = MockCursor::new(); - mock.expect_seek_to_last() - .return_once(|| Ok(CursorResult::Ok(()))); - mock.expect_rowid() - .return_once(|| Ok(Some(i64::MAX as u64))); - mock.expect_seek() - .with(predicate::always(), predicate::always()) - .returning(|rowid, _| { - if rowid == SeekKey::TableRowId(50) { - Ok(CursorResult::Ok(false)) - } else { - Ok(CursorResult::Ok(true)) - } - }); - - // Mock the random number generation - let new_rowid = - get_new_rowid(&mut (Box::new(mock) as Box), StepRng::new(1, 1))?; - assert_eq!(new_rowid, CursorResult::Ok(50)); - - // Test case 3: IO error - let mut mock = MockCursor::new(); - mock.expect_seek_to_last() - .return_once(|| Ok(CursorResult::Ok(()))); - mock.expect_rowid() - .return_once(|| Ok(Some(i64::MAX as u64))); - mock.expect_seek() - .with(predicate::always(), predicate::always()) - .return_once(|_, _| Ok(CursorResult::IO)); - - let result = get_new_rowid(&mut (Box::new(mock) as Box), thread_rng()); - assert!(matches!(result, Ok(CursorResult::IO))); - - // Test case 4: Failure to generate new rowid - let mut mock = MockCursor::new(); - mock.expect_seek_to_last() - .return_once(|| Ok(CursorResult::Ok(()))); - mock.expect_rowid() - .return_once(|| Ok(Some(i64::MAX as u64))); - mock.expect_seek() - .with(predicate::always(), predicate::always()) - .returning(|_, _| Ok(CursorResult::Ok(true))); - - // Mock the random number generation - let result = get_new_rowid(&mut (Box::new(mock) as Box), StepRng::new(1, 1)); - assert!(matches!(result, Err(LimboError::InternalError(_)))); - - Ok(()) - } + use std::{collections::HashMap, rc::Rc}; #[test] fn test_length() {