diff --git a/core/btree.rs b/core/btree.rs index 01c5b331e..0f644ebad 100644 --- a/core/btree.rs +++ b/core/btree.rs @@ -36,6 +36,7 @@ pub struct Cursor { pager: Arc, root_page: usize, page: RefCell>>, + rowid: RefCell>, record: RefCell>, } @@ -45,6 +46,7 @@ impl Cursor { pager, root_page, page: RefCell::new(None), + rowid: RefCell::new(None), record: RefCell::new(None), } } @@ -56,14 +58,16 @@ impl Cursor { pub fn rewind(&mut self) -> Result<()> { let mem_page = MemPage::new(None, self.root_page, 0); self.page.replace(Some(Arc::new(mem_page))); - let record = self.get_next_record()?; - self.record.replace(record); + let (rowid, next) = self.get_next_record()?; + self.rowid.replace(rowid); + self.record.replace(next); Ok(()) } pub fn next(&mut self) -> Result> { let result = self.record.take(); - let next = self.get_next_record()?; + let (rowid, next) = self.get_next_record()?; + self.rowid.replace(rowid); self.record.replace(next); Ok(result) } @@ -73,6 +77,10 @@ impl Cursor { Ok(()) } + pub fn rowid(&self) -> Result>> { + Ok(self.rowid.borrow()) + } + pub fn record(&self) -> Result>> { Ok(self.record.borrow()) } @@ -81,7 +89,7 @@ impl Cursor { self.record.borrow().is_some() } - fn get_next_record(&mut self) -> Result> { + fn get_next_record(&mut self) -> Result<(Option, Option)> { loop { let mem_page = { let mem_page = self.page.borrow(); @@ -96,7 +104,7 @@ impl Cursor { self.page.replace(Some(parent.clone())); continue; } - None => return Ok(None), + None => return Ok((None, None)), } } let cell = &page.cells[mem_page.cell_idx()]; @@ -114,7 +122,7 @@ impl Cursor { BTreeCell::TableLeafCell(TableLeafCell { _rowid, _payload }) => { mem_page.advance(); let record = crate::sqlite3_ondisk::read_record(_payload)?; - return Ok(Some(record)); + return Ok((Some(*_rowid), Some(record))); } } } diff --git a/core/schema.rs b/core/schema.rs index 7853b2110..b1ff248ad 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -71,7 +71,17 @@ impl Table { } None => Type::Null, }; - cols.push(Column { name, ty }); + let primary_key = column.constraints.iter().any(|c| { + matches!( + c.constraint, + sqlite3_parser::ast::ColumnConstraint::PrimaryKey { .. } + ) + }); + cols.push(Column { + name, + ty, + primary_key, + }); } } CreateTableBody::AsSelect(_) => todo!(), @@ -116,6 +126,7 @@ impl Table { pub struct Column { pub name: String, pub ty: Type, + pub primary_key: bool, } pub enum Type { @@ -149,22 +160,27 @@ pub fn sqlite_schema_table() -> Table { Column { name: "type".to_string(), ty: Type::Text, + primary_key: false, }, Column { name: "name".to_string(), ty: Type::Text, + primary_key: false, }, Column { name: "tbl_name".to_string(), ty: Type::Text, + primary_key: false, }, Column { name: "rootpage".to_string(), ty: Type::Integer, + primary_key: false, }, Column { name: "sql".to_string(), ty: Type::Text, + primary_key: false, }, ], } diff --git a/core/translate.rs b/core/translate.rs index 223761b40..c8b2b0b82 100644 --- a/core/translate.rs +++ b/core/translate.rs @@ -1,5 +1,5 @@ -use crate::vdbe::{Insn, Program, ProgramBuilder}; use crate::schema::Schema; +use crate::vdbe::{Insn, Program, ProgramBuilder}; use anyhow::Result; use sqlite3_parser::ast::{OneSelect, Select, Stmt}; @@ -154,13 +154,20 @@ fn translate_columns( sqlite3_parser::ast::Expr::Variable(_) => todo!(), }, sqlite3_parser::ast::ResultColumn::Star => { - for i in 0..table.unwrap().columns.len() { + for (i, col) in table.unwrap().columns.iter().enumerate() { let dest = program.alloc_register(); - program.emit_insn(Insn::Column { - column: i, - dest, - cursor_id: cursor_id.unwrap(), - }); + if col.primary_key { + program.emit_insn(Insn::RowId { + cursor_id: cursor_id.unwrap(), + dest, + }); + } else { + program.emit_insn(Insn::Column { + column: i, + dest, + cursor_id: cursor_id.unwrap(), + }); + } } } sqlite3_parser::ast::ResultColumn::TableStar(_) => todo!(), @@ -168,4 +175,4 @@ fn translate_columns( } let register_end = program.next_free_register(); (register_start, register_end) -} \ No newline at end of file +} diff --git a/core/vdbe.rs b/core/vdbe.rs index a730fecd3..87c63a973 100644 --- a/core/vdbe.rs +++ b/core/vdbe.rs @@ -78,6 +78,12 @@ pub enum Insn { value: i64, dest: usize, }, + + // Read the rowid of the current row. + RowId { + cursor_id: CursorID, + dest: usize, + }, } pub struct ProgramBuilder { @@ -267,6 +273,15 @@ impl Program { state.registers[*dest] = Value::Integer(*value); state.pc += 1; } + Insn::RowId { cursor_id, dest } => { + let cursor = state.cursors.get_mut(cursor_id).unwrap(); + if let Some(ref rowid) = *cursor.rowid()? { + state.registers[*dest] = Value::Integer(*rowid as i64); + } else { + todo!(); + } + state.pc += 1; + } } } } @@ -354,6 +369,7 @@ fn insn_to_str(addr: usize, insn: &Insn) -> String { Insn::Integer { value, dest } => { ("Integer", *dest, *value as usize, 0, "", 0, "".to_string()) } + Insn::RowId { cursor_id, dest } => ("RowId", *cursor_id, *dest, 0, "", 0, "".to_string()), }; format!( "{:<4} {:<13} {:<4} {:<4} {:<4} {:<13} {:<2} {}",