From 3fe14f37a5ade297c1f1d37fe3b43ffe869ffd0b Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Wed, 26 Mar 2025 12:47:25 -0400 Subject: [PATCH] Create plan for Update queries --- core/translate/emitter.rs | 183 ++++++++++++++++++++++++-- core/translate/main_loop.rs | 32 ++++- core/translate/optimizer.rs | 45 ++++++- core/translate/plan.rs | 102 ++++++++++++++- core/translate/update.rs | 250 +++++++++++++++--------------------- 5 files changed, 438 insertions(+), 174 deletions(-) diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index a47b1e46b..b2cea311f 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -11,12 +11,11 @@ use crate::vdbe::{insn::Insn, BranchOffset}; use crate::{Result, SymbolTable}; use super::aggregation::emit_ungrouped_aggregation; -use super::expr::{translate_condition_expr, ConditionMetadata}; +use super::expr::{translate_condition_expr, translate_expr, ConditionMetadata}; use super::group_by::{emit_group_by, init_group_by, GroupByMetadata}; use super::main_loop::{close_loop, emit_loop, init_loop, open_loop, LeftJoinMetadata, LoopLabels}; use super::order_by::{emit_order_by, init_order_by, SortMetadata}; -use super::plan::Operation; -use super::plan::{SelectPlan, TableReference}; +use super::plan::{Operation, SelectPlan, TableReference, UpdatePlan}; use super::subquery::emit_subqueries; #[derive(Debug)] @@ -89,7 +88,7 @@ pub struct TranslateCtx<'a> { /// Used to distinguish database operations #[allow(clippy::upper_case_acronyms, dead_code)] -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Copy)] pub enum OperationMode { SELECT, INSERT, @@ -138,6 +137,7 @@ fn epilogue( program: &mut ProgramBuilder, init_label: BranchOffset, start_offset: BranchOffset, + write: bool, ) -> Result<()> { program.emit_insn(Insn::Halt { err_code: 0, @@ -145,7 +145,7 @@ fn epilogue( }); program.resolve_label(init_label, program.offset()); - program.emit_insn(Insn::Transaction { write: false }); + program.emit_insn(Insn::Transaction { write }); program.emit_constant_insns(); program.emit_insn(Insn::Goto { @@ -161,6 +161,7 @@ pub fn emit_program(program: &mut ProgramBuilder, plan: Plan, syms: &SymbolTable match plan { Plan::Select(plan) => emit_program_for_select(program, plan, syms), Plan::Delete(plan) => emit_program_for_delete(program, plan, syms), + Plan::Update(plan) => emit_program_for_update(program, plan, syms), } } @@ -179,7 +180,7 @@ fn emit_program_for_select( // Trivial exit on LIMIT 0 if let Some(limit) = plan.limit { if limit == 0 { - epilogue(program, init_label, start_offset)?; + epilogue(program, init_label, start_offset, false)?; program.result_columns = plan.result_columns; program.table_references = plan.table_references; return Ok(()); @@ -188,7 +189,7 @@ fn emit_program_for_select( // Emit main parts of query emit_query(program, &mut plan, &mut t_ctx)?; // Finalize program - epilogue(program, init_label, start_offset)?; + epilogue(program, init_label, start_offset, false)?; program.result_columns = plan.result_columns; program.table_references = plan.table_references; Ok(()) @@ -240,7 +241,7 @@ pub fn emit_query<'a>( program, t_ctx, &plan.table_references, - &OperationMode::SELECT, + OperationMode::SELECT, )?; for where_term in plan.where_clause.iter().filter(|wt| wt.is_constant()) { @@ -255,7 +256,7 @@ pub fn emit_query<'a>( &plan.table_references, &where_term.expr, condition_metadata, - &mut t_ctx.resolver, + &t_ctx.resolver, )?; program.resolve_label(jump_target_when_true, program.offset()); } @@ -305,7 +306,7 @@ fn emit_program_for_delete( // exit early if LIMIT 0 if let Some(0) = plan.limit { - epilogue(program, init_label, start_offset)?; + epilogue(program, init_label, start_offset, true)?; program.result_columns = plan.result_columns; program.table_references = plan.table_references; return Ok(()); @@ -325,7 +326,7 @@ fn emit_program_for_delete( program, &mut t_ctx, &plan.table_references, - &OperationMode::DELETE, + OperationMode::DELETE, )?; // Set up main query execution loop @@ -343,7 +344,7 @@ fn emit_program_for_delete( program.resolve_label(after_main_loop_label, program.offset()); // Finalize program - epilogue(program, init_label, start_offset)?; + epilogue(program, init_label, start_offset, true)?; program.result_columns = plan.result_columns; program.table_references = plan.table_references; Ok(()) @@ -409,3 +410,161 @@ fn emit_delete_insns( Ok(()) } + +fn emit_program_for_update( + program: &mut ProgramBuilder, + plan: UpdatePlan, + syms: &SymbolTable, +) -> Result<()> { + let (mut t_ctx, init_label, start_offset) = prologue( + program, + syms, + plan.table_references.len(), + plan.returning.as_ref().map_or(0, |r| r.len()), + )?; + + let after_main_loop_label = program.allocate_label(); + t_ctx.label_main_loop_end = Some(after_main_loop_label); + if plan.contains_constant_false_condition { + program.emit_insn(Insn::Goto { + target_pc: after_main_loop_label, + }); + } + let skip_label = program.allocate_label(); + init_loop( + program, + &mut t_ctx, + &plan.table_references, + OperationMode::UPDATE, + )?; + open_loop( + program, + &mut t_ctx, + &plan.table_references, + &plan.where_clause, + )?; + emit_update_insns(&plan, &t_ctx, program)?; + program.resolve_label(skip_label, program.offset()); + close_loop(program, &mut t_ctx, &plan.table_references)?; + + program.resolve_label(after_main_loop_label, program.offset()); + + // Finalize program + epilogue(program, init_label, start_offset, true)?; + program.result_columns = plan.returning.unwrap_or_default(); + program.table_references = plan.table_references; + Ok(()) +} + +fn emit_update_insns( + plan: &UpdatePlan, + t_ctx: &TranslateCtx, + program: &mut ProgramBuilder, +) -> crate::Result<()> { + let table_ref = &plan.table_references.first().unwrap(); + let (cursor_id, index) = match &table_ref.op { + Operation::Scan { .. } => (program.resolve_cursor_id(&table_ref.identifier), None), + Operation::Search(search) => match search { + &Search::RowidEq { .. } | Search::RowidSearch { .. } => { + (program.resolve_cursor_id(&table_ref.identifier), None) + } + Search::IndexSearch { index, .. } => ( + program.resolve_cursor_id(&table_ref.identifier), + Some((index.clone(), program.resolve_cursor_id(&index.name))), + ), + }, + _ => return Ok(()), + }; + + for cond in plan.where_clause.iter().filter(|c| c.is_constant()) { + let jump_target = program.allocate_label(); + let meta = ConditionMetadata { + jump_if_condition_is_true: false, + jump_target_when_true: jump_target, + jump_target_when_false: t_ctx.label_main_loop_end.unwrap(), + }; + translate_condition_expr( + program, + &plan.table_references, + &cond.expr, + meta, + &t_ctx.resolver, + )?; + program.resolve_label(jump_target, program.offset()); + } + let first_col_reg = program.alloc_registers(table_ref.table.columns().len()); + let rowid_reg = program.alloc_register(); + program.emit_insn(Insn::RowId { + cursor_id, + dest: rowid_reg, + }); + // if no rowid, we're done + program.emit_insn(Insn::IsNull { + reg: rowid_reg, + target_pc: t_ctx.label_main_loop_end.unwrap(), + }); + + // we scan a column at a time, loading either the column's values, or the new value + // from the Set expression, into registers so we can emit a MakeRecord and update the row. + for idx in 0..table_ref.columns().len() { + if let Some((idx, expr)) = plan.set_clauses.iter().find(|(i, _)| *i == idx) { + let target_reg = first_col_reg + idx; + translate_expr( + program, + Some(&plan.table_references), + expr, + target_reg, + &t_ctx.resolver, + )?; + } else { + let table_column = table_ref.table.columns().get(idx).unwrap(); + let column_idx_in_index = index.as_ref().and_then(|(idx, _)| { + idx.columns + .iter() + .position(|c| Some(&c.name) == table_column.name.as_ref()) + }); + program.emit_insn(Insn::Column { + cursor_id: *index + .as_ref() + .and_then(|(_, id)| { + if column_idx_in_index.is_some() { + Some(id) + } else { + None + } + }) + .unwrap_or(&cursor_id), + column: column_idx_in_index.unwrap_or(idx), + dest: first_col_reg + idx, + }); + } + } + let record_reg = program.alloc_register(); + program.emit_insn(Insn::MakeRecord { + start_reg: first_col_reg, + count: table_ref.columns().len(), + dest_reg: record_reg, + }); + program.emit_insn(Insn::InsertAsync { + cursor: cursor_id, + key_reg: rowid_reg, + record_reg, + flag: 0, + }); + program.emit_insn(Insn::InsertAwait { cursor_id }); + + if let Some(limit) = plan.limit { + let limit_reg = program.alloc_register(); + program.emit_insn(Insn::Integer { + value: limit as i64, + dest: limit_reg, + }); + program.mark_last_insn_constant(); + program.emit_insn(Insn::DecrJumpZero { + reg: limit_reg, + target_pc: t_ctx.label_main_loop_end.unwrap(), + }) + } + // TODO(pthorpe): handle RETURNING clause + Ok(()) +} diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index b70541c77..7b51a2328 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -59,7 +59,7 @@ pub fn init_loop( program: &mut ProgramBuilder, t_ctx: &mut TranslateCtx, tables: &[TableReference], - mode: &OperationMode, + mode: OperationMode, ) -> Result<()> { assert!( t_ctx.meta_left_joins.len() == tables.len(), @@ -90,16 +90,24 @@ pub fn init_loop( }, ); match (mode, &table.table) { - (OperationMode::SELECT, Table::BTree(_)) => { - let root_page = table.btree().unwrap().root_page; + (OperationMode::SELECT, Table::BTree(btree)) => { + let root_page = btree.root_page; program.emit_insn(Insn::OpenReadAsync { cursor_id, root_page, }); program.emit_insn(Insn::OpenReadAwait {}); } - (OperationMode::DELETE, Table::BTree(_)) => { - let root_page = table.btree().unwrap().root_page; + (OperationMode::DELETE, Table::BTree(btree)) => { + let root_page = btree.root_page; + program.emit_insn(Insn::OpenWriteAsync { + cursor_id, + root_page, + }); + program.emit_insn(Insn::OpenWriteAwait {}); + } + (OperationMode::UPDATE, Table::BTree(btree)) => { + let root_page = btree.root_page; program.emit_insn(Insn::OpenWriteAsync { cursor_id, root_page, @@ -140,6 +148,13 @@ pub fn init_loop( }); program.emit_insn(Insn::OpenWriteAwait {}); } + OperationMode::UPDATE => { + program.emit_insn(Insn::OpenWriteAsync { + cursor_id: table_cursor_id, + root_page: table.table.get_root_page(), + }); + program.emit_insn(Insn::OpenWriteAwait {}); + } _ => { unimplemented!() } @@ -166,6 +181,13 @@ pub fn init_loop( }); program.emit_insn(Insn::OpenWriteAwait {}); } + OperationMode::UPDATE => { + program.emit_insn(Insn::OpenWriteAsync { + cursor_id: index_cursor_id, + root_page: index.root_page, + }); + program.emit_insn(Insn::OpenWriteAwait {}); + } _ => { unimplemented!() } diff --git a/core/translate/optimizer.rs b/core/translate/optimizer.rs index fed01c738..6fa7f9619 100644 --- a/core/translate/optimizer.rs +++ b/core/translate/optimizer.rs @@ -9,13 +9,14 @@ use crate::{ use super::plan::{ DeletePlan, Direction, IterationDirection, Operation, Plan, Search, SelectPlan, TableReference, - WhereTerm, + UpdatePlan, WhereTerm, }; pub fn optimize_plan(plan: &mut Plan, schema: &Schema) -> Result<()> { match plan { Plan::Select(plan) => optimize_select_plan(plan, schema), Plan::Delete(plan) => optimize_delete_plan(plan, schema), + Plan::Update(plan) => optimize_update_plan(plan, schema), } } @@ -63,6 +64,22 @@ fn optimize_delete_plan(plan: &mut DeletePlan, schema: &Schema) -> Result<()> { Ok(()) } +fn optimize_update_plan(plan: &mut UpdatePlan, schema: &Schema) -> Result<()> { + rewrite_exprs_update(plan)?; + if let ConstantConditionEliminationResult::ImpossibleCondition = + eliminate_constant_conditions(&mut plan.where_clause)? + { + plan.contains_constant_false_condition = true; + return Ok(()); + } + use_indexes( + &mut plan.table_references, + &schema.indexes, + &mut plan.where_clause, + )?; + Ok(()) +} + fn optimize_subqueries(plan: &mut SelectPlan, schema: &Schema) -> Result<()> { for table in plan.table_references.iter_mut() { if let Operation::Subquery { plan, .. } = &mut table.op { @@ -89,7 +106,7 @@ fn query_is_already_ordered_by( Search::RowidEq { .. } => Ok(key.is_rowid_alias_of(0)), Search::RowidSearch { .. } => Ok(key.is_rowid_alias_of(0)), Search::IndexSearch { index, .. } => { - let index_rc = key.check_index_scan(0, &table_reference, available_indexes)?; + let index_rc = key.check_index_scan(0, table_reference, available_indexes)?; let index_is_the_same = index_rc .map(|irc| Arc::ptr_eq(index, &irc)) .unwrap_or(false); @@ -104,7 +121,7 @@ fn eliminate_unnecessary_orderby(plan: &mut SelectPlan, schema: &Schema) -> Resu if plan.order_by.is_none() { return Ok(()); } - if plan.table_references.len() == 0 { + if plan.table_references.is_empty() { return Ok(()); } @@ -154,7 +171,7 @@ fn use_indexes( if let Some(index_search) = try_extract_index_search_expression( cond, table_index, - &table_reference, + table_reference, available_indexes, )? { where_clause.remove(i); @@ -249,6 +266,26 @@ fn rewrite_exprs_delete(plan: &mut DeletePlan) -> Result<()> { Ok(()) } +fn rewrite_exprs_update(plan: &mut UpdatePlan) -> Result<()> { + if let Some(rc) = plan.returning.as_mut() { + for rc in rc.iter_mut() { + rewrite_expr(&mut rc.expr)?; + } + } + for (_, expr) in plan.set_clauses.iter_mut() { + rewrite_expr(expr)?; + } + for cond in plan.where_clause.iter_mut() { + rewrite_expr(&mut cond.expr)?; + } + if let Some(order_by) = &mut plan.order_by { + for (expr, _) in order_by.iter_mut() { + rewrite_expr(expr)?; + } + } + Ok(()) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ConstantPredicate { AlwaysTrue, diff --git a/core/translate/plan.rs b/core/translate/plan.rs index 44f9c2dbd..666dec84f 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -7,16 +7,13 @@ use std::{ sync::Arc, }; +use crate::schema::{PseudoTable, Type}; use crate::{ function::AggFunc, schema::{BTreeTable, Column, Index, Table}, vdbe::BranchOffset, VirtualTable, }; -use crate::{ - schema::{PseudoTable, Type}, - translate::plan::Plan::{Delete, Select}, -}; #[derive(Debug, Clone)] pub struct ResultSetColumn { @@ -112,6 +109,7 @@ impl Ord for EvalAt { pub enum Plan { Select(SelectPlan), Delete(DeletePlan), + Update(UpdatePlan), } /// The type of the query, either top level or subquery @@ -170,6 +168,23 @@ pub struct DeletePlan { pub contains_constant_false_condition: bool, } +#[derive(Debug, Clone)] +pub struct UpdatePlan { + // list of table references, table being updated is always first + pub table_references: Vec, + // which columns are being updated and what they are being set to + pub set_clauses: Vec<(usize, ast::Expr)>, // (column_index, expression) + // where clause split into a vec at 'AND' boundaries. + pub where_clause: Vec, + pub order_by: Option>, + // TODO: support optional LIMIT + pub limit: Option, + // optional RETURNING clause + pub returning: Option>, + // whether the WHERE clause is always false + pub contains_constant_false_condition: bool, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub enum IterationDirection { Forwards, @@ -370,8 +385,9 @@ impl Display for Aggregate { impl Display for Plan { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { - Select(select_plan) => select_plan.fmt(f), - Delete(delete_plan) => delete_plan.fmt(f), + Self::Select(select_plan) => select_plan.fmt(f), + Self::Delete(delete_plan) => delete_plan.fmt(f), + Self::Update(update_plan) => update_plan.fmt(f), } } } @@ -461,3 +477,77 @@ impl Display for DeletePlan { Ok(()) } } + +impl fmt::Display for UpdatePlan { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "QUERY PLAN")?; + + for (i, reference) in self.table_references.iter().enumerate() { + let is_last = i == self.table_references.len() - 1; + let indent = if i == 0 { + if is_last { "`--" } else { "|--" }.to_string() + } else { + format!( + " {}{}", + "| ".repeat(i - 1), + if is_last { "`--" } else { "|--" } + ) + }; + + match &reference.op { + Operation::Scan { .. } => { + let table_name = if reference.table.get_name() == reference.identifier { + reference.identifier.clone() + } else { + format!("{} AS {}", reference.table.get_name(), reference.identifier) + }; + + if i == 0 { + writeln!(f, "{}UPDATE {}", indent, table_name)?; + } else { + writeln!(f, "{}SCAN {}", indent, table_name)?; + } + } + Operation::Search(search) => match search { + Search::RowidEq { .. } | Search::RowidSearch { .. } => { + writeln!( + f, + "{}SEARCH {} USING INTEGER PRIMARY KEY (rowid=?)", + indent, reference.identifier + )?; + } + Search::IndexSearch { index, .. } => { + writeln!( + f, + "{}SEARCH {} USING INDEX {}", + indent, reference.identifier, index.name + )?; + } + }, + Operation::Subquery { plan, .. } => { + writeln!(f, "{}SUBQUERY {}", indent, reference.identifier)?; + for line in format!("{}", plan).lines() { + writeln!(f, "{} {}", indent, line)?; + } + } + } + } + if let Some(order_by) = &self.order_by { + writeln!(f, "ORDER BY:")?; + for (expr, dir) in order_by { + writeln!(f, " - {} {}", expr, dir)?; + } + } + if let Some(limit) = self.limit { + writeln!(f, "LIMIT: {}", limit)?; + } + if let Some(ret) = &self.returning { + writeln!(f, "RETURNING:")?; + for col in ret { + writeln!(f, " - {}", col.expr)?; + } + } + + Ok(()) + } +} diff --git a/core/translate/update.rs b/core/translate/update.rs index ac6fb3999..e552777d6 100644 --- a/core/translate/update.rs +++ b/core/translate/update.rs @@ -3,20 +3,17 @@ use crate::{ bail_parse_error, schema::{Schema, Table}, util::normalize_ident, - vdbe::{ - builder::{CursorType, ProgramBuilder, ProgramBuilderOpts, QueryMode}, - insn::Insn, - }, + vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode}, SymbolTable, }; -use limbo_sqlite3_parser::ast::Update; +use limbo_sqlite3_parser::ast::{self, Expr, ResultColumn, SortOrder, Update}; -use super::planner::bind_column_references; -use super::{ - emitter::Resolver, - expr::{translate_condition_expr, translate_expr, ConditionMetadata}, - plan::TableReference, +use super::emitter::emit_program; +use super::optimizer::optimize_plan; +use super::plan::{ + Direction, IterationDirection, Plan, ResultSetColumn, TableReference, UpdatePlan, }; +use super::planner::{bind_column_references, parse_limit, parse_where}; /* * Update is simple. By default we scan the table, and for each row, we check the WHERE @@ -53,6 +50,8 @@ pub fn translate_update( body: &mut Update, syms: &SymbolTable, ) -> crate::Result { + let mut plan = prepare_update_plan(schema, body)?; + optimize_plan(&mut plan, schema)?; // TODO: freestyling these numbers let mut program = ProgramBuilder::new(ProgramBuilderOpts { query_mode, @@ -60,157 +59,114 @@ pub fn translate_update( approx_num_insns: 20, approx_num_labels: 4, }); + emit_program(&mut program, plan, syms)?; + Ok(program) +} - if body.with.is_some() { - bail_parse_error!("WITH clause is not supported"); - } - if body.or_conflict.is_some() { - bail_parse_error!("ON CONFLICT clause is not supported"); - } +pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result { let table_name = &body.tbl_name.name; let table = match schema.get_table(table_name.0.as_str()) { Some(table) => table, None => bail_parse_error!("Parse error: no such table: {}", table_name), }; - if let Table::Virtual(_) = table.as_ref() { - bail_parse_error!("vtable update not yet supported"); - } - let resolver = Resolver::new(syms); - - let init_label = program.allocate_label(); - program.emit_insn(Insn::Init { - target_pc: init_label, - }); - let start_offset = program.offset(); let Some(btree_table) = table.btree() else { - crate::bail_corrupt_error!("Parse error: no such table: {}", table_name); + bail_parse_error!("Parse error: no such table: {}", table_name); }; - let cursor_id = program.alloc_cursor_id( - Some(table_name.0.clone()), - CursorType::BTreeTable(btree_table.clone()), - ); - let root_page = btree_table.root_page; - program.emit_insn(Insn::OpenWriteAsync { - cursor_id, - root_page, - }); - program.emit_insn(Insn::OpenWriteAwait {}); - - let end_label = program.allocate_label(); - program.emit_insn(Insn::RewindAsync { cursor_id }); - program.emit_insn(Insn::RewindAwait { - cursor_id, - pc_if_empty: end_label, - }); - let first_col_reg = program.alloc_registers(btree_table.columns.len()); - let referenced_tables = vec![TableReference { + let mut iter_dir = None; + if let Some(order_by) = body.order_by.as_ref() { + if !order_by.is_empty() { + if let Some(order) = order_by.first().unwrap().order { + iter_dir = Some(match order { + SortOrder::Asc => IterationDirection::Forwards, + SortOrder::Desc => IterationDirection::Backwards, + }); + } + } + } + let table_references = vec![TableReference { table: Table::BTree(btree_table.clone()), identifier: table_name.0.clone(), - op: Operation::Scan { iter_dir: None }, + op: Operation::Scan { iter_dir }, join_info: None, }]; + let set_clauses = body + .sets + .iter_mut() + .map(|set| { + let ident = normalize_ident(set.col_names[0].0.as_str()); + let col_index = btree_table + .columns + .iter() + .enumerate() + .find(|(_, col)| { + col.name + .as_ref() + .unwrap_or(&String::new()) + .eq_ignore_ascii_case(&ident) + }) + .map(|(i, _)| i) + .unwrap(); + let _ = bind_column_references(&mut set.expr, &table_references, None); + (col_index, set.expr.clone()) + }) + .collect::>(); - // store the (col_index, Expr value) of each 'Set' - // if a column declared here isn't found: error - let mut update_idxs = Vec::with_capacity(body.sets.len()); - for s in body.sets.iter_mut() { - let ident = normalize_ident(s.col_names[0].0.as_str()); - if let Some((i, _)) = btree_table.columns.iter().enumerate().find(|(_, col)| { - col.name - .as_ref() - .unwrap_or(&String::new()) - .eq_ignore_ascii_case(&ident) - }) { - bind_column_references(&mut s.expr, &referenced_tables, None)?; - update_idxs.push((i, &s.expr)); - } else { - bail_parse_error!("column {} not found", ident); + let mut where_clause = vec![]; + let mut result_columns = vec![]; + if let Some(returning) = &mut body.returning { + for rc in returning.iter_mut() { + if let ResultColumn::Expr(expr, alias) = rc { + bind_column_references(expr, &table_references, None)?; + result_columns.push(ResultSetColumn { + expr: expr.clone(), + alias: alias.as_ref().and_then(|a| { + if let ast::As::As(name) = a { + Some(name.to_string()) + } else { + None + } + }), + contains_aggregates: false, + }); + } else { + bail_parse_error!("Only expressions are allowed in RETURNING clause"); + } } } - - let loop_start = program.offset(); - let skip_label = program.allocate_label(); - if let Some(where_clause) = body.where_clause.as_mut() { - bind_column_references(where_clause, &referenced_tables, None)?; - translate_condition_expr( - &mut program, - &referenced_tables, - where_clause, - ConditionMetadata { - jump_if_condition_is_true: false, - jump_target_when_true: crate::vdbe::BranchOffset::Placeholder, - jump_target_when_false: skip_label, - }, - &resolver, - )?; - } - let rowid_reg = program.alloc_register(); - program.emit_insn(Insn::RowId { - cursor_id, - dest: rowid_reg, + let order_by = body.order_by.as_ref().map(|order| { + order + .iter() + .map(|o| { + ( + o.expr.clone(), + o.order + .map(|s| match s { + SortOrder::Asc => Direction::Ascending, + SortOrder::Desc => Direction::Descending, + }) + .unwrap_or(Direction::Ascending), + ) + }) + .collect() }); - // if no rowid, we're done - program.emit_insn(Insn::IsNull { - reg: rowid_reg, - target_pc: end_label, - }); - - // we scan a column at a time, loading either the column's values, or the new value - // from the Set expression, into registers so we can emit a MakeRecord and update the row. - for idx in 0..btree_table.columns.len() { - if let Some((idx, expr)) = update_idxs.iter().find(|(i, _)| *i == idx) { - let target_reg = first_col_reg + idx; - translate_expr( - &mut program, - Some(&referenced_tables), - expr, - target_reg, - &resolver, - )?; - } else { - program.emit_insn(Insn::Column { - cursor_id, - column: idx, - dest: first_col_reg + idx, - }); - } - } - let record_reg = program.alloc_register(); - program.emit_insn(Insn::MakeRecord { - start_reg: first_col_reg, - count: btree_table.columns.len(), - dest_reg: record_reg, - }); - - program.emit_insn(Insn::InsertAsync { - cursor: cursor_id, - key_reg: rowid_reg, - record_reg, - flag: 0, - }); - program.emit_insn(Insn::InsertAwait { cursor_id }); - - // label for false `WHERE` clause: proceed to next row - program.resolve_label(skip_label, program.offset()); - program.emit_insn(Insn::NextAsync { cursor_id }); - program.emit_insn(Insn::NextAwait { - cursor_id, - pc_if_next: loop_start, - }); - - // cleanup/halt - program.resolve_label(end_label, program.offset()); - program.emit_insn(Insn::Halt { - err_code: 0, - description: String::new(), - }); - program.resolve_label(init_label, program.offset()); - program.emit_insn(Insn::Transaction { write: true }); - - program.emit_constant_insns(); - program.emit_insn(Insn::Goto { - target_pc: start_offset, - }); - program.table_references = referenced_tables.clone(); - Ok(program) + parse_where( + body.where_clause.as_ref().map(|w| *w.clone()), + &table_references, + Some(&result_columns), + &mut where_clause, + )?; + let limit = if let Some(Ok((_, limit))) = body.limit.as_ref().map(|l| parse_limit(*l.clone())) { + limit + } else { + None + }; + Ok(Plan::Update(UpdatePlan { + table_references, + set_clauses, + where_clause, + returning: Some(result_columns), + order_by, + limit, + contains_constant_false_condition: false, + })) }