diff --git a/core/translate/insert.rs b/core/translate/insert.rs index 11724e742..2a143b166 100644 --- a/core/translate/insert.rs +++ b/core/translate/insert.rs @@ -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::(); + 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::(); - 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::(); - 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,