diff --git a/core/translate/compound_select.rs b/core/translate/compound_select.rs index 8066dd459..619d07585 100644 --- a/core/translate/compound_select.rs +++ b/core/translate/compound_select.rs @@ -6,7 +6,7 @@ use crate::translate::plan::{Plan, QueryDestination, SelectPlan}; use crate::vdbe::builder::{CursorType, ProgramBuilder}; use crate::vdbe::insn::Insn; use crate::vdbe::BranchOffset; -use crate::{emit_explain, QueryMode, SymbolTable}; +use crate::{emit_explain, LimboError, QueryMode, SymbolTable}; use std::sync::Arc; use tracing::instrument; use turso_parser::ast::{CompoundOperator, Expr, Literal, SortOrder}; @@ -40,56 +40,66 @@ pub fn emit_program_for_compound_select( // Each subselect shares the same limit_ctx and offset, because the LIMIT, OFFSET applies to // the entire compound select, not just a single subselect. - let limit_ctx = limit.as_ref().map(|limit| { - let reg = program.alloc_register(); - match limit.as_ref() { - Expr::Literal(Literal::Numeric(n)) => { - if let Ok(value) = n.parse::() { - program.add_comment(program.offset(), "LIMIT counter"); - program.emit_insn(Insn::Integer { value, dest: reg }); - } else { - let value = n.parse::().unwrap(); - program.emit_insn(Insn::Real { value, dest: reg }); + let limit_ctx = limit + .as_ref() + .map(|limit| { + let reg = program.alloc_register(); + match limit.as_ref() { + Expr::Literal(Literal::Numeric(n)) => { + if let Ok(value) = n.parse::() { + program.add_comment(program.offset(), "LIMIT counter"); + program.emit_insn(Insn::Integer { value, dest: reg }); + } else { + let value = n + .parse::() + .map_err(|_| LimboError::ParseError("invalid limit".to_string()))?; + program.emit_insn(Insn::Real { value, dest: reg }); + program.add_comment(program.offset(), "LIMIT counter"); + program.emit_insn(Insn::MustBeInt { reg }); + } + } + _ => { + _ = translate_expr(program, None, limit, reg, &right_most_ctx.resolver); program.add_comment(program.offset(), "LIMIT counter"); program.emit_insn(Insn::MustBeInt { reg }); } } - _ => { - _ = translate_expr(program, None, limit, reg, &right_most_ctx.resolver); - program.add_comment(program.offset(), "LIMIT counter"); - program.emit_insn(Insn::MustBeInt { reg }); - } - } - LimitCtx::new_shared(reg) - }); - let offset_reg = offset.as_ref().map(|offset_expr| { - let reg = program.alloc_register(); - match offset_expr.as_ref() { - Expr::Literal(Literal::Numeric(n)) => { - // Compile-time constant offset - if let Ok(value) = n.parse::() { - program.emit_insn(Insn::Integer { value, dest: reg }); - } else { - let value = n.parse::().unwrap(); - program.emit_insn(Insn::Real { value, dest: reg }); + Ok::<_, LimboError>(LimitCtx::new_shared(reg)) + }) + .transpose()?; + let offset_reg = offset + .as_ref() + .map(|offset_expr| { + let reg = program.alloc_register(); + match offset_expr.as_ref() { + Expr::Literal(Literal::Numeric(n)) => { + // Compile-time constant offset + if let Ok(value) = n.parse::() { + program.emit_insn(Insn::Integer { value, dest: reg }); + } else { + let value = n + .parse::() + .map_err(|_| LimboError::ParseError("invalid offset".to_string()))?; + program.emit_insn(Insn::Real { value, dest: reg }); + } + } + _ => { + _ = translate_expr(program, None, offset_expr, reg, &right_most_ctx.resolver); } } - _ => { - _ = translate_expr(program, None, offset_expr, reg, &right_most_ctx.resolver); - } - } - program.add_comment(program.offset(), "OFFSET counter"); - program.emit_insn(Insn::MustBeInt { reg }); - let combined_reg = program.alloc_register(); - program.add_comment(program.offset(), "OFFSET + LIMIT"); - program.emit_insn(Insn::OffsetLimit { - offset_reg: reg, - combined_reg, - limit_reg: limit_ctx.as_ref().unwrap().reg_limit, - }); + program.add_comment(program.offset(), "OFFSET counter"); + program.emit_insn(Insn::MustBeInt { reg }); + let combined_reg = program.alloc_register(); + program.add_comment(program.offset(), "OFFSET + LIMIT"); + program.emit_insn(Insn::OffsetLimit { + offset_reg: reg, + combined_reg, + limit_reg: limit_ctx.as_ref().unwrap().reg_limit, + }); - reg - }); + Ok::<_, LimboError>(reg) + }) + .transpose()?; // When a compound SELECT is part of a query that yields results to a coroutine (e.g. within an INSERT clause), // we must allocate registers for the result columns to be yielded. Each subselect will then yield to diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 6d635d763..0564927bd 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -256,7 +256,7 @@ pub fn emit_query<'a>( let after_main_loop_label = program.allocate_label(); t_ctx.label_main_loop_end = Some(after_main_loop_label); - init_limit(program, t_ctx, &plan.limit, &plan.offset); + init_limit(program, t_ctx, &plan.limit, &plan.offset)?; if !plan.values.is_empty() { let reg_result_cols_start = emit_values(program, plan, t_ctx)?; @@ -427,7 +427,7 @@ fn emit_program_for_delete( let after_main_loop_label = program.allocate_label(); t_ctx.label_main_loop_end = Some(after_main_loop_label); - init_limit(program, &mut t_ctx, &plan.limit, &None); + init_limit(program, &mut t_ctx, &plan.limit, &None)?; // No rows will be read from source table loops if there is a constant false condition eg. WHERE 0 if plan.contains_constant_false_condition { @@ -879,7 +879,7 @@ fn emit_program_for_update( let after_main_loop_label = program.allocate_label(); t_ctx.label_main_loop_end = Some(after_main_loop_label); - init_limit(program, &mut t_ctx, &plan.limit, &plan.offset); + init_limit(program, &mut t_ctx, &plan.limit, &plan.offset)?; // No rows will be read from source table loops if there is a constant false condition eg. WHERE 0 if plan.contains_constant_false_condition { @@ -1953,12 +1953,12 @@ fn init_limit( t_ctx: &mut TranslateCtx, limit: &Option>, offset: &Option>, -) { +) -> Result<()> { if t_ctx.limit_ctx.is_none() && limit.is_some() { t_ctx.limit_ctx = Some(LimitCtx::new(program)); } let Some(limit_ctx) = &t_ctx.limit_ctx else { - return; + return Ok(()); }; if limit_ctx.initialize_counter { @@ -2004,7 +2004,7 @@ fn init_limit( dest: offset_reg, }); } else { - let value = n.parse::().unwrap(); + let value = n.parse::()?; program.emit_insn(Insn::Real { value, dest: limit_ctx.reg_limit, @@ -2041,6 +2041,8 @@ fn init_limit( target_pc: main_loop_end, jump_if_null: false, }); + + Ok(()) } /// We have `Expr`s which have *not* had column references bound to them,