From e4ebb6d9e13e47b0878dc09dd40e4d6c77437773 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 11:10:52 +0200 Subject: [PATCH 01/10] core/vdbe: Add ProgramState::get_vtab_cursor() helper --- core/vdbe/mod.rs | 33 ++++++++++++++------------------- 1 file changed, 14 insertions(+), 19 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 709129381..9cef73520 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -229,19 +229,6 @@ fn get_cursor_as_sorter_mut<'long, 'short>( cursor } -fn get_cursor_as_virtual_mut<'long, 'short>( - cursors: &'short mut RefMut<'long, Vec>>, - cursor_id: CursorID, -) -> &'short mut VTabOpaqueCursor { - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_virtual_mut(); - cursor -} - struct Bitfield([u64; N]); impl Bitfield { @@ -352,6 +339,17 @@ impl ProgramState { self.interrupted = false; self.parameters.clear(); } + + pub fn get_vtab_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut VTabOpaqueCursor { + let mut cursors = self.cursors.borrow_mut(); + let cursor = cursors + .get_mut(cursor_id) + .expect("cursor id out of bounds") + .as_mut() + .expect("cursor not allocated") + .as_virtual_mut(); + unsafe { std::mem::transmute(cursor) } + } } macro_rules! must_be_btree_cursor { @@ -893,8 +891,7 @@ impl Program { let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VFilter on non-virtual table cursor"); }; - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_virtual_mut(&mut cursors, *cursor_id); + let cursor = state.get_vtab_cursor(*cursor_id); let mut args = Vec::new(); for i in 0..*arg_count { args.push(state.registers[args_reg + i].clone()); @@ -915,8 +912,7 @@ impl Program { let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VColumn on non-virtual table cursor"); }; - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_virtual_mut(&mut cursors, *cursor_id); + let cursor = state.get_vtab_cursor(*cursor_id); state.registers[*dest] = virtual_table.column(cursor, *column)?; state.pc += 1; } @@ -981,8 +977,7 @@ impl Program { let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VNextAsync on non-virtual table cursor"); }; - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_virtual_mut(&mut cursors, *cursor_id); + let cursor = state.get_vtab_cursor(*cursor_id); let has_more = virtual_table.next(cursor)?; if has_more { state.pc = pc_if_next.to_offset_int(); From 06446b768b5050cb124454441b70eb4ec0f8bd7a Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 11:16:13 +0200 Subject: [PATCH 02/10] core/vdbe: Add ProgramState::get_sorter() helper --- core/vdbe/mod.rs | 39 +++++++++++++++++---------------------- 1 file changed, 17 insertions(+), 22 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 9cef73520..77de6e9be 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -216,19 +216,6 @@ fn get_cursor_as_pseudo_mut<'long, 'short>( cursor } -fn get_cursor_as_sorter_mut<'long, 'short>( - cursors: &'short mut RefMut<'long, Vec>>, - cursor_id: CursorID, -) -> &'short mut Sorter { - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_sorter_mut(); - cursor -} - struct Bitfield([u64; N]); impl Bitfield { @@ -340,6 +327,17 @@ impl ProgramState { self.parameters.clear(); } + pub fn get_sorter<'a>(&'a self, cursor_id: CursorID) -> &'a mut Sorter { + let mut cursors = self.cursors.borrow_mut(); + let cursor = cursors + .get_mut(cursor_id) + .expect("cursor id out of bounds") + .as_mut() + .expect("cursor not allocated") + .as_sorter_mut(); + unsafe { std::mem::transmute(cursor) } + } + pub fn get_vtab_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut VTabOpaqueCursor { let mut cursors = self.cursors.borrow_mut(); let cursor = cursors @@ -1081,7 +1079,7 @@ impl Program { } } CursorType::Sorter => { - let cursor = get_cursor_as_sorter_mut(&mut cursors, *cursor_id); + let cursor = state.get_sorter(*cursor_id); if let Some(record) = cursor.record() { state.registers[*dest] = record.get_value(*column).clone(); } else { @@ -1932,8 +1930,7 @@ impl Program { dest_reg, pseudo_cursor, } => { - let mut cursors = state.cursors.borrow_mut(); - let sorter_cursor = get_cursor_as_sorter_mut(&mut cursors, *cursor_id); + let sorter_cursor = state.get_sorter(*cursor_id); let record = match sorter_cursor.record() { Some(record) => record.clone(), None => { @@ -1942,6 +1939,7 @@ impl Program { } }; state.registers[*dest_reg] = OwnedValue::Record(record.clone()); + let mut cursors = state.cursors.borrow_mut(); let pseudo_cursor = get_cursor_as_pseudo_mut(&mut cursors, *pseudo_cursor); pseudo_cursor.insert(record); state.pc += 1; @@ -1950,8 +1948,7 @@ impl Program { cursor_id, record_reg, } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_sorter_mut(&mut cursors, *cursor_id); + let cursor = state.get_sorter(*cursor_id); let record = match &state.registers[*record_reg] { OwnedValue::Record(record) => record, _ => unreachable!("SorterInsert on non-record register"), @@ -1963,8 +1960,7 @@ impl Program { cursor_id, pc_if_empty, } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_sorter_mut(&mut cursors, *cursor_id); + let cursor = state.get_sorter(*cursor_id); if cursor.is_empty() { state.pc = pc_if_empty.to_offset_int(); } else { @@ -1977,8 +1973,7 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_sorter_mut(&mut cursors, *cursor_id); + let cursor = state.get_sorter(*cursor_id); cursor.next(); if cursor.has_more() { state.pc = pc_if_next.to_offset_int(); From 222808ab6caa853c684dbbe754e6ded37d2783e4 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 11:20:19 +0200 Subject: [PATCH 03/10] ore/vdbe: Add ProgramState::get_pseudo_cursor() helper --- core/vdbe/mod.rs | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 77de6e9be..e61dc5888 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -203,19 +203,6 @@ fn get_cursor_as_index_mut<'long, 'short>( cursor } -fn get_cursor_as_pseudo_mut<'long, 'short>( - cursors: &'short mut RefMut<'long, Vec>>, - cursor_id: CursorID, -) -> &'short mut PseudoCursor { - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_pseudo_mut(); - cursor -} - struct Bitfield([u64; N]); impl Bitfield { @@ -327,6 +314,17 @@ impl ProgramState { self.parameters.clear(); } + pub fn get_pseudo_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut PseudoCursor { + let mut cursors = self.cursors.borrow_mut(); + let cursor = cursors + .get_mut(cursor_id) + .expect("cursor id out of bounds") + .as_mut() + .expect("cursor not allocated") + .as_pseudo_mut(); + unsafe { std::mem::transmute(cursor) } + } + pub fn get_sorter<'a>(&'a self, cursor_id: CursorID) -> &'a mut Sorter { let mut cursors = self.cursors.borrow_mut(); let cursor = cursors @@ -1087,7 +1085,7 @@ impl Program { } } CursorType::Pseudo(_) => { - let cursor = get_cursor_as_pseudo_mut(&mut cursors, *cursor_id); + let cursor = state.get_pseudo_cursor(*cursor_id); if let Some(record) = cursor.record() { state.registers[*dest] = record.get_value(*column).clone(); } else { @@ -1939,8 +1937,7 @@ impl Program { } }; state.registers[*dest_reg] = OwnedValue::Record(record.clone()); - let mut cursors = state.cursors.borrow_mut(); - let pseudo_cursor = get_cursor_as_pseudo_mut(&mut cursors, *pseudo_cursor); + let pseudo_cursor = state.get_pseudo_cursor(*pseudo_cursor); pseudo_cursor.insert(record); state.pc += 1; } From 3aeb11b673c821966a36bf2801131b8532335598 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 11:23:25 +0200 Subject: [PATCH 04/10] core/vdbe: Add ProgramStatem::get_btree_{table,index}_cursor() helpers --- core/vdbe/mod.rs | 271 ++++++++++++++++++++++------------------------- 1 file changed, 129 insertions(+), 142 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index e61dc5888..206b66dbb 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -67,7 +67,7 @@ use rand::{thread_rng, Rng}; use regex::{Regex, RegexBuilder}; use sorter::Sorter; use std::borrow::BorrowMut; -use std::cell::{Cell, RefCell, RefMut}; +use std::cell::{Cell, RefCell}; use std::collections::HashMap; use std::ffi::c_void; use std::num::NonZero; @@ -177,32 +177,6 @@ impl RegexCache { } } -fn get_cursor_as_table_mut<'long, 'short>( - cursors: &'short mut RefMut<'long, Vec>>, - cursor_id: CursorID, -) -> &'short mut BTreeCursor { - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_table_mut(); - cursor -} - -fn get_cursor_as_index_mut<'long, 'short>( - cursors: &'short mut RefMut<'long, Vec>>, - cursor_id: CursorID, -) -> &'short mut BTreeCursor { - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_index_mut(); - cursor -} - struct Bitfield([u64; N]); impl Bitfield { @@ -314,6 +288,28 @@ impl ProgramState { self.parameters.clear(); } + pub fn get_btree_table_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut BTreeCursor { + let mut cursors = self.cursors.borrow_mut(); + let cursor = cursors + .get_mut(cursor_id) + .expect("cursor id out of bounds") + .as_mut() + .expect("cursor not allocated") + .as_table_mut(); + unsafe { std::mem::transmute(cursor) } + } + + pub fn get_btree_index_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut BTreeCursor { + let mut cursors = self.cursors.borrow_mut(); + let cursor = cursors + .get_mut(cursor_id) + .expect("cursor id out of bounds") + .as_mut() + .expect("cursor not allocated") + .as_index_mut(); + unsafe { std::mem::transmute(cursor) } + } + pub fn get_pseudo_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut PseudoCursor { let mut cursors = self.cursors.borrow_mut(); let cursor = cursors @@ -349,11 +345,11 @@ impl ProgramState { } macro_rules! must_be_btree_cursor { - ($cursor_id:expr, $cursor_ref:expr, $cursors:expr, $insn_name:expr) => {{ + ($cursor_id:expr, $cursor_ref:expr, $state:expr, $insn_name:expr) => {{ let (_, cursor_type) = $cursor_ref.get($cursor_id).unwrap(); let cursor = match cursor_type { - CursorType::BTreeTable(_) => get_cursor_as_table_mut(&mut $cursors, $cursor_id), - CursorType::BTreeIndex(_) => get_cursor_as_index_mut(&mut $cursors, $cursor_id), + CursorType::BTreeTable(_) => $state.get_btree_table_cursor($cursor_id), + CursorType::BTreeIndex(_) => $state.get_btree_index_cursor($cursor_id), CursorType::Pseudo(_) => panic!("{} on pseudo cursor", $insn_name), CursorType::Sorter => panic!("{} on sorter cursor", $insn_name), CursorType::VirtualTable(_) => panic!("{} on virtual table cursor", $insn_name), @@ -483,9 +479,8 @@ impl Program { state.pc += 1; } Insn::NullRow { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "NullRow"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NullRow"); cursor.set_null_flag(true); state.pc += 1; } @@ -995,16 +990,14 @@ impl Program { state.pc += 1; } Insn::RewindAsync { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "RewindAsync"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "RewindAsync"); return_if_io!(cursor.rewind()); state.pc += 1; } Insn::LastAsync { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "LastAsync"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "LastAsync"); return_if_io!(cursor.last()); state.pc += 1; } @@ -1013,9 +1006,8 @@ impl Program { pc_if_empty, } => { assert!(pc_if_empty.is_offset()); - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "LastAwait"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "LastAwait"); cursor.wait_for_completion()?; if cursor.is_empty() { state.pc = pc_if_empty.to_offset_int(); @@ -1028,9 +1020,8 @@ impl Program { pc_if_empty, } => { assert!(pc_if_empty.is_offset()); - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "RewindAwait"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "RewindAwait"); cursor.wait_for_completion()?; if cursor.is_empty() { state.pc = pc_if_empty.to_offset_int(); @@ -1043,11 +1034,10 @@ impl Program { column, dest, } => { - let mut cursors = state.cursors.borrow_mut(); if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() { - let index_cursor = get_cursor_as_index_mut(&mut cursors, index_cursor_id); + let index_cursor = state.get_btree_index_cursor(index_cursor_id); let rowid = index_cursor.rowid()?; - let table_cursor = get_cursor_as_table_mut(&mut cursors, table_cursor_id); + let table_cursor = state.get_btree_table_cursor(table_cursor_id); match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? { CursorResult::Ok(_) => {} CursorResult::IO => { @@ -1059,22 +1049,25 @@ impl Program { let (_, cursor_type) = self.cursor_ref.get(*cursor_id).unwrap(); match cursor_type { CursorType::BTreeTable(_) | CursorType::BTreeIndex(_) => { - let cursor = must_be_btree_cursor!( - *cursor_id, - self.cursor_ref, - cursors, - "Column" - ); - let record = cursor.record(); - if let Some(record) = record.as_ref() { - state.registers[*dest] = if cursor.get_null_flag() { - OwnedValue::Null + let value = { + let cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + state, + "Column" + ); + let record = cursor.record(); + if let Some(record) = record.as_ref() { + if cursor.get_null_flag() { + OwnedValue::Null + } else { + record.get_value(*column).clone() + } } else { - record.get_value(*column).clone() - }; - } else { - state.registers[*dest] = OwnedValue::Null; - } + OwnedValue::Null + } + }; + state.registers[*dest] = value; } CursorType::Sorter => { let cursor = state.get_sorter(*cursor_id); @@ -1117,17 +1110,15 @@ impl Program { return Ok(StepResult::Row); } Insn::NextAsync { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "NextAsync"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NextAsync"); cursor.set_null_flag(false); return_if_io!(cursor.next()); state.pc += 1; } Insn::PrevAsync { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "PrevAsync"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "PrevAsync"); cursor.set_null_flag(false); return_if_io!(cursor.prev()); state.pc += 1; @@ -1136,10 +1127,9 @@ impl Program { cursor_id, pc_if_next, } => { - let mut cursors = state.cursors.borrow_mut(); assert!(pc_if_next.is_offset()); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "PrevAwait"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "PrevAwait"); cursor.wait_for_completion()?; if !cursor.is_empty() { state.pc = pc_if_next.to_offset_int(); @@ -1152,9 +1142,8 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, cursors, "NextAwait"); + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NextAwait"); cursor.wait_for_completion()?; if !cursor.is_empty() { state.pc = pc_if_next.to_offset_int(); @@ -1288,11 +1277,10 @@ impl Program { state.pc += 1; } Insn::RowId { cursor_id, dest } => { - let mut cursors = state.cursors.borrow_mut(); if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() { - let index_cursor = get_cursor_as_index_mut(&mut cursors, index_cursor_id); + let index_cursor = state.get_btree_index_cursor(index_cursor_id); let rowid = index_cursor.rowid()?; - let table_cursor = get_cursor_as_table_mut(&mut cursors, table_cursor_id); + let table_cursor = state.get_btree_table_cursor(table_cursor_id); match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? { CursorResult::Ok(_) => {} CursorResult::IO => { @@ -1301,7 +1289,7 @@ impl Program { } } } - + let mut cursors = state.cursors.borrow_mut(); if let Some(Cursor::Table(btree_cursor)) = cursors.get_mut(*cursor_id).unwrap() { if let Some(ref rowid) = btree_cursor.rowid()? { @@ -1335,8 +1323,7 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor_id); + let cursor = state.get_btree_table_cursor(*cursor_id); let rowid = match &state.registers[*src_reg] { OwnedValue::Integer(rowid) => *rowid as u64, OwnedValue::Null => { @@ -1371,9 +1358,8 @@ impl Program { is_index, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); if *is_index { - let cursor = get_cursor_as_index_mut(&mut cursors, *cursor_id); + let cursor = state.get_btree_index_cursor(*cursor_id); let record_from_regs: Record = make_owned_record(&state.registers, start_reg, num_regs); let found = return_if_io!( @@ -1385,7 +1371,7 @@ impl Program { state.pc += 1; } } else { - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor_id); + let cursor = state.get_btree_table_cursor(*cursor_id); let rowid = match &state.registers[*start_reg] { OwnedValue::Null => { // All integer values are greater than null so we just rewind the cursor @@ -1417,9 +1403,8 @@ impl Program { is_index, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); if *is_index { - let cursor = get_cursor_as_index_mut(&mut cursors, *cursor_id); + let cursor = state.get_btree_index_cursor(*cursor_id); let record_from_regs: Record = make_owned_record(&state.registers, start_reg, num_regs); let found = return_if_io!( @@ -1431,7 +1416,7 @@ impl Program { state.pc += 1; } } else { - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor_id); + let cursor = state.get_btree_table_cursor(*cursor_id); let rowid = match &state.registers[*start_reg] { OwnedValue::Null => { // All integer values are greater than null so we just rewind the cursor @@ -1462,22 +1447,24 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_index_mut(&mut cursors, *cursor_id); - let record_from_regs: Record = - make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { - // Compare against the same number of values - if idx_record.get_values()[..record_from_regs.len()] - >= record_from_regs.get_values()[..] - { - state.pc = target_pc.to_offset_int(); + let pc = { + let cursor = state.get_btree_index_cursor(*cursor_id); + let record_from_regs: Record = + make_owned_record(&state.registers, start_reg, num_regs); + if let Some(ref idx_record) = *cursor.record() { + // Compare against the same number of values + if idx_record.get_values()[..record_from_regs.len()] + >= record_from_regs.get_values()[..] + { + target_pc.to_offset_int() + } else { + state.pc + 1 + } } else { - state.pc += 1; + target_pc.to_offset_int() } - } else { - state.pc = target_pc.to_offset_int(); }; + state.pc = pc; } Insn::IdxLE { cursor_id, @@ -1486,22 +1473,24 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_index_mut(&mut cursors, *cursor_id); - let record_from_regs: Record = - make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { - // Compare against the same number of values - if idx_record.get_values()[..record_from_regs.len()] - <= record_from_regs.get_values()[..] - { - state.pc = target_pc.to_offset_int(); + let pc = { + let cursor = state.get_btree_index_cursor(*cursor_id); + let record_from_regs: Record = + make_owned_record(&state.registers, start_reg, num_regs); + if let Some(ref idx_record) = *cursor.record() { + // Compare against the same number of values + if idx_record.get_values()[..record_from_regs.len()] + <= record_from_regs.get_values()[..] + { + target_pc.to_offset_int() + } else { + state.pc + 1 + } } else { - state.pc += 1; + target_pc.to_offset_int() } - } else { - state.pc = target_pc.to_offset_int(); }; + state.pc = pc; } Insn::IdxGT { cursor_id, @@ -1510,22 +1499,24 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_index_mut(&mut cursors, *cursor_id); - let record_from_regs: Record = - make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { - // Compare against the same number of values - if idx_record.get_values()[..record_from_regs.len()] - > record_from_regs.get_values()[..] - { - state.pc = target_pc.to_offset_int(); + let pc = { + let cursor = state.get_btree_index_cursor(*cursor_id); + let record_from_regs: Record = + make_owned_record(&state.registers, start_reg, num_regs); + if let Some(ref idx_record) = *cursor.record() { + // Compare against the same number of values + if idx_record.get_values()[..record_from_regs.len()] + > record_from_regs.get_values()[..] + { + target_pc.to_offset_int() + } else { + state.pc + 1 + } } else { - state.pc += 1; + target_pc.to_offset_int() } - } else { - state.pc = target_pc.to_offset_int(); }; + state.pc = pc; } Insn::IdxLT { cursor_id, @@ -1534,22 +1525,24 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_index_mut(&mut cursors, *cursor_id); - let record_from_regs: Record = - make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { - // Compare against the same number of values - if idx_record.get_values()[..record_from_regs.len()] - < record_from_regs.get_values()[..] - { - state.pc = target_pc.to_offset_int(); + let pc = { + let cursor = state.get_btree_index_cursor(*cursor_id); + let record_from_regs: Record = + make_owned_record(&state.registers, start_reg, num_regs); + if let Some(ref idx_record) = *cursor.record() { + // Compare against the same number of values + if idx_record.get_values()[..record_from_regs.len()] + < record_from_regs.get_values()[..] + { + target_pc.to_offset_int() + } else { + state.pc + 1 + } } else { - state.pc += 1; + target_pc.to_offset_int() } - } else { - state.pc = target_pc.to_offset_int(); }; + state.pc = pc; } Insn::DecrJumpZero { reg, target_pc } => { assert!(target_pc.is_offset()); @@ -2676,8 +2669,7 @@ impl Program { record_reg, flag: _, } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor); + let cursor = &mut state.get_btree_table_cursor(*cursor); let record = match &state.registers[*record_reg] { OwnedValue::Record(r) => r, _ => unreachable!("Not a record! Cannot insert a non record value."), @@ -2690,8 +2682,7 @@ impl Program { state.pc += 1; } Insn::InsertAwait { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor_id); + let cursor = &mut state.get_btree_table_cursor(*cursor_id); cursor.wait_for_completion()?; // Only update last_insert_rowid for regular table inserts, not schema modifications if cursor.root_page() != 1 { @@ -2706,14 +2697,12 @@ impl Program { state.pc += 1; } Insn::DeleteAsync { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor_id); + let cursor = &mut state.get_btree_table_cursor(*cursor_id); return_if_io!(cursor.delete()); state.pc += 1; } Insn::DeleteAwait { cursor_id } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor_id); + let cursor = &mut state.get_btree_table_cursor(*cursor_id); cursor.wait_for_completion()?; let prev_changes = self.n_change.get(); self.n_change.set(prev_changes + 1); @@ -2722,8 +2711,7 @@ impl Program { Insn::NewRowid { cursor, rowid_reg, .. } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = get_cursor_as_table_mut(&mut cursors, *cursor); + let cursor = &mut state.get_btree_table_cursor(*cursor); // TODO: make io handle rng let rowid = return_if_io!(get_new_rowid(cursor, thread_rng())); state.registers[*rowid_reg] = OwnedValue::Integer(rowid); @@ -2768,9 +2756,8 @@ impl Program { rowid_reg, target_pc, } => { - let mut cursors = state.cursors.borrow_mut(); let cursor = - must_be_btree_cursor!(*cursor, self.cursor_ref, cursors, "NotExists"); + must_be_btree_cursor!(*cursor, self.cursor_ref, state, "NotExists"); let exists = return_if_io!(cursor.exists(&state.registers[*rowid_reg])); if exists { state.pc += 1; From 085f93ce79e130f227c7966bdc318938ada85af9 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 12:23:35 +0200 Subject: [PATCH 05/10] core/vdbe: Add ProgramState::get_cursor() helper --- core/vdbe/mod.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 206b66dbb..6594cf430 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -342,6 +342,16 @@ impl ProgramState { .as_virtual_mut(); unsafe { std::mem::transmute(cursor) } } + + pub fn get_cursor<'a>(&'a self, cursor_id: CursorID) -> std::cell::RefMut<'a, Cursor> { + let cursors = self.cursors.borrow_mut(); + std::cell::RefMut::map(cursors, |c| { + c.get_mut(cursor_id) + .expect("cursor id out of bounds") + .as_mut() + .expect("cursor not allocated") + }) + } } macro_rules! must_be_btree_cursor { From 45539a4fe520af4829ef692f6e983a25aadccc92 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 13:43:49 +0200 Subject: [PATCH 06/10] core/vdbe: Replace get_vtab_cursor() calls with get_cursor() --- core/vdbe/mod.rs | 41 ++++++++++++++++++++--------------------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 6594cf430..00bbd22ec 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -332,17 +332,6 @@ impl ProgramState { unsafe { std::mem::transmute(cursor) } } - pub fn get_vtab_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut VTabOpaqueCursor { - let mut cursors = self.cursors.borrow_mut(); - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_virtual_mut(); - unsafe { std::mem::transmute(cursor) } - } - pub fn get_cursor<'a>(&'a self, cursor_id: CursorID) -> std::cell::RefMut<'a, Cursor> { let cursors = self.cursors.borrow_mut(); std::cell::RefMut::map(cursors, |c| { @@ -892,12 +881,15 @@ impl Program { let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VFilter on non-virtual table cursor"); }; - let cursor = state.get_vtab_cursor(*cursor_id); - let mut args = Vec::new(); - for i in 0..*arg_count { - args.push(state.registers[args_reg + i].clone()); - } - let has_rows = virtual_table.filter(cursor, *arg_count, args)?; + let has_rows = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_virtual_mut(); + let mut args = Vec::new(); + for i in 0..*arg_count { + args.push(state.registers[args_reg + i].clone()); + } + virtual_table.filter(cursor, *arg_count, args)? + }; if !has_rows { state.pc = pc_if_empty.to_offset_int(); } else { @@ -913,8 +905,12 @@ impl Program { let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VColumn on non-virtual table cursor"); }; - let cursor = state.get_vtab_cursor(*cursor_id); - state.registers[*dest] = virtual_table.column(cursor, *column)?; + let value = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_virtual_mut(); + virtual_table.column(cursor, *column)? + }; + state.registers[*dest] = value; state.pc += 1; } Insn::VUpdate { @@ -978,8 +974,11 @@ impl Program { let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VNextAsync on non-virtual table cursor"); }; - let cursor = state.get_vtab_cursor(*cursor_id); - let has_more = virtual_table.next(cursor)?; + let has_more = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_virtual_mut(); + virtual_table.next(cursor)? + }; if has_more { state.pc = pc_if_next.to_offset_int(); } else { From c12f2aeca40cc62f28ae7a3e90573b96df7e0904 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 13:44:24 +0200 Subject: [PATCH 07/10] core/vdbe: Replace get_sorter() calls with get_cursor() --- core/vdbe/mod.rs | 67 ++++++++++++++++++++++++++++-------------------- 1 file changed, 39 insertions(+), 28 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 00bbd22ec..df6b35ef7 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -321,17 +321,6 @@ impl ProgramState { unsafe { std::mem::transmute(cursor) } } - pub fn get_sorter<'a>(&'a self, cursor_id: CursorID) -> &'a mut Sorter { - let mut cursors = self.cursors.borrow_mut(); - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_sorter_mut(); - unsafe { std::mem::transmute(cursor) } - } - pub fn get_cursor<'a>(&'a self, cursor_id: CursorID) -> std::cell::RefMut<'a, Cursor> { let cursors = self.cursors.borrow_mut(); std::cell::RefMut::map(cursors, |c| { @@ -1079,8 +1068,12 @@ impl Program { state.registers[*dest] = value; } CursorType::Sorter => { - let cursor = state.get_sorter(*cursor_id); - if let Some(record) = cursor.record() { + let record = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_sorter_mut(); + cursor.record().map(|r| r.clone()) + }; + if let Some(record) = record { state.registers[*dest] = record.get_value(*column).clone(); } else { state.registers[*dest] = OwnedValue::Null; @@ -1930,9 +1923,13 @@ impl Program { dest_reg, pseudo_cursor, } => { - let sorter_cursor = state.get_sorter(*cursor_id); - let record = match sorter_cursor.record() { - Some(record) => record.clone(), + let record = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_sorter_mut(); + cursor.record().map(|r| r.clone()) + }; + let record = match record { + Some(record) => record, None => { state.pc += 1; continue; @@ -1947,23 +1944,33 @@ impl Program { cursor_id, record_reg, } => { - let cursor = state.get_sorter(*cursor_id); - let record = match &state.registers[*record_reg] { - OwnedValue::Record(record) => record, - _ => unreachable!("SorterInsert on non-record register"), - }; - cursor.insert(record); + { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_sorter_mut(); + let record = match &state.registers[*record_reg] { + OwnedValue::Record(record) => record, + _ => unreachable!("SorterInsert on non-record register"), + }; + cursor.insert(record); + } state.pc += 1; } Insn::SorterSort { cursor_id, pc_if_empty, } => { - let cursor = state.get_sorter(*cursor_id); - if cursor.is_empty() { + let is_empty = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_sorter_mut(); + let is_empty = cursor.is_empty(); + if !is_empty { + cursor.sort(); + } + is_empty + }; + if is_empty { state.pc = pc_if_empty.to_offset_int(); } else { - cursor.sort(); state.pc += 1; } } @@ -1972,9 +1979,13 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let cursor = state.get_sorter(*cursor_id); - cursor.next(); - if cursor.has_more() { + let has_more = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_sorter_mut(); + cursor.next(); + cursor.has_more() + }; + if has_more { state.pc = pc_if_next.to_offset_int(); } else { state.pc += 1; From 1c0d9c3b466efaf1b69912d3e323ac38b3499059 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 14:15:52 +0200 Subject: [PATCH 08/10] core/vdbe: Replace get_pseudo_cursor() calls with get_cursor() --- core/vdbe/mod.rs | 33 ++++++++++++++------------------- 1 file changed, 14 insertions(+), 19 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index df6b35ef7..2c516d285 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -310,17 +310,6 @@ impl ProgramState { unsafe { std::mem::transmute(cursor) } } - pub fn get_pseudo_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut PseudoCursor { - let mut cursors = self.cursors.borrow_mut(); - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_pseudo_mut(); - unsafe { std::mem::transmute(cursor) } - } - pub fn get_cursor<'a>(&'a self, cursor_id: CursorID) -> std::cell::RefMut<'a, Cursor> { let cursors = self.cursors.borrow_mut(); std::cell::RefMut::map(cursors, |c| { @@ -1080,12 +1069,16 @@ impl Program { } } CursorType::Pseudo(_) => { - let cursor = state.get_pseudo_cursor(*cursor_id); - if let Some(record) = cursor.record() { - state.registers[*dest] = record.get_value(*column).clone(); - } else { - state.registers[*dest] = OwnedValue::Null; - } + let value = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_pseudo_mut(); + if let Some(record) = cursor.record() { + record.get_value(*column).clone() + } else { + OwnedValue::Null + } + }; + state.registers[*dest] = value; } CursorType::VirtualTable(_) => { panic!( @@ -1936,8 +1929,10 @@ impl Program { } }; state.registers[*dest_reg] = OwnedValue::Record(record.clone()); - let pseudo_cursor = state.get_pseudo_cursor(*pseudo_cursor); - pseudo_cursor.insert(record); + { + let mut pseudo_cursor = state.get_cursor(*pseudo_cursor); + pseudo_cursor.as_pseudo_mut().insert(record); + } state.pc += 1; } Insn::SorterInsert { From cdcaebb87845e8227cefc2c93a43aa30ad0d667f Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 14:34:12 +0200 Subject: [PATCH 09/10] core/vdbe: Unify B-Tree cursors --- core/types.rs | 24 ++++++------------------ core/vdbe/mod.rs | 14 +++++++------- 2 files changed, 13 insertions(+), 25 deletions(-) diff --git a/core/types.rs b/core/types.rs index 93162a01e..2b3bd2e38 100644 --- a/core/types.rs +++ b/core/types.rs @@ -619,20 +619,15 @@ impl Record { } pub enum Cursor { - Table(BTreeCursor), - Index(BTreeCursor), + BTree(BTreeCursor), Pseudo(PseudoCursor), Sorter(Sorter), Virtual(VTabOpaqueCursor), } impl Cursor { - pub fn new_table(cursor: BTreeCursor) -> Self { - Self::Table(cursor) - } - - pub fn new_index(cursor: BTreeCursor) -> Self { - Self::Index(cursor) + pub fn new_btree(cursor: BTreeCursor) -> Self { + Self::BTree(cursor) } pub fn new_pseudo(cursor: PseudoCursor) -> Self { @@ -643,17 +638,10 @@ impl Cursor { Self::Sorter(cursor) } - pub fn as_table_mut(&mut self) -> &mut BTreeCursor { + pub fn as_btree_mut(&mut self) -> &mut BTreeCursor { match self { - Self::Table(cursor) => cursor, - _ => panic!("Cursor is not a table"), - } - } - - pub fn as_index_mut(&mut self) -> &mut BTreeCursor { - match self { - Self::Index(cursor) => cursor, - _ => panic!("Cursor is not an index"), + Self::BTree(cursor) => cursor, + _ => panic!("Cursor is not a btree"), } } diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 2c516d285..d852b4e18 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -295,7 +295,7 @@ impl ProgramState { .expect("cursor id out of bounds") .as_mut() .expect("cursor not allocated") - .as_table_mut(); + .as_btree_mut(); unsafe { std::mem::transmute(cursor) } } @@ -306,7 +306,7 @@ impl ProgramState { .expect("cursor id out of bounds") .as_mut() .expect("cursor not allocated") - .as_index_mut(); + .as_btree_mut(); unsafe { std::mem::transmute(cursor) } } @@ -770,13 +770,13 @@ impl Program { cursors .get_mut(*cursor_id) .unwrap() - .replace(Cursor::new_table(cursor)); + .replace(Cursor::new_btree(cursor)); } CursorType::BTreeIndex(_) => { cursors .get_mut(*cursor_id) .unwrap() - .replace(Cursor::new_index(cursor)); + .replace(Cursor::new_btree(cursor)); } CursorType::Pseudo(_) => { panic!("OpenReadAsync on pseudo cursor"); @@ -1285,7 +1285,7 @@ impl Program { } } let mut cursors = state.cursors.borrow_mut(); - if let Some(Cursor::Table(btree_cursor)) = cursors.get_mut(*cursor_id).unwrap() + if let Some(Cursor::BTree(btree_cursor)) = cursors.get_mut(*cursor_id).unwrap() { if let Some(ref rowid) = btree_cursor.rowid()? { state.registers[*dest] = OwnedValue::Integer(*rowid as i64); @@ -2826,12 +2826,12 @@ impl Program { cursors .get_mut(*cursor_id) .unwrap() - .replace(Cursor::new_index(cursor)); + .replace(Cursor::new_btree(cursor)); } else { cursors .get_mut(*cursor_id) .unwrap() - .replace(Cursor::new_table(cursor)); + .replace(Cursor::new_btree(cursor)); } state.pc += 1; } From f3ee86d784008df4cd847c1785c61e3e61a2ea97 Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Tue, 4 Mar 2025 14:35:57 +0200 Subject: [PATCH 10/10] core/vdbe: Replace get_btree_{table,index}_cursor() calls with get_cursor() --- core/vdbe/mod.rs | 490 ++++++++++++++++++++++++++++------------------- 1 file changed, 298 insertions(+), 192 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index d852b4e18..f2c34c6e9 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -288,28 +288,6 @@ impl ProgramState { self.parameters.clear(); } - pub fn get_btree_table_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut BTreeCursor { - let mut cursors = self.cursors.borrow_mut(); - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_btree_mut(); - unsafe { std::mem::transmute(cursor) } - } - - pub fn get_btree_index_cursor<'a>(&'a self, cursor_id: CursorID) -> &'a mut BTreeCursor { - let mut cursors = self.cursors.borrow_mut(); - let cursor = cursors - .get_mut(cursor_id) - .expect("cursor id out of bounds") - .as_mut() - .expect("cursor not allocated") - .as_btree_mut(); - unsafe { std::mem::transmute(cursor) } - } - pub fn get_cursor<'a>(&'a self, cursor_id: CursorID) -> std::cell::RefMut<'a, Cursor> { let cursors = self.cursors.borrow_mut(); std::cell::RefMut::map(cursors, |c| { @@ -325,8 +303,8 @@ macro_rules! must_be_btree_cursor { ($cursor_id:expr, $cursor_ref:expr, $state:expr, $insn_name:expr) => {{ let (_, cursor_type) = $cursor_ref.get($cursor_id).unwrap(); let cursor = match cursor_type { - CursorType::BTreeTable(_) => $state.get_btree_table_cursor($cursor_id), - CursorType::BTreeIndex(_) => $state.get_btree_index_cursor($cursor_id), + CursorType::BTreeTable(_) => $state.get_cursor($cursor_id), + CursorType::BTreeIndex(_) => $state.get_cursor($cursor_id), CursorType::Pseudo(_) => panic!("{} on pseudo cursor", $insn_name), CursorType::Sorter => panic!("{} on sorter cursor", $insn_name), CursorType::VirtualTable(_) => panic!("{} on virtual table cursor", $insn_name), @@ -456,9 +434,12 @@ impl Program { state.pc += 1; } Insn::NullRow { cursor_id } => { - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NullRow"); - cursor.set_null_flag(true); + { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NullRow"); + let cursor = cursor.as_btree_mut(); + cursor.set_null_flag(true); + } state.pc += 1; } Insn::Compare { @@ -968,24 +949,36 @@ impl Program { content_reg: _, num_fields: _, } => { - let mut cursors = state.cursors.borrow_mut(); - let cursor = PseudoCursor::new(); - cursors - .get_mut(*cursor_id) - .unwrap() - .replace(Cursor::new_pseudo(cursor)); + { + let mut cursors = state.cursors.borrow_mut(); + let cursor = PseudoCursor::new(); + cursors + .get_mut(*cursor_id) + .unwrap() + .replace(Cursor::new_pseudo(cursor)); + } state.pc += 1; } Insn::RewindAsync { cursor_id } => { - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "RewindAsync"); - return_if_io!(cursor.rewind()); + { + let mut cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + state, + "RewindAsync" + ); + let cursor = cursor.as_btree_mut(); + return_if_io!(cursor.rewind()); + } state.pc += 1; } Insn::LastAsync { cursor_id } => { - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "LastAsync"); - return_if_io!(cursor.last()); + { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "LastAsync"); + let cursor = cursor.as_btree_mut(); + return_if_io!(cursor.last()); + } state.pc += 1; } Insn::LastAwait { @@ -993,10 +986,14 @@ impl Program { pc_if_empty, } => { assert!(pc_if_empty.is_offset()); - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "LastAwait"); - cursor.wait_for_completion()?; - if cursor.is_empty() { + let is_empty = { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "LastAwait"); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + cursor.is_empty() + }; + if is_empty { state.pc = pc_if_empty.to_offset_int(); } else { state.pc += 1; @@ -1007,10 +1004,18 @@ impl Program { pc_if_empty, } => { assert!(pc_if_empty.is_offset()); - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "RewindAwait"); - cursor.wait_for_completion()?; - if cursor.is_empty() { + let is_empty = { + let mut cursor = must_be_btree_cursor!( + *cursor_id, + self.cursor_ref, + state, + "RewindAwait" + ); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + cursor.is_empty() + }; + if is_empty { state.pc = pc_if_empty.to_offset_int(); } else { state.pc += 1; @@ -1022,27 +1027,38 @@ impl Program { dest, } => { if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() { - let index_cursor = state.get_btree_index_cursor(index_cursor_id); - let rowid = index_cursor.rowid()?; - let table_cursor = state.get_btree_table_cursor(table_cursor_id); - match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? { - CursorResult::Ok(_) => {} - CursorResult::IO => { - state.deferred_seek = Some((index_cursor_id, table_cursor_id)); - return Ok(StepResult::IO); + let deferred_seek = { + let rowid = { + let mut index_cursor = state.get_cursor(index_cursor_id); + let index_cursor = index_cursor.as_btree_mut(); + let rowid = index_cursor.rowid()?; + rowid + }; + let mut table_cursor = state.get_cursor(table_cursor_id); + let table_cursor = table_cursor.as_btree_mut(); + match table_cursor + .seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? + { + CursorResult::Ok(_) => None, + CursorResult::IO => Some((index_cursor_id, table_cursor_id)), } + }; + if let Some(deferred_seek) = deferred_seek { + state.deferred_seek = Some(deferred_seek); + return Ok(StepResult::IO); } } let (_, cursor_type) = self.cursor_ref.get(*cursor_id).unwrap(); match cursor_type { CursorType::BTreeTable(_) | CursorType::BTreeIndex(_) => { let value = { - let cursor = must_be_btree_cursor!( + let mut cursor = must_be_btree_cursor!( *cursor_id, self.cursor_ref, state, "Column" ); + let cursor = cursor.as_btree_mut(); let record = cursor.record(); if let Some(record) = record.as_ref() { if cursor.get_null_flag() { @@ -1105,17 +1121,23 @@ impl Program { return Ok(StepResult::Row); } Insn::NextAsync { cursor_id } => { - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NextAsync"); - cursor.set_null_flag(false); - return_if_io!(cursor.next()); + { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NextAsync"); + let cursor = cursor.as_btree_mut(); + cursor.set_null_flag(false); + return_if_io!(cursor.next()); + } state.pc += 1; } Insn::PrevAsync { cursor_id } => { - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "PrevAsync"); - cursor.set_null_flag(false); - return_if_io!(cursor.prev()); + { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "PrevAsync"); + let cursor = cursor.as_btree_mut(); + cursor.set_null_flag(false); + return_if_io!(cursor.prev()); + } state.pc += 1; } Insn::PrevAwait { @@ -1123,10 +1145,14 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "PrevAwait"); - cursor.wait_for_completion()?; - if !cursor.is_empty() { + let is_empty = { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "PrevAwait"); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + cursor.is_empty() + }; + if !is_empty { state.pc = pc_if_next.to_offset_int(); } else { state.pc += 1; @@ -1137,10 +1163,14 @@ impl Program { pc_if_next, } => { assert!(pc_if_next.is_offset()); - let cursor = - must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NextAwait"); - cursor.wait_for_completion()?; - if !cursor.is_empty() { + let is_empty = { + let mut cursor = + must_be_btree_cursor!(*cursor_id, self.cursor_ref, state, "NextAwait"); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + cursor.is_empty() + }; + if !is_empty { state.pc = pc_if_next.to_offset_int(); } else { state.pc += 1; @@ -1273,15 +1303,26 @@ impl Program { } Insn::RowId { cursor_id, dest } => { if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() { - let index_cursor = state.get_btree_index_cursor(index_cursor_id); - let rowid = index_cursor.rowid()?; - let table_cursor = state.get_btree_table_cursor(table_cursor_id); - match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? { - CursorResult::Ok(_) => {} - CursorResult::IO => { - state.deferred_seek = Some((index_cursor_id, table_cursor_id)); - return Ok(StepResult::IO); - } + let deferred_seek = { + let rowid = { + let mut index_cursor = state.get_cursor(index_cursor_id); + let index_cursor = index_cursor.as_btree_mut(); + let rowid = index_cursor.rowid()?; + rowid + }; + let mut table_cursor = state.get_cursor(table_cursor_id); + let table_cursor = table_cursor.as_btree_mut(); + let deferred_seek = match table_cursor + .seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? + { + CursorResult::Ok(_) => None, + CursorResult::IO => Some((index_cursor_id, table_cursor_id)), + }; + deferred_seek + }; + if let Some(deferred_seek) = deferred_seek { + state.deferred_seek = Some(deferred_seek); + return Ok(StepResult::IO); } } let mut cursors = state.cursors.borrow_mut(); @@ -1318,25 +1359,33 @@ impl Program { target_pc, } => { assert!(target_pc.is_offset()); - let cursor = state.get_btree_table_cursor(*cursor_id); - let rowid = match &state.registers[*src_reg] { - OwnedValue::Integer(rowid) => *rowid as u64, - OwnedValue::Null => { - state.pc = target_pc.to_offset_int(); - continue; - } - other => { - return Err(LimboError::InternalError( - format!("SeekRowid: the value in the register is not an integer or NULL: {}", other) - )); + let pc = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + let rowid = match &state.registers[*src_reg] { + OwnedValue::Integer(rowid) => Some(*rowid as u64), + OwnedValue::Null => None, + other => { + return Err(LimboError::InternalError( + format!("SeekRowid: the value in the register is not an integer or NULL: {}", other) + )); + } + }; + match rowid { + Some(rowid) => { + let found = return_if_io!( + cursor.seek(SeekKey::TableRowId(rowid), SeekOp::EQ) + ); + if !found { + target_pc.to_offset_int() + } else { + state.pc + 1 + } + } + None => target_pc.to_offset_int(), } }; - let found = return_if_io!(cursor.seek(SeekKey::TableRowId(rowid), SeekOp::EQ)); - if !found { - state.pc = target_pc.to_offset_int(); - } else { - state.pc += 1; - } + state.pc = pc; } Insn::DeferredSeek { index_cursor_id, @@ -1354,40 +1403,54 @@ impl Program { } => { assert!(target_pc.is_offset()); if *is_index { - let cursor = state.get_btree_index_cursor(*cursor_id); - let record_from_regs: Record = - make_owned_record(&state.registers, start_reg, num_regs); - let found = return_if_io!( - cursor.seek(SeekKey::IndexKey(&record_from_regs), SeekOp::GE) - ); + let found = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + let record_from_regs: Record = + make_owned_record(&state.registers, start_reg, num_regs); + let found = return_if_io!( + cursor.seek(SeekKey::IndexKey(&record_from_regs), SeekOp::GE) + ); + found + }; if !found { state.pc = target_pc.to_offset_int(); } else { state.pc += 1; } } else { - let cursor = state.get_btree_table_cursor(*cursor_id); - let rowid = match &state.registers[*start_reg] { - OwnedValue::Null => { - // All integer values are greater than null so we just rewind the cursor - return_if_io!(cursor.rewind()); - state.pc += 1; - continue; - } - OwnedValue::Integer(rowid) => *rowid as u64, - _ => { - return Err(LimboError::InternalError( - "SeekGE: the value in the register is not an integer".into(), - )); + let pc = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + let rowid = match &state.registers[*start_reg] { + OwnedValue::Null => { + // All integer values are greater than null so we just rewind the cursor + return_if_io!(cursor.rewind()); + None + } + OwnedValue::Integer(rowid) => Some(*rowid as u64), + _ => { + return Err(LimboError::InternalError( + "SeekGE: the value in the register is not an integer" + .into(), + )); + } + }; + match rowid { + Some(rowid) => { + let found = return_if_io!( + cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GE) + ); + if !found { + target_pc.to_offset_int() + } else { + state.pc + 1 + } + } + None => state.pc + 1, } }; - let found = - return_if_io!(cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GE)); - if !found { - state.pc = target_pc.to_offset_int(); - } else { - state.pc += 1; - } + state.pc = pc; } } Insn::SeekGT { @@ -1399,40 +1462,55 @@ impl Program { } => { assert!(target_pc.is_offset()); if *is_index { - let cursor = state.get_btree_index_cursor(*cursor_id); - let record_from_regs: Record = - make_owned_record(&state.registers, start_reg, num_regs); - let found = return_if_io!( - cursor.seek(SeekKey::IndexKey(&record_from_regs), SeekOp::GT) - ); + let found = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + let record_from_regs: Record = + make_owned_record(&state.registers, start_reg, num_regs); + let found = return_if_io!( + cursor.seek(SeekKey::IndexKey(&record_from_regs), SeekOp::GT) + ); + found + }; if !found { state.pc = target_pc.to_offset_int(); } else { state.pc += 1; } } else { - let cursor = state.get_btree_table_cursor(*cursor_id); - let rowid = match &state.registers[*start_reg] { - OwnedValue::Null => { - // All integer values are greater than null so we just rewind the cursor - return_if_io!(cursor.rewind()); - state.pc += 1; - continue; - } - OwnedValue::Integer(rowid) => *rowid as u64, - _ => { - return Err(LimboError::InternalError( - "SeekGT: the value in the register is not an integer".into(), - )); - } + let pc = { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + let rowid = match &state.registers[*start_reg] { + OwnedValue::Null => { + // All integer values are greater than null so we just rewind the cursor + return_if_io!(cursor.rewind()); + None + } + OwnedValue::Integer(rowid) => Some(*rowid as u64), + _ => { + return Err(LimboError::InternalError( + "SeekGT: the value in the register is not an integer" + .into(), + )); + } + }; + let found = match rowid { + Some(rowid) => { + let found = return_if_io!( + cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GT) + ); + if !found { + target_pc.to_offset_int() + } else { + state.pc + 1 + } + } + None => state.pc + 1, + }; + found }; - let found = - return_if_io!(cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GT)); - if !found { - state.pc = target_pc.to_offset_int(); - } else { - state.pc += 1; - } + state.pc = pc; } } Insn::IdxGE { @@ -1443,10 +1521,11 @@ impl Program { } => { assert!(target_pc.is_offset()); let pc = { - let cursor = state.get_btree_index_cursor(*cursor_id); + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); let record_from_regs: Record = make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { + let pc = if let Some(ref idx_record) = *cursor.record() { // Compare against the same number of values if idx_record.get_values()[..record_from_regs.len()] >= record_from_regs.get_values()[..] @@ -1457,7 +1536,8 @@ impl Program { } } else { target_pc.to_offset_int() - } + }; + pc }; state.pc = pc; } @@ -1469,10 +1549,11 @@ impl Program { } => { assert!(target_pc.is_offset()); let pc = { - let cursor = state.get_btree_index_cursor(*cursor_id); + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); let record_from_regs: Record = make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { + let pc = if let Some(ref idx_record) = *cursor.record() { // Compare against the same number of values if idx_record.get_values()[..record_from_regs.len()] <= record_from_regs.get_values()[..] @@ -1483,7 +1564,8 @@ impl Program { } } else { target_pc.to_offset_int() - } + }; + pc }; state.pc = pc; } @@ -1495,10 +1577,11 @@ impl Program { } => { assert!(target_pc.is_offset()); let pc = { - let cursor = state.get_btree_index_cursor(*cursor_id); + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); let record_from_regs: Record = make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { + let pc = if let Some(ref idx_record) = *cursor.record() { // Compare against the same number of values if idx_record.get_values()[..record_from_regs.len()] > record_from_regs.get_values()[..] @@ -1509,7 +1592,8 @@ impl Program { } } else { target_pc.to_offset_int() - } + }; + pc }; state.pc = pc; } @@ -1521,10 +1605,11 @@ impl Program { } => { assert!(target_pc.is_offset()); let pc = { - let cursor = state.get_btree_index_cursor(*cursor_id); + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); let record_from_regs: Record = make_owned_record(&state.registers, start_reg, num_regs); - if let Some(ref idx_record) = *cursor.record() { + let pc = if let Some(ref idx_record) = *cursor.record() { // Compare against the same number of values if idx_record.get_values()[..record_from_regs.len()] < record_from_regs.get_values()[..] @@ -1535,7 +1620,8 @@ impl Program { } } else { target_pc.to_offset_int() - } + }; + pc }; state.pc = pc; } @@ -2684,41 +2770,53 @@ impl Program { record_reg, flag: _, } => { - let cursor = &mut state.get_btree_table_cursor(*cursor); - let record = match &state.registers[*record_reg] { - OwnedValue::Record(r) => r, - _ => unreachable!("Not a record! Cannot insert a non record value."), - }; - let key = &state.registers[*key_reg]; - // NOTE(pere): Sending moved_before == true is okay because we moved before but - // if we were to set to false after starting a balance procedure, it might - // leave undefined state. - return_if_io!(cursor.insert(key, record, true)); + { + let mut cursor = state.get_cursor(*cursor); + let cursor = cursor.as_btree_mut(); + let record = match &state.registers[*record_reg] { + OwnedValue::Record(r) => r, + _ => unreachable!("Not a record! Cannot insert a non record value."), + }; + let key = &state.registers[*key_reg]; + // NOTE(pere): Sending moved_before == true is okay because we moved before but + // if we were to set to false after starting a balance procedure, it might + // leave undefined state. + return_if_io!(cursor.insert(key, record, true)); + } state.pc += 1; } Insn::InsertAwait { cursor_id } => { - let cursor = &mut state.get_btree_table_cursor(*cursor_id); - cursor.wait_for_completion()?; - // Only update last_insert_rowid for regular table inserts, not schema modifications - if cursor.root_page() != 1 { - if let Some(rowid) = cursor.rowid()? { - if let Some(conn) = self.connection.upgrade() { - conn.update_last_rowid(rowid); + { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + // Only update last_insert_rowid for regular table inserts, not schema modifications + if cursor.root_page() != 1 { + if let Some(rowid) = cursor.rowid()? { + if let Some(conn) = self.connection.upgrade() { + conn.update_last_rowid(rowid); + } + let prev_changes = self.n_change.get(); + self.n_change.set(prev_changes + 1); } - let prev_changes = self.n_change.get(); - self.n_change.set(prev_changes + 1); } } state.pc += 1; } Insn::DeleteAsync { cursor_id } => { - let cursor = &mut state.get_btree_table_cursor(*cursor_id); - return_if_io!(cursor.delete()); + { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + return_if_io!(cursor.delete()); + } state.pc += 1; } Insn::DeleteAwait { cursor_id } => { - let cursor = &mut state.get_btree_table_cursor(*cursor_id); - cursor.wait_for_completion()?; + { + let mut cursor = state.get_cursor(*cursor_id); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + } let prev_changes = self.n_change.get(); self.n_change.set(prev_changes + 1); state.pc += 1; @@ -2726,9 +2824,13 @@ impl Program { Insn::NewRowid { cursor, rowid_reg, .. } => { - let cursor = &mut state.get_btree_table_cursor(*cursor); - // TODO: make io handle rng - let rowid = return_if_io!(get_new_rowid(cursor, thread_rng())); + let rowid = { + let mut cursor = state.get_cursor(*cursor); + let cursor = cursor.as_btree_mut(); + // TODO: make io handle rng + let rowid = return_if_io!(get_new_rowid(cursor, thread_rng())); + rowid + }; state.registers[*rowid_reg] = OwnedValue::Integer(rowid); state.pc += 1; } @@ -2771,9 +2873,13 @@ impl Program { rowid_reg, target_pc, } => { - let cursor = - must_be_btree_cursor!(*cursor, self.cursor_ref, state, "NotExists"); - let exists = return_if_io!(cursor.exists(&state.registers[*rowid_reg])); + let exists = { + let mut cursor = + must_be_btree_cursor!(*cursor, self.cursor_ref, state, "NotExists"); + let cursor = cursor.as_btree_mut(); + let exists = return_if_io!(cursor.exists(&state.registers[*rowid_reg])); + exists + }; if exists { state.pc += 1; } else {