diff --git a/core/translate/insert.rs b/core/translate/insert.rs index bca56aae6..df8e14fd7 100644 --- a/core/translate/insert.rs +++ b/core/translate/insert.rs @@ -949,6 +949,123 @@ fn populate_column_registers( Ok(()) } +/// Represents how a table should be populated during an INSERT. +#[derive(Debug)] +struct Insertion<'a> { + /// The integer key ("rowid") provided to the VDBE. + key: InsertionKey<'a>, + /// The column values that will be fed to the MakeRecord instruction to insert the row. + /// If the table has a rowid alias column, it will also be included in this record, + /// but a NULL will be stored for it. + col_mappings: Vec>, + /// The register that will contain the record built using the MakeRecord instruction. + record_reg: usize, +} + +impl<'a> Insertion<'a> { + /// Return the register that contains the rowid. + pub fn key_register(&self) -> usize { + self.key.register() + } + + /// Return the first register of the values that used to build the record + /// for the main table insert. + pub fn first_col_register(&self) -> usize { + self.col_mappings + .first() + .expect("columns must be present") + .register + } + + /// Return the register that contains the record built using the MakeRecord instruction. + pub fn record_register(&self) -> usize { + self.record_reg + } + + /// Returns the column mapping for a given column name. + pub fn get_col_mapping_by_name(&self, name: &str) -> &ColMapping<'a> { + if let InsertionKey::RowidAlias(mapping) = &self.key { + // If the key is a rowid alias, a NULL is emitted as the column value, + // so we need to return the key mapping instead so that the non-NULL rowid is used + // for the index insert. + if mapping + .column + .name + .as_ref() + .is_some_and(|n| n.eq_ignore_ascii_case(name)) + { + return mapping; + } + } + self.col_mappings + .iter() + .find(|col| { + col.column + .name + .as_ref() + .is_some_and(|n| n.eq_ignore_ascii_case(name)) + }) + .unwrap_or_else(|| panic!("column {name} not found in insertion")) + } +} + +#[derive(Debug)] +enum InsertionKey<'a> { + /// Rowid is not provided by user and will be autogenerated. + Autogenerated { register: usize }, + /// Rowid is provided via the 'rowid' keyword. + LiteralRowid { + value_index: Option, + register: usize, + }, + /// Rowid is provided via a rowid alias column. + RowidAlias(ColMapping<'a>), +} + +impl InsertionKey<'_> { + fn register(&self) -> usize { + match self { + InsertionKey::Autogenerated { register } => *register, + InsertionKey::LiteralRowid { register, .. } => *register, + InsertionKey::RowidAlias(x) => x.register, + } + } + fn is_provided_by_user(&self) -> bool { + !matches!(self, InsertionKey::Autogenerated { .. }) + } + + fn column_name(&self) -> &str { + match self { + InsertionKey::RowidAlias(x) => x + .column + .name + .as_ref() + .expect("rowid alias column must be present") + .as_str(), + InsertionKey::LiteralRowid { .. } => ROWID, + InsertionKey::Autogenerated { .. } => ROWID, + } + } +} + +/// Represents how a column in a table should be populated during an INSERT. +/// In a vector of [ColMapping], the index of a given [ColMapping] is +/// the position of the column in the table. +#[derive(Debug)] +struct ColMapping<'a> { + /// Column definition + column: &'a Column, + /// Index of the value to use from a tuple in the insert statement. + /// This is needed because the values in the insert statement are not necessarily + /// in the same order as the columns in the table, nor do they necessarily contain + /// all of the columns in the table. + /// If None, a NULL will be emitted for the column, unless it has a default value. + /// A NULL rowid alias column's value will be autogenerated. + value_index: Option, + /// Register where the value will be stored for insertion into the table. + register: usize, +} + // TODO: comeback here later to apply the same improvements on select fn translate_virtual_table_insert( mut program: ProgramBuilder,