aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2026-02-23 00:45:37 +01:00
Code Issues Packages Projects Releases Wiki Activity

remove assumption that translate_select is always called from a top-level context + adjust insert to use translate_select when needed

Browse Source
This commit is contained in:
pedrocarlo
2025-05-23 15:02:03 -03:00
parent fd9e0db5cc
commit bb7da39c72
5 changed files with 140 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
core/lib.rs
View File

@@ -420,6 +420,7 @@ impl Connection {
&syms,
None,
&mut table_ref_counter,
translate::plan::QueryDestination::ResultRows,
)?;
optimize_plan(
&mut plan,

173
core/translate/insert.rs
View File

@@ -1,4 +1,3 @@
use std::ops::Deref;
use std::rc::Rc;
use limbo_sqlite3_parser::ast::{
@@ -22,6 +21,8 @@ use crate::{Result, SymbolTable, VirtualTable};
use super::emitter::Resolver;
use super::expr::{translate_expr_no_constant_opt, NoConstantOptReason};
use super::optimizer::rewrite_expr;
use super::plan::QueryDestination;
use super::select::translate_select;
#[allow(clippy::too_many_arguments)]
pub fn translate_insert(
@@ -59,8 +60,8 @@ pub fn translate_insert(
let resolver = Resolver::new(syms);
if let Some(virtual_table) = &table.virtual_table() {
translate_virtual_table_insert(
&mut program,
program = translate_virtual_table_insert(
program,
virtual_table.clone(),
columns,
body,
@@ -100,7 +101,7 @@ pub fn translate_insert(
.collect::<Vec<(&String, usize, usize)>>();
let root_page = btree_table.root_page;
let inserting_multiple_rows = match body {
let inserting_multiple_rows = match &body {
InsertBody::Select(select, _) => match select.body.select.as_ref() {
OneSelect::Values(values) => values.len() > 1,
OneSelect::Select(..) => true,
@@ -108,20 +109,80 @@ pub fn translate_insert(
InsertBody::DefaultValues => false,
};
let values = match body {
InsertBody::Select(ref mut select, _) => match select.body.select.as_mut() {
OneSelect::Values(ref mut values) => values,
_ => todo!(),
},
InsertBody::DefaultValues => &mut vec![vec![]],
};
let mut param_idx = 1;
for expr in values.iter_mut().flat_map(|v| v.iter_mut()) {
rewrite_expr(expr, &mut param_idx)?;
}
let halt_label = program.allocate_label();
let loop_start_label = program.allocate_label();
let column_mappings = resolve_columns_for_insert(&table, columns, values)?;
let index_col_mappings = resolve_indicies_for_insert(schema, table.as_ref(), &column_mappings)?;
let (num_values, value) = match body {
// TODO: upsert
InsertBody::Select(select, _) => {
// Simple Common case of INSERT INTO <table> VALUES (...)
if matches!(select.body.select.as_ref(), OneSelect::Values(values) if values.len() <= 1)
{
let OneSelect::Values(mut values) = *select.body.select else {
unreachable!();
};
if values.is_empty() {
crate::bail_parse_error!("no values to insert");
}
let mut param_idx = 1;
for expr in values.iter_mut().flat_map(|v| v.iter_mut()) {
rewrite_expr(expr, &mut param_idx)?;
}
(values[0].len(), values.pop())
} else {
let yield_reg = program.alloc_register();
let jump_on_definition_label = program.allocate_label();
let start_offset_label = program.allocate_label();
program.emit_insn(Insn::InitCoroutine {
yield_reg,
jump_on_definition: jump_on_definition_label,
start_offset: start_offset_label,
});
program.preassign_label_to_next_insn(start_offset_label);
let query_destination = QueryDestination::CoroutineYield {
yield_reg,
coroutine_implementation_start: halt_label,
};
program.incr_nesting();
let result = translate_select(
query_mode,
schema,
*select,
syms,
program,
query_destination,
)?;
program = result.program;
program.decr_nesting();
program.emit_insn(Insn::EndCoroutine { yield_reg });
program.preassign_label_to_next_insn(jump_on_definition_label);
program.emit_insn(Insn::OpenWrite {
cursor_id,
root_page: RegisterOrLiteral::Literal(root_page),
name: table_name.0.clone(),
});
// Main loop
// FIXME: rollback is not implemented. E.g. if you insert 2 rows and one fails to unique constraint violation,
// the other row will still be inserted.
program.resolve_label(loop_start_label, program.offset());
program.emit_insn(Insn::Yield {
yield_reg,
end_offset: halt_label,
});
(result.num_result_cols, None)
}
}
InsertBody::DefaultValues => (0, None),
};
let column_mappings = resolve_columns_for_insert(&table, &columns, num_values)?;
// Check if rowid was provided (through INTEGER PRIMARY KEY as a rowid alias)
let rowid_alias_index = btree_table.columns.iter().position(|c| c.is_rowid_alias);
let has_user_provided_rowid = {
@@ -147,55 +208,9 @@ pub fn translate_insert(
};
let record_register = program.alloc_register();
let halt_label = program.allocate_label();
let loop_start_label = program.allocate_label();
// Multiple rows - use coroutine for value population
if inserting_multiple_rows {
let yield_reg = program.alloc_register();
let jump_on_definition_label = program.allocate_label();
let start_offset_label = program.allocate_label();
program.emit_insn(Insn::InitCoroutine {
yield_reg,
jump_on_definition: jump_on_definition_label,
start_offset: start_offset_label,
});
program.preassign_label_to_next_insn(start_offset_label);
for value in values.iter() {
populate_column_registers(
&mut program,
value,
&column_mappings,
column_registers_start,
true,
rowid_reg,
&resolver,
)?;
program.emit_insn(Insn::Yield {
yield_reg,
end_offset: halt_label,
});
}
program.emit_insn(Insn::EndCoroutine { yield_reg });
program.preassign_label_to_next_insn(jump_on_definition_label);
program.emit_insn(Insn::OpenWrite {
cursor_id,
root_page: RegisterOrLiteral::Literal(root_page),
name: table_name.0.clone(),
});
// Main loop
// FIXME: rollback is not implemented. E.g. if you insert 2 rows and one fails to unique constraint violation,
// the other row will still be inserted.
program.resolve_label(loop_start_label, program.offset());
program.emit_insn(Insn::Yield {
yield_reg,
end_offset: halt_label,
});
} else {
if !inserting_multiple_rows {
// Single row - populate registers directly
program.emit_insn(Insn::OpenWrite {
cursor_id,
@@ -205,7 +220,7 @@ pub fn translate_insert(
populate_column_registers(
&mut program,
&values[0],
&value.unwrap(),
&column_mappings,
column_registers_start,
false,
@@ -299,6 +314,7 @@ pub fn translate_insert(
_ => (),
}
let index_col_mappings = resolve_indicies_for_insert(schema, table.as_ref(), &column_mappings)?;
for index_col_mapping in index_col_mappings {
// find which cursor we opened earlier for this index
let idx_cursor_id = idx_cursors
@@ -640,24 +656,25 @@ fn populate_column_registers(
Ok(())
}
// TODO: comeback here later to apply the same improvements on select
fn translate_virtual_table_insert(
program: &mut ProgramBuilder,
mut program: ProgramBuilder,
virtual_table: Rc<VirtualTable>,
columns: &Option<DistinctNames>,
body: &InsertBody,
on_conflict: &Option<ResolveType>,
columns: Option<DistinctNames>,
mut body: InsertBody,
on_conflict: Option<ResolveType>,
resolver: &Resolver,
) -> Result<()> {
let values = match body {
InsertBody::Select(select, None) => match &select.body.select.deref() {
OneSelect::Values(values) => values,
) -> Result<ProgramBuilder> {
let (num_values, value) = match &mut body {
InsertBody::Select(select, None) => match select.body.select.as_mut() {
OneSelect::Values(values) => (1, values.pop().unwrap()),
_ => crate::bail_parse_error!("Virtual tables only support VALUES clause in INSERT"),
},
InsertBody::DefaultValues => &vec![],
InsertBody::DefaultValues => (0, vec![]),
_ => crate::bail_parse_error!("Unsupported INSERT body for virtual tables"),
};
let table = Table::Virtual(virtual_table.clone());
let column_mappings = resolve_columns_for_insert(&table, columns, values)?;
let column_mappings = resolve_columns_for_insert(&table, &columns, num_values)?;
let registers_start = program.alloc_registers(2);
/* *
@@ -674,8 +691,8 @@ fn translate_virtual_table_insert(
let values_reg = program.alloc_registers(column_mappings.len());
populate_column_registers(
program,
&values[0],
&mut program,
&value,
&column_mappings,
values_reg,
false,
@@ -699,5 +716,5 @@ fn translate_virtual_table_insert(
let halt_label = program.allocate_label();
program.resolve_label(halt_label, program.offset());
Ok(())
Ok(program)
}

10
core/translate/mod.rs
View File

@@ -233,7 +233,15 @@ pub fn translate_inner(
ast::Stmt::Rollback { .. } => bail_parse_error!("ROLLBACK not supported yet"),
ast::Stmt::Savepoint(_) => bail_parse_error!("SAVEPOINT not supported yet"),
ast::Stmt::Select(select) => {
translate_select(query_mode, schema, *select, syms, program)?.program
translate_select(
query_mode,
schema,
*select,
syms,
program,
plan::QueryDestination::ResultRows,
)?
.program
}
ast::Stmt::Update(mut update) => translate_update(
query_mode,

39
core/translate/planner.rs
View File

@@ -293,15 +293,21 @@ fn parse_from_clause_table<'a>(
crate::bail_parse_error!("Table {} not found", normalized_qualified_name);
}
ast::SelectTable::Select(subselect, maybe_alias) => {
let Plan::Select(mut subplan) =
prepare_select_plan(schema, *subselect, syms, Some(scope), table_ref_counter)?
else {
crate::bail_parse_error!("Only non-compound SELECT queries are currently supported in FROM clause subqueries");
};
subplan.query_destination = QueryDestination::CoroutineYield {
let query_destination = QueryDestination::CoroutineYield {
yield_reg: usize::MAX, // will be set later in bytecode emission
coroutine_implementation_start: BranchOffset::Placeholder, // will be set later in bytecode emission
};
let Plan::Select(subplan) = prepare_select_plan(
schema,
*subselect,
syms,
Some(scope),
table_ref_counter,
query_destination,
)?
else {
crate::bail_parse_error!("Only non-compound SELECT queries are currently supported in FROM clause subqueries");
};
let cur_table_index = scope.tables.len();
let identifier = maybe_alias
.map(|a| match a {
@@ -449,16 +455,23 @@ pub fn parse_from<'a>(
crate::bail_parse_error!("duplicate WITH table name {}", cte.tbl_name.0);
}
// CTE can refer to other CTEs that came before it, plus any schema tables or tables in the outer scope.
let cte_plan =
prepare_select_plan(schema, *cte.select, syms, Some(&scope), table_ref_counter)?;
let Plan::Select(mut cte_plan) = cte_plan else {
crate::bail_parse_error!("Only SELECT queries are currently supported in CTEs");
};
cte_plan.query_destination = QueryDestination::CoroutineYield {
// CTE can be rewritten as a subquery.
let query_destination = QueryDestination::CoroutineYield {
yield_reg: usize::MAX, // will be set later in bytecode emission
coroutine_implementation_start: BranchOffset::Placeholder, // will be set later in bytecode emission
};
// CTE can refer to other CTEs that came before it, plus any schema tables or tables in the outer scope.
let cte_plan = prepare_select_plan(
schema,
*cte.select,
syms,
Some(&scope),
table_ref_counter,
query_destination,
)?;
let Plan::Select(cte_plan) = cte_plan else {
crate::bail_parse_error!("Only SELECT queries are currently supported in CTEs");
};
scope.ctes.push(Cte {
name: cte_name_normalized,
plan: cte_plan,

11
core/translate/select.rs
View File

@@ -30,6 +30,7 @@ pub fn translate_select(
select: ast::Select,
syms: &SymbolTable,
mut program: ProgramBuilder,
query_destination: QueryDestination,
) -> Result<TranslateSelectResult> {
let mut select_plan = prepare_select_plan(
schema,
@@ -37,6 +38,7 @@ pub fn translate_select(
syms,
None,
&mut program.table_reference_counter,
query_destination,
)?;
optimize_plan(&mut select_plan, schema)?;
let opts = match &select_plan {
@@ -82,6 +84,7 @@ pub fn prepare_select_plan<'a>(
syms: &SymbolTable,
outer_scope: Option<&'a Scope<'a>>,
table_ref_counter: &mut TableRefIdCounter,
query_destination: QueryDestination,
) -> Result<Plan> {
let compounds = select.body.compounds.take();
match compounds {
@@ -96,6 +99,7 @@ pub fn prepare_select_plan<'a>(
syms,
outer_scope,
table_ref_counter,
query_destination,
)?))
}
Some(compounds) => {
@@ -108,6 +112,7 @@ pub fn prepare_select_plan<'a>(
syms,
outer_scope,
table_ref_counter,
query_destination.clone(),
)?;
let mut rest = Vec::with_capacity(compounds.len());
for CompoundSelect { select, operator } in compounds {
@@ -128,6 +133,7 @@ pub fn prepare_select_plan<'a>(
syms,
outer_scope,
table_ref_counter,
query_destination.clone(),
)?;
rest.push((plan, operator));
}
@@ -177,6 +183,7 @@ fn prepare_one_select_plan<'a>(
syms: &SymbolTable,
outer_scope: Option<&'a Scope<'a>>,
table_ref_counter: &mut TableRefIdCounter,
query_destination: QueryDestination,
) -> Result<SelectPlan> {
match select {
ast::OneSelect::Select(select_inner) => {
@@ -246,7 +253,7 @@ fn prepare_one_select_plan<'a>(
limit: None,
offset: None,
contains_constant_false_condition: false,
query_destination: QueryDestination::ResultRows,
query_destination,
distinctness: Distinctness::from_ast(distinctness.as_ref()),
values: vec![],
};
@@ -560,7 +567,7 @@ fn prepare_one_select_plan<'a>(
limit: None,
offset: None,
contains_constant_false_condition: false,
query_destination: QueryDestination::ResultRows,
query_destination,
distinctness: Distinctness::NonDistinct,
values,
};