Add functions for constructing and translating Insertions

core/translate/insert.rs

@@ -1066,6 +1066,268 @@ struct ColMapping<'a> {
     register: usize,
 }
 
+/// Resolves how each column in a table should be populated during an INSERT.
+/// Returns an [Insertion] struct that contains the key and record for the insertion.
+fn build_insertion<'a>(
+    program: &mut ProgramBuilder,
+    table: &'a Table,
+    columns: &Option<DistinctNames>,
+    num_values: usize,
+) -> Result<Insertion<'a>> {
+    let table_columns = table.columns();
+    let rowid_register = program.alloc_register();
+    let mut insertion_key = InsertionKey::Autogenerated {
+        register: rowid_register,
+    };
+    let mut column_mappings = table
+        .columns()
+        .iter()
+        .map(|c| ColMapping {
+            column: c,
+            value_index: None,
+            register: program.alloc_register(),
+        })
+        .collect::<Vec<_>>();
+
+    if columns.is_none() {
+        // Case 1: No columns specified - map values to columns in order
+        let mut value_idx = 0;
+        for (i, col) in table_columns.iter().enumerate() {
+            if value_idx == num_values {
+                // If there are less values than columns, the rest will have value_index = None,
+                // meaning NULLs will be emitted for them.
+                break;
+            }
+            if col.hidden {
+                // Hidden columns are not taken into account.
+                continue;
+            }
+            if col.is_rowid_alias {
+                insertion_key = InsertionKey::RowidAlias(ColMapping {
+                    column: col,
+                    value_index: Some(value_idx),
+                    register: rowid_register,
+                });
+            } else {
+                column_mappings[i].value_index = Some(value_idx);
+            }
+            value_idx += 1;
+        }
+
+        if value_idx != num_values {
+            crate::bail_parse_error!(
+                "table {} has {} columns but {} values were supplied",
+                &table.get_name(),
+                num_values,
+                value_idx
+            );
+        }
+    } else {
+        // Case 2: Columns specified - map named columns to their values
+        // Map each named column to its value index
+        for (value_index, column_name) in columns.as_ref().unwrap().iter().enumerate() {
+            let column_name = normalize_ident(column_name.as_str());
+            if let Some((idx_in_table, col_in_table)) = table.get_column_by_name(&column_name) {
+                // Named column
+                if col_in_table.is_rowid_alias {
+                    insertion_key = InsertionKey::RowidAlias(ColMapping {
+                        column: col_in_table,
+                        value_index: Some(value_index),
+                        register: rowid_register,
+                    });
+                } else {
+                    column_mappings[idx_in_table].value_index = Some(value_index);
+                }
+            } else if column_name == ROWID {
+                // Explicit use of the 'rowid' keyword
+                if let Some(col_in_table) = table.columns().iter().find(|c| c.is_rowid_alias) {
+                    insertion_key = InsertionKey::RowidAlias(ColMapping {
+                        column: col_in_table,
+                        value_index: Some(value_index),
+                        register: rowid_register,
+                    });
+                } else {
+                    insertion_key = InsertionKey::LiteralRowid {
+                        value_index: Some(value_index),
+                        register: rowid_register,
+                    };
+                }
+            } else {
+                crate::bail_parse_error!(
+                    "table {} has no column named {}",
+                    &table.get_name(),
+                    column_name
+                );
+            }
+        }
+    }
+
+    Ok(Insertion {
+        key: insertion_key,
+        col_mappings: column_mappings,
+        record_reg: program.alloc_register(),
+    })
+}
+
+/// Populates the column registers with values for multiple rows.
+/// This is used for INSERT INTO <table> VALUES (...), (...), ... or INSERT INTO <table> SELECT ...
+/// which use either a coroutine or an ephemeral table as the value source.
+fn translate_rows_multiple<'short, 'long: 'short>(
+    program: &mut ProgramBuilder,
+    insertion: &'short Insertion<'long>,
+    yield_reg: usize,
+    resolver: &Resolver,
+    temp_table_ctx: &Option<TempTableCtx>,
+) -> Result<()> {
+    if let Some(ref temp_table_ctx) = temp_table_ctx {
+        // Rewind loop to read from ephemeral table
+        program.emit_insn(Insn::Rewind {
+            cursor_id: temp_table_ctx.cursor_id,
+            pc_if_empty: temp_table_ctx.loop_end_label,
+        });
+        program.preassign_label_to_next_insn(temp_table_ctx.loop_start_label);
+    }
+    let translate_value_fn =
+        |prg: &mut ProgramBuilder, value_index: usize, column_register: usize| {
+            if let Some(temp_table_ctx) = temp_table_ctx {
+                prg.emit_column(temp_table_ctx.cursor_id, value_index, column_register);
+            } else {
+                prg.emit_insn(Insn::Copy {
+                    src_reg: yield_reg + value_index,
+                    dst_reg: column_register,
+                    extra_amount: 0,
+                });
+            }
+            Ok(())
+        };
+    translate_rows_base(program, insertion, translate_value_fn, resolver)
+}
+/// Populates the column registers with values for a single row
+#[allow(clippy::too_many_arguments)]
+fn translate_rows_single(
+    program: &mut ProgramBuilder,
+    value: &[Expr],
+    insertion: &Insertion,
+    resolver: &Resolver,
+) -> Result<()> {
+    let translate_value_fn =
+        |prg: &mut ProgramBuilder, value_index: usize, column_register: usize| -> Result<()> {
+            translate_expr_no_constant_opt(
+                prg,
+                None,
+                value.get(value_index).unwrap_or_else(|| {
+                    panic!("value index out of bounds: {value_index} for value: {value:?}")
+                }),
+                column_register,
+                resolver,
+                NoConstantOptReason::RegisterReuse,
+            )?;
+            Ok(())
+        };
+    translate_rows_base(program, insertion, translate_value_fn, resolver)
+}
+
+/// Translate the key and the columns of the insertion.
+/// This function is called by both [translate_rows_single] and [translate_rows_multiple],
+/// each providing a different [translate_value_fn] implementation, because for multiple rows
+/// we need to emit the values in a loop, from either an ephemeral table or a coroutine,
+/// whereas for the single row the translation happens in a single pass without looping.
+fn translate_rows_base<'short, 'long: 'short>(
+    program: &mut ProgramBuilder,
+    insertion: &'short Insertion<'long>,
+    mut translate_value_fn: impl FnMut(&mut ProgramBuilder, usize, usize) -> Result<()>,
+    resolver: &Resolver,
+) -> Result<()> {
+    translate_key(program, insertion, &mut translate_value_fn, resolver)?;
+    for col in insertion.col_mappings.iter() {
+        translate_column(
+            program,
+            col.column,
+            col.register,
+            col.value_index,
+            &mut translate_value_fn,
+            resolver,
+        )?;
+    }
+
+    Ok(())
+}
+
+/// Translate the [InsertionKey].
+fn translate_key(
+    program: &mut ProgramBuilder,
+    insertion: &Insertion,
+    mut translate_value_fn: impl FnMut(&mut ProgramBuilder, usize, usize) -> Result<()>,
+    resolver: &Resolver,
+) -> Result<()> {
+    match &insertion.key {
+        InsertionKey::RowidAlias(rowid_alias_column) => translate_column(
+            program,
+            rowid_alias_column.column,
+            rowid_alias_column.register,
+            rowid_alias_column.value_index,
+            &mut translate_value_fn,
+            resolver,
+        ),
+        InsertionKey::LiteralRowid {
+            value_index,
+            register,
+        } => translate_column(
+            program,
+            ROWID_COLUMN,
+            *register,
+            *value_index,
+            &mut translate_value_fn,
+            resolver,
+        ),
+        InsertionKey::Autogenerated { .. } => Ok(()), // will be populated later
+    }
+}
+
+fn translate_column(
+    program: &mut ProgramBuilder,
+    column: &Column,
+    column_register: usize,
+    value_index: Option<usize>,
+    translate_value_fn: &mut impl FnMut(&mut ProgramBuilder, usize, usize) -> Result<()>,
+    resolver: &Resolver,
+) -> Result<()> {
+    if let Some(value_index) = value_index {
+        translate_value_fn(program, value_index, column_register)?;
+    } else if column.is_rowid_alias {
+        // Although a non-NULL integer key is used for the insertion key,
+        // the rowid alias column is emitted as NULL.
+        program.emit_insn(Insn::SoftNull {
+            reg: column_register,
+        });
+    } else if column.hidden {
+        // Emit NULL for not-explicitly-mentioned hidden columns, even ignoring DEFAULT.
+        program.emit_insn(Insn::Null {
+            dest: column_register,
+            dest_end: None,
+        });
+    } else if let Some(default_expr) = column.default.as_ref() {
+        translate_expr(program, None, default_expr, column_register, resolver)?;
+    } else {
+        let nullable = !column.notnull && !column.primary_key && !column.unique;
+        if !nullable {
+            crate::bail_parse_error!(
+                "column {} is not nullable",
+                column
+                    .name
+                    .as_ref()
+                    .expect("column name must be present")
+                    .as_str()
+            );
+        }
+        program.emit_insn(Insn::Null {
+            dest: column_register,
+            dest_end: None,
+        });
+    }
+    Ok(())
+}
+
 // TODO: comeback here later to apply the same improvements on select
 fn translate_virtual_table_insert(
     mut program: ProgramBuilder,