merge autoincrement into translate insert

This commit is contained in:
Pavan-Nambi
2025-09-13 21:10:51 +05:30
parent 0effb981e6
commit 255cfb10e6

View File

@@ -391,6 +391,70 @@ pub fn translate_insert(
let key_ready_for_uniqueness_check_label = program.allocate_label();
let key_generation_label = program.allocate_label();
let mut autoincrement_meta = None;
if btree_table.has_autoincrement {
let seq_table = schema.get_btree_table("sqlite_sequence").ok_or_else(|| {
crate::error::LimboError::InternalError("sqlite_sequence table not found".to_string())
})?;
let seq_cursor_id = program.alloc_cursor_id(CursorType::BTreeTable(seq_table.clone()));
program.emit_insn(Insn::OpenWrite {
cursor_id: seq_cursor_id,
root_page: seq_table.root_page.into(),
db: 0,
});
let table_name_reg = program.emit_string8_new_reg(btree_table.name.clone());
let r_seq = program.alloc_register();
let r_seq_rowid = program.alloc_register();
autoincrement_meta = Some((seq_cursor_id, r_seq, r_seq_rowid, table_name_reg));
program.emit_insn(Insn::Integer {
dest: r_seq,
value: 0,
});
program.emit_insn(Insn::Null {
dest: r_seq_rowid,
dest_end: None,
});
let loop_start_label = program.allocate_label();
let loop_end_label = program.allocate_label();
let found_label = program.allocate_label();
program.emit_insn(Insn::Rewind {
cursor_id: seq_cursor_id,
pc_if_empty: loop_end_label,
});
program.preassign_label_to_next_insn(loop_start_label);
let name_col_reg = program.alloc_register();
program.emit_column_or_rowid(seq_cursor_id, 0, name_col_reg);
program.emit_insn(Insn::Ne {
lhs: table_name_reg,
rhs: name_col_reg,
target_pc: found_label,
flags: Default::default(),
collation: None,
});
program.emit_column_or_rowid(seq_cursor_id, 1, r_seq);
program.emit_insn(Insn::RowId {
cursor_id: seq_cursor_id,
dest: r_seq_rowid,
});
program.emit_insn(Insn::Goto {
target_pc: loop_end_label,
});
program.preassign_label_to_next_insn(found_label);
program.emit_insn(Insn::Next {
cursor_id: seq_cursor_id,
pc_if_next: loop_start_label,
});
program.preassign_label_to_next_insn(loop_end_label);
}
if has_user_provided_rowid {
let must_be_int_label = program.allocate_label();
@@ -415,13 +479,67 @@ pub fn translate_insert(
program.preassign_label_to_next_insn(key_generation_label);
if btree_table.has_autoincrement {
emit_autoincrement_logic(
&mut program,
schema,
&btree_table,
cursor_id,
insertion.key_register(),
)?;
let (_, r_seq, _, _) = autoincrement_meta.unwrap();
let r_max = program.alloc_register();
let dummy_reg = program.alloc_register();
program.emit_insn(Insn::NewRowid {
cursor: cursor_id,
rowid_reg: dummy_reg,
prev_largest_reg: r_max,
});
let r_max_is_bigger_label = program.allocate_label();
let continue_label = program.allocate_label();
program.emit_insn(Insn::Ge {
lhs: r_max,
rhs: r_seq,
target_pc: r_max_is_bigger_label,
flags: Default::default(),
collation: None,
});
program.emit_insn(Insn::Copy {
src_reg: r_seq,
dst_reg: insertion.key_register(),
extra_amount: 0,
});
program.emit_insn(Insn::Goto {
target_pc: continue_label,
});
program.preassign_label_to_next_insn(r_max_is_bigger_label);
program.emit_insn(Insn::Copy {
src_reg: r_max,
dst_reg: insertion.key_register(),
extra_amount: 0,
});
program.preassign_label_to_next_insn(continue_label);
let no_overflow_label = program.allocate_label();
let max_i64_reg = program.alloc_register();
program.emit_insn(Insn::Integer {
dest: max_i64_reg,
value: i64::MAX,
});
program.emit_insn(Insn::Ne {
lhs: insertion.key_register(),
rhs: max_i64_reg,
target_pc: no_overflow_label,
flags: Default::default(),
collation: None,
});
program.emit_insn(Insn::Halt {
err_code: crate::error::SQLITE_FULL,
description: "database or disk is full".to_string(),
});
program.preassign_label_to_next_insn(no_overflow_label);
program.emit_insn(Insn::AddImm {
register: insertion.key_register(),
value: 1,
});
} else {
program.emit_insn(Insn::NewRowid {
cursor: cursor_id,
@@ -691,13 +809,79 @@ pub fn translate_insert(
table_name: table_name.to_string(),
});
if btree_table.has_autoincrement && has_user_provided_rowid {
update_sequence_after_user_insert(
&mut program,
schema,
&btree_table,
insertion.key_register(),
)?;
if let Some((seq_cursor_id, r_seq, r_seq_rowid, table_name_reg)) = autoincrement_meta {
let no_update_needed_label = program.allocate_label();
program.emit_insn(Insn::Le {
lhs: insertion.key_register(),
rhs: r_seq,
target_pc: no_update_needed_label,
flags: Default::default(),
collation: None,
});
let record_reg = program.alloc_register();
let record_start_reg = program.alloc_registers(2);
program.emit_insn(Insn::Copy {
src_reg: table_name_reg,
dst_reg: record_start_reg,
extra_amount: 0,
});
program.emit_insn(Insn::Copy {
src_reg: insertion.key_register(),
dst_reg: record_start_reg + 1,
extra_amount: 0,
});
let seq_table = schema.get_btree_table("sqlite_sequence").unwrap();
let affinity_str = seq_table
.columns
.iter()
.map(|col| col.affinity().aff_mask())
.collect::<String>();
program.emit_insn(Insn::MakeRecord {
start_reg: record_start_reg,
count: 2,
dest_reg: record_reg,
index_name: None,
affinity_str: Some(affinity_str),
});
let update_existing_label = program.allocate_label();
let end_update_label = program.allocate_label();
program.emit_insn(Insn::NotNull {
reg: r_seq_rowid,
target_pc: update_existing_label,
});
program.emit_insn(Insn::NewRowid {
cursor: seq_cursor_id,
rowid_reg: r_seq_rowid,
prev_largest_reg: 0,
});
program.emit_insn(Insn::Insert {
cursor: seq_cursor_id,
key_reg: r_seq_rowid,
record_reg,
flag: InsertFlags::new(),
table_name: "sqlite_sequence".to_string(),
});
program.emit_insn(Insn::Goto {
target_pc: end_update_label,
});
program.preassign_label_to_next_insn(update_existing_label);
program.emit_insn(Insn::Insert {
cursor: seq_cursor_id,
key_reg: r_seq_rowid,
record_reg,
flag: InsertFlags(turso_parser::ast::ResolveType::Replace.bit_value() as u8),
table_name: "sqlite_sequence".to_string(),
});
program.preassign_label_to_next_insn(end_update_label);
program.preassign_label_to_next_insn(no_update_needed_label);
program.emit_insn(Insn::Close {
cursor_id: seq_cursor_id,
});
}
// Emit update in the CDC table if necessary (after the INSERT updated the table)
@@ -1204,342 +1388,7 @@ fn translate_virtual_table_insert(
Ok(program)
}
fn emit_autoincrement_logic(
program: &mut ProgramBuilder,
schema: &Schema,
table: &schema::BTreeTable,
main_cursor_id: usize,
key_register: usize,
) -> Result<()> {
let seq_table = schema.get_btree_table("sqlite_sequence").ok_or_else(|| {
crate::error::LimboError::InternalError("sqlite_sequence table not found".to_string())
})?;
let seq_cursor_id = program.alloc_cursor_id(CursorType::BTreeTable(seq_table.clone()));
program.emit_insn(Insn::OpenWrite {
cursor_id: seq_cursor_id,
root_page: seq_table.root_page.into(),
db: 0,
});
let table_name_reg = program.emit_string8_new_reg(table.name.clone());
let r_seq = program.alloc_register();
let r_max = program.alloc_register();
let r_seq_rowid = program.alloc_register();
program.emit_insn(Insn::Integer {
dest: r_seq,
value: 0,
});
program.emit_insn(Insn::Null {
dest: r_seq_rowid,
dest_end: None,
});
let loop_start_label = program.allocate_label();
let loop_end_label = program.allocate_label();
let found_label = program.allocate_label();
program.emit_insn(Insn::Rewind {
cursor_id: seq_cursor_id,
pc_if_empty: loop_end_label,
});
program.preassign_label_to_next_insn(loop_start_label);
let name_col_reg = program.alloc_register();
program.emit_column_or_rowid(seq_cursor_id, 0, name_col_reg);
program.emit_insn(Insn::Ne {
lhs: table_name_reg,
rhs: name_col_reg,
target_pc: found_label,
flags: Default::default(),
collation: None,
});
program.emit_column_or_rowid(seq_cursor_id, 1, r_seq);
program.emit_insn(Insn::RowId {
cursor_id: seq_cursor_id,
dest: r_seq_rowid,
});
program.emit_insn(Insn::Goto {
target_pc: loop_end_label,
});
program.preassign_label_to_next_insn(found_label);
program.emit_insn(Insn::Next {
cursor_id: seq_cursor_id,
pc_if_next: loop_start_label,
});
program.preassign_label_to_next_insn(loop_end_label);
let dummy_reg = program.alloc_register();
program.emit_insn(Insn::NewRowid {
cursor: main_cursor_id,
rowid_reg: dummy_reg,
prev_largest_reg: r_max,
});
let r_max_is_bigger_label = program.allocate_label();
let continue_label = program.allocate_label();
program.emit_insn(Insn::Ge {
lhs: r_max,
rhs: r_seq,
target_pc: r_max_is_bigger_label,
flags: Default::default(),
collation: None,
});
program.emit_insn(Insn::Copy {
src_reg: r_seq,
dst_reg: key_register,
extra_amount: 0,
});
program.emit_insn(Insn::Goto {
target_pc: continue_label,
});
program.preassign_label_to_next_insn(r_max_is_bigger_label);
program.emit_insn(Insn::Copy {
src_reg: r_max,
dst_reg: key_register,
extra_amount: 0,
});
program.preassign_label_to_next_insn(continue_label);
let no_overflow_label = program.allocate_label();
let max_i64_reg = program.alloc_register();
program.emit_insn(Insn::Integer {
dest: max_i64_reg,
value: i64::MAX,
});
program.emit_insn(Insn::Ne {
lhs: key_register,
rhs: max_i64_reg,
target_pc: no_overflow_label,
flags: Default::default(),
collation: None,
});
program.emit_insn(Insn::Halt {
err_code: crate::error::SQLITE_FULL,
description: "database or disk is full".to_string(),
});
program.preassign_label_to_next_insn(no_overflow_label);
program.emit_insn(Insn::AddImm {
register: key_register,
value: 1,
});
let record_reg = program.alloc_register();
let record_start_reg = program.alloc_registers(2);
program.emit_insn(Insn::Copy {
src_reg: table_name_reg,
dst_reg: record_start_reg,
extra_amount: 0,
});
program.emit_insn(Insn::Copy {
src_reg: key_register,
dst_reg: record_start_reg + 1,
extra_amount: 0,
});
let affinity_str = seq_table
.columns
.iter()
.map(|col| col.affinity().aff_mask())
.collect::<String>();
program.emit_insn(Insn::MakeRecord {
start_reg: record_start_reg,
count: 2,
dest_reg: record_reg,
index_name: None,
affinity_str: Some(affinity_str),
});
let update_existing_label = program.allocate_label();
let end_update_label = program.allocate_label();
program.emit_insn(Insn::NotNull {
reg: r_seq_rowid,
target_pc: update_existing_label,
});
program.emit_insn(Insn::NewRowid {
cursor: seq_cursor_id,
rowid_reg: r_seq_rowid,
prev_largest_reg: 0,
});
program.emit_insn(Insn::Insert {
cursor: seq_cursor_id,
key_reg: r_seq_rowid,
record_reg,
flag: InsertFlags::new(),
table_name: "sqlite_sequence".to_string(),
});
program.emit_insn(Insn::Goto {
target_pc: end_update_label,
});
program.preassign_label_to_next_insn(update_existing_label);
program.emit_insn(Insn::Insert {
cursor: seq_cursor_id,
key_reg: r_seq_rowid,
record_reg,
flag: InsertFlags(turso_parser::ast::ResolveType::Replace.bit_value() as u8),
table_name: "sqlite_sequence".to_string(),
});
program.preassign_label_to_next_insn(end_update_label);
program.emit_insn(Insn::Close {
cursor_id: seq_cursor_id,
});
Ok(())
}
fn update_sequence_after_user_insert(
program: &mut ProgramBuilder,
schema: &Schema,
table: &schema::BTreeTable,
user_key_register: usize,
) -> Result<()> {
let seq_table = schema.get_btree_table("sqlite_sequence").ok_or_else(|| {
crate::error::LimboError::InternalError("sqlite_sequence table not found".to_string())
})?;
let seq_cursor_id = program.alloc_cursor_id(CursorType::BTreeTable(seq_table.clone()));
program.emit_insn(Insn::OpenWrite {
cursor_id: seq_cursor_id,
root_page: seq_table.root_page.into(),
db: 0,
});
let table_name_reg = program.emit_string8_new_reg(table.name.clone());
let r_seq = program.alloc_register();
let r_seq_rowid = program.alloc_register();
program.emit_insn(Insn::Integer {
dest: r_seq,
value: 0,
});
program.emit_insn(Insn::Null {
dest: r_seq_rowid,
dest_end: None,
});
let loop_start_label = program.allocate_label();
let loop_end_label = program.allocate_label();
let found_label = program.allocate_label();
program.emit_insn(Insn::Rewind {
cursor_id: seq_cursor_id,
pc_if_empty: loop_end_label,
});
program.preassign_label_to_next_insn(loop_start_label);
let name_col_reg = program.alloc_register();
program.emit_column_or_rowid(seq_cursor_id, 0, name_col_reg);
program.emit_insn(Insn::Ne {
lhs: table_name_reg,
rhs: name_col_reg,
target_pc: found_label,
flags: Default::default(),
collation: None,
});
program.emit_column_or_rowid(seq_cursor_id, 1, r_seq);
program.emit_insn(Insn::RowId {
cursor_id: seq_cursor_id,
dest: r_seq_rowid,
});
program.emit_insn(Insn::Goto {
target_pc: loop_end_label,
});
program.preassign_label_to_next_insn(found_label);
program.emit_insn(Insn::Next {
cursor_id: seq_cursor_id,
pc_if_next: loop_start_label,
});
program.preassign_label_to_next_insn(loop_end_label);
let no_update_needed_label = program.allocate_label();
program.emit_insn(Insn::Le {
lhs: user_key_register,
rhs: r_seq,
target_pc: no_update_needed_label,
flags: Default::default(),
collation: None,
});
let record_reg = program.alloc_register();
let record_start_reg = program.alloc_registers(2);
program.emit_insn(Insn::Copy {
src_reg: table_name_reg,
dst_reg: record_start_reg,
extra_amount: 0,
});
program.emit_insn(Insn::Copy {
src_reg: user_key_register,
dst_reg: record_start_reg + 1,
extra_amount: 0,
});
let affinity_str = seq_table
.columns
.iter()
.map(|col| col.affinity().aff_mask())
.collect::<String>();
program.emit_insn(Insn::MakeRecord {
start_reg: record_start_reg,
count: 2,
dest_reg: record_reg,
index_name: None,
affinity_str: Some(affinity_str),
});
let update_existing_label = program.allocate_label();
let end_update_label = program.allocate_label();
program.emit_insn(Insn::NotNull {
reg: r_seq_rowid,
target_pc: update_existing_label,
});
program.emit_insn(Insn::NewRowid {
cursor: seq_cursor_id,
rowid_reg: r_seq_rowid,
prev_largest_reg: 0,
});
program.emit_insn(Insn::Insert {
cursor: seq_cursor_id,
key_reg: r_seq_rowid,
record_reg,
flag: InsertFlags::new(),
table_name: "sqlite_sequence".to_string(),
});
program.emit_insn(Insn::Goto {
target_pc: end_update_label,
});
// Update existing row
program.preassign_label_to_next_insn(update_existing_label);
program.emit_insn(Insn::Insert {
cursor: seq_cursor_id,
key_reg: r_seq_rowid,
record_reg,
flag: InsertFlags(turso_parser::ast::ResolveType::Replace.bit_value() as u8),
table_name: "sqlite_sequence".to_string(),
});
program.preassign_label_to_next_insn(end_update_label);
program.preassign_label_to_next_insn(no_update_needed_label);
program.emit_insn(Insn::Close {
cursor_id: seq_cursor_id,
});
Ok(())
}
/// makes sure that an AUTOINCREMENT table has a sequence row in `sqlite_sequence`, inserting one with 0 if missing.
fn ensure_sequence_initialized(
program: &mut ProgramBuilder,
schema: &Schema,
@@ -1550,6 +1399,7 @@ fn ensure_sequence_initialized(
})?;
let seq_cursor_id = program.alloc_cursor_id(CursorType::BTreeTable(seq_table.clone()));
program.emit_insn(Insn::OpenWrite {
cursor_id: seq_cursor_id,
root_page: seq_table.root_page.into(),
@@ -1566,9 +1416,13 @@ fn ensure_sequence_initialized(
cursor_id: seq_cursor_id,
pc_if_empty: insert_new_label,
});
program.preassign_label_to_next_insn(loop_start_label);
let name_col_reg = program.alloc_register();
program.emit_column_or_rowid(seq_cursor_id, 0, name_col_reg);
program.emit_insn(Insn::Eq {
lhs: table_name_reg,
rhs: name_col_reg,
@@ -1576,6 +1430,7 @@ fn ensure_sequence_initialized(
flags: Default::default(),
collation: None,
});
program.emit_insn(Insn::Next {
cursor_id: seq_cursor_id,
pc_if_next: loop_start_label,
@@ -1586,6 +1441,7 @@ fn ensure_sequence_initialized(
let record_reg = program.alloc_register();
let record_start_reg = program.alloc_registers(2);
let zero_reg = program.alloc_register();
program.emit_insn(Insn::Integer {
dest: zero_reg,
value: 0,
@@ -1596,6 +1452,7 @@ fn ensure_sequence_initialized(
dst_reg: record_start_reg,
extra_amount: 0,
});
program.emit_insn(Insn::Copy {
src_reg: zero_reg,
dst_reg: record_start_reg + 1,
@@ -1607,6 +1464,7 @@ fn ensure_sequence_initialized(
.iter()
.map(|c| c.affinity().aff_mask())
.collect();
program.emit_insn(Insn::MakeRecord {
start_reg: record_start_reg,
count: 2,