Implement reading primary key columns

core/btree.rs

@@ -36,6 +36,7 @@ pub struct Cursor {
     pager: Arc<Pager>,
     root_page: usize,
     page: RefCell<Option<Arc<MemPage>>>,
+    rowid: RefCell<Option<u64>>,
     record: RefCell<Option<Record>>,
 }
 
@@ -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<Option<Record>> {
         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<Ref<Option<u64>>> {
+        Ok(self.rowid.borrow())
+    }
+
     pub fn record(&self) -> Result<Ref<Option<Record>>> {
         Ok(self.record.borrow())
     }
@@ -81,7 +89,7 @@ impl Cursor {
         self.record.borrow().is_some()
     }
 
-    fn get_next_record(&mut self) -> Result<Option<Record>> {
+    fn get_next_record(&mut self) -> Result<(Option<u64>, Option<Record>)> {
         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)));
                 }
             }
         }

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,
             },
         ],
     }

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)
-}
+}

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}  {}",