From ffcadd00aefcfa8b506b5e2e48d338f95377ea37 Mon Sep 17 00:00:00 2001 From: bit-aloo Date: Thu, 21 Aug 2025 20:10:39 +0530 Subject: [PATCH] evaluate limit or offset expr --- core/translate/compound_select.rs | 84 ++++++++++++--- core/translate/display.rs | 2 +- core/translate/emitter.rs | 126 +++++++++++++++------- core/translate/planner.rs | 43 ++------ core/translate/result_row.rs | 170 +++++++++++++++++++++++++++++- 5 files changed, 331 insertions(+), 94 deletions(-) diff --git a/core/translate/compound_select.rs b/core/translate/compound_select.rs index 15fb4d841..5316a744f 100644 --- a/core/translate/compound_select.rs +++ b/core/translate/compound_select.rs @@ -1,6 +1,7 @@ use crate::schema::{Index, IndexColumn, Schema}; use crate::translate::emitter::{emit_query, LimitCtx, TranslateCtx}; use crate::translate::plan::{Plan, QueryDestination, SelectPlan}; +use crate::translate::result_row::{build_limit_offset_expr, try_fold_expr_to_i64}; use crate::vdbe::builder::{CursorType, ProgramBuilder}; use crate::vdbe::insn::Insn; use crate::vdbe::BranchOffset; @@ -31,8 +32,8 @@ pub fn emit_program_for_compound_select( let right_plan = right_most.clone(); // Trivial exit on LIMIT 0 - if let Some(limit) = limit { - if *limit == 0 { + if let Some(expr) = limit { + if let Some(0) = try_fold_expr_to_i64(expr) { program.result_columns = right_plan.result_columns; program.table_references.extend(right_plan.table_references); return Ok(()); @@ -41,26 +42,79 @@ 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.map(|limit| { + let limit_ctx = limit.clone().map(|limit| { let reg = program.alloc_register(); - program.emit_insn(Insn::Integer { - value: limit as i64, - dest: reg, - }); + if let Some(val) = try_fold_expr_to_i64(&limit) { + // Compile-time constant limit + program.emit_insn(Insn::Integer { + value: val, + dest: reg, + }); + } else { + program.add_comment(program.offset(), "OFFSET expr"); + + let label_zero = program.allocate_label(); + + build_limit_offset_expr(program, reg, &limit); + + program.emit_insn(Insn::MustBeInt { reg }); + + program.emit_insn(Insn::IfNeg { + reg, + target_pc: label_zero, + }); + program.emit_insn(Insn::IsNull { + reg, + target_pc: label_zero, + }); + + program.preassign_label_to_next_insn(label_zero); + program.emit_insn(Insn::Integer { + value: 0, + dest: reg, + }); + } LimitCtx::new_shared(reg) }); - let offset_reg = offset.map(|offset| { + let offset_reg = offset.as_ref().map(|offset_expr| { let reg = program.alloc_register(); - program.emit_insn(Insn::Integer { - value: offset as i64, - dest: reg, - }); + + if let Some(val) = try_fold_expr_to_i64(offset_expr) { + // Compile-time constant offset + program.emit_insn(Insn::Integer { + value: val, + dest: reg, + }); + } else { + program.add_comment(program.offset(), "OFFSET expr"); + + let label_zero = program.allocate_label(); + + build_limit_offset_expr(program, reg, &offset_expr); + + program.emit_insn(Insn::MustBeInt { reg }); + + program.emit_insn(Insn::IfNeg { + reg, + target_pc: label_zero, + }); + program.emit_insn(Insn::IsNull { + reg, + target_pc: label_zero, + }); + + program.preassign_label_to_next_insn(label_zero); + program.emit_insn(Insn::Integer { + value: 0, + dest: reg, + }); + } let combined_reg = program.alloc_register(); program.emit_insn(Insn::OffsetLimit { offset_reg: reg, combined_reg, - limit_reg: limit_ctx.unwrap().reg_limit, + limit_reg: limit_ctx.as_ref().unwrap().reg_limit, }); reg @@ -137,8 +191,8 @@ fn emit_compound_select( let compound_select = Plan::CompoundSelect { left, right_most: plan, - limit, - offset, + limit: limit.clone(), + offset: offset.clone(), order_by, }; emit_compound_select( diff --git a/core/translate/display.rs b/core/translate/display.rs index 631e5a295..3e79ef223 100644 --- a/core/translate/display.rs +++ b/core/translate/display.rs @@ -217,7 +217,7 @@ impl fmt::Display for UpdatePlan { )?; } } - if let Some(limit) = self.limit { + if let Some(limit) = self.limit.clone() { writeln!(f, "LIMIT: {limit}")?; } if let Some(ret) = &self.returning { diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 29cbd210a..70b2f8310 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -26,6 +26,7 @@ use crate::schema::{BTreeTable, Column, Schema, Table}; use crate::translate::compound_select::emit_program_for_compound_select; use crate::translate::expr::{emit_returning_results, ReturningValueRegisters}; use crate::translate::plan::{DeletePlan, Plan, QueryDestination, Search}; +use crate::translate::result_row::{build_limit_offset_expr, try_fold_expr_to_i64}; use crate::translate::values::emit_values; use crate::util::exprs_are_equivalent; use crate::vdbe::builder::{CursorKey, CursorType, ProgramBuilder}; @@ -227,8 +228,8 @@ fn emit_program_for_select( ); // Trivial exit on LIMIT 0 - if let Some(limit) = plan.limit { - if limit == 0 { + if let Some(limit) = plan.limit.as_ref().and_then(try_fold_expr_to_i64) { + if limit <= 0 { program.result_columns = plan.result_columns; program.table_references.extend(plan.table_references); return Ok(()); @@ -256,7 +257,7 @@ pub fn emit_query<'a>( // Emit subqueries first so the results can be read in the main query loop. emit_subqueries(program, t_ctx, &mut plan.table_references)?; - init_limit(program, t_ctx, plan.limit, plan.offset); + init_limit(program, t_ctx, plan.limit.clone(), plan.offset.clone()); // No rows will be read from source table loops if there is a constant false condition eg. WHERE 0 // however an aggregation might still happen, @@ -404,10 +405,12 @@ fn emit_program_for_delete( ); // exit early if LIMIT 0 - if let Some(0) = plan.limit { - program.result_columns = plan.result_columns; - program.table_references.extend(plan.table_references); - return Ok(()); + if let Some(limit) = plan.limit.as_ref().and_then(try_fold_expr_to_i64) { + if limit <= 0 { + program.result_columns = plan.result_columns; + program.table_references.extend(plan.table_references); + return Ok(()); + } } init_limit(program, &mut t_ctx, plan.limit, None); @@ -660,13 +663,15 @@ fn emit_program_for_update( ); // Exit on LIMIT 0 - if let Some(0) = plan.limit { - program.result_columns = plan.returning.unwrap_or_default(); - program.table_references.extend(plan.table_references); - return Ok(()); + if let Some(limit) = plan.limit.as_ref().and_then(try_fold_expr_to_i64) { + if limit <= 0 { + program.result_columns = plan.returning.unwrap_or_default(); + program.table_references.extend(plan.table_references); + return Ok(()); + } } - init_limit(program, &mut t_ctx, plan.limit, plan.offset); + init_limit(program, &mut t_ctx, plan.limit.clone(), plan.offset.clone()); 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 { @@ -1541,41 +1546,90 @@ pub fn emit_cdc_insns( }); Ok(()) } - /// Initialize the limit/offset counters and registers. /// In case of compound SELECTs, the limit counter is initialized only once, /// hence [LimitCtx::initialize_counter] being false in those cases. fn init_limit( program: &mut ProgramBuilder, t_ctx: &mut TranslateCtx, - limit: Option, - offset: Option, + limit: Option, + offset: Option, ) { - if t_ctx.limit_ctx.is_none() { - t_ctx.limit_ctx = limit.map(|_| LimitCtx::new(program)); + 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 { + let Some(limit_ctx) = &t_ctx.limit_ctx else { return; }; if limit_ctx.initialize_counter { - program.emit_insn(Insn::Integer { - value: limit.expect("limit must be Some if limit_ctx is Some") as i64, - dest: limit_ctx.reg_limit, - }); + if let Some(expr) = limit { + if let Some(value) = try_fold_expr_to_i64(&expr) { + program.emit_insn(Insn::Integer { + value, + dest: limit_ctx.reg_limit, + }); + } else { + let r = limit_ctx.reg_limit; + program.add_comment(program.offset(), "OFFSET expr"); + let label_zero = program.allocate_label(); + build_limit_offset_expr(program, r, &expr); + program.emit_insn(Insn::MustBeInt { reg: r }); + program.emit_insn(Insn::IfNeg { + reg: r, + target_pc: label_zero, + }); + program.emit_insn(Insn::IsNull { + reg: r, + target_pc: label_zero, + }); + program.preassign_label_to_next_insn(label_zero); + program.emit_insn(Insn::Integer { value: 0, dest: r }); + } + } } - if t_ctx.reg_offset.is_none() && offset.is_some_and(|n| n.ne(&0)) { - let reg = program.alloc_register(); - t_ctx.reg_offset = Some(reg); - program.emit_insn(Insn::Integer { - value: offset.unwrap() as i64, - dest: reg, - }); - let combined_reg = program.alloc_register(); - t_ctx.reg_limit_offset_sum = Some(combined_reg); - program.emit_insn(Insn::OffsetLimit { - limit_reg: t_ctx.limit_ctx.unwrap().reg_limit, - offset_reg: reg, - combined_reg, - }); + + if t_ctx.reg_offset.is_none() { + if let Some(expr) = offset { + if let Some(value) = try_fold_expr_to_i64(&expr) { + if value != 0 { + let reg = program.alloc_register(); + t_ctx.reg_offset = Some(reg); + program.emit_insn(Insn::Integer { value, dest: reg }); + let combined_reg = program.alloc_register(); + t_ctx.reg_limit_offset_sum = Some(combined_reg); + program.emit_insn(Insn::OffsetLimit { + limit_reg: limit_ctx.reg_limit, + offset_reg: reg, + combined_reg, + }); + } + } else { + + let reg = program.alloc_register(); + t_ctx.reg_offset = Some(reg); + let r = reg; + program.add_comment(program.offset(), "OFFSET expr"); + let label_zero = program.allocate_label(); + build_limit_offset_expr(program, r, &expr); + program.emit_insn(Insn::MustBeInt { reg: r }); + program.emit_insn(Insn::IfNeg { + reg: r, + target_pc: label_zero, + }); + program.emit_insn(Insn::IsNull { + reg: r, + target_pc: label_zero, + }); + program.preassign_label_to_next_insn(label_zero); + program.emit_insn(Insn::Integer { value: 0, dest: r }); + let combined_reg = program.alloc_register(); + t_ctx.reg_limit_offset_sum = Some(combined_reg); + program.emit_insn(Insn::OffsetLimit { + limit_reg: limit_ctx.reg_limit, + offset_reg: reg, + combined_reg, + }); + } + } } } diff --git a/core/translate/planner.rs b/core/translate/planner.rs index 43f012875..c8c34d4ca 100644 --- a/core/translate/planner.rs +++ b/core/translate/planner.rs @@ -22,7 +22,7 @@ use crate::{ use turso_parser::ast::Literal::Null; use turso_parser::ast::{ self, As, Expr, FromClause, JoinType, Limit, Literal, Materialized, QualifiedName, - TableInternalId, UnaryOperator, With, + TableInternalId, With, }; pub const ROWID: &str = "rowid"; @@ -1106,43 +1106,12 @@ fn parse_join( Ok(()) } -pub fn parse_limit(limit: &Limit) -> Result<(Option, Option)> { - let offset_val = match &limit.offset { - Some(offset_expr) => match offset_expr.as_ref() { - Expr::Literal(ast::Literal::Numeric(n)) => n.parse().ok(), - // If OFFSET is negative, the result is as if OFFSET is zero - Expr::Unary(UnaryOperator::Negative, expr) => { - if let Expr::Literal(ast::Literal::Numeric(ref n)) = &**expr { - n.parse::().ok().map(|num| -num) - } else { - crate::bail_parse_error!("Invalid OFFSET clause"); - } - } - _ => crate::bail_parse_error!("Invalid OFFSET clause"), - }, - None => Some(0), - }; +pub fn parse_limit(limit: &Limit) -> Result<(Option, Option)> { + let limit_expr = Some(limit.expr.clone()); - if let Expr::Literal(ast::Literal::Numeric(n)) = limit.expr.as_ref() { - Ok((n.parse().ok(), offset_val)) - } else if let Expr::Unary(UnaryOperator::Negative, expr) = limit.expr.as_ref() { - if let Expr::Literal(ast::Literal::Numeric(n)) = expr.as_ref() { - let limit_val = n.parse::().ok().map(|num| -num); - Ok((limit_val, offset_val)) - } else { - crate::bail_parse_error!("Invalid LIMIT clause"); - } - } else if let Expr::Id(id) = limit.expr.as_ref() { - if id.as_str().eq_ignore_ascii_case("true") { - Ok((Some(1), offset_val)) - } else if id.as_str().eq_ignore_ascii_case("false") { - Ok((Some(0), offset_val)) - } else { - crate::bail_parse_error!("Invalid LIMIT clause"); - } - } else { - crate::bail_parse_error!("Invalid LIMIT clause"); - } + let offset_expr = limit.offset.clone(); + + Ok((limit_expr, offset_expr)) } pub fn break_predicate_at_and_boundaries(predicate: &Expr, out_predicates: &mut Vec) { diff --git a/core/translate/result_row.rs b/core/translate/result_row.rs index f2b722988..e2950a435 100644 --- a/core/translate/result_row.rs +++ b/core/translate/result_row.rs @@ -1,7 +1,10 @@ +use turso_sqlite3_parser::ast::{Expr, Literal, Operator, UnaryOperator}; + use crate::{ + error::SQLITE_CONSTRAINT, vdbe::{ builder::ProgramBuilder, - insn::{IdxInsertFlags, InsertFlags, Insn}, + insn::{CmpInsFlags, IdxInsertFlags, InsertFlags, Insn}, BranchOffset, }, Result, @@ -164,15 +167,172 @@ pub fn emit_offset( jump_to: BranchOffset, reg_offset: Option, ) { - match plan.offset { - Some(offset) if offset > 0 => { - program.add_comment(program.offset(), "OFFSET"); + let Some(offset_expr) = &plan.offset else { + return; + }; + + if let Some(val) = try_fold_expr_to_i64(offset_expr) { + if val > 0 { + program.add_comment(program.offset(), "OFFSET const"); program.emit_insn(Insn::IfPos { reg: reg_offset.expect("reg_offset must be Some"), target_pc: jump_to, decrement_by: 1, }); } - _ => {} + return; + } + + let r = reg_offset.expect("reg_offset must be Some"); + + program.add_comment(program.offset(), "OFFSET expr"); + + let label_zero = program.allocate_label(); + + build_limit_offset_expr(program, r, offset_expr); + + program.emit_insn(Insn::MustBeInt { reg: r }); + + program.emit_insn(Insn::IfNeg { + reg: r, + target_pc: label_zero, + }); + program.emit_insn(Insn::IsNull { + reg: r, + target_pc: label_zero, + }); + + program.emit_insn(Insn::IfPos { + reg: r, + target_pc: jump_to, + decrement_by: 1, + }); + + program.preassign_label_to_next_insn(label_zero); + program.emit_insn(Insn::Integer { value: 0, dest: r }); +} + +pub fn build_limit_offset_expr(program: &mut ProgramBuilder, r: usize, expr: &Expr) { + match expr { + Expr::Literal(Literal::Numeric(n)) => { + let value = n.parse::().unwrap_or_else(|_| { + program.emit_insn(Insn::Halt { + err_code: SQLITE_CONSTRAINT, + description: "invalid numeric literal".into(), + }); + 0 + }); + program.emit_int(value, r); + } + Expr::Unary(UnaryOperator::Negative, inner) => { + let inner_reg = program.alloc_register(); + build_limit_offset_expr(program, inner_reg, inner); + + let neg_one_reg = program.alloc_register(); + program.emit_int(-1, neg_one_reg); + + program.emit_insn(Insn::Multiply { + lhs: inner_reg, + rhs: neg_one_reg, + dest: r, + }); + } + Expr::Unary(UnaryOperator::Positive, inner) => { + let inner_reg = program.alloc_register(); + build_limit_offset_expr(program, inner_reg, inner); + program.emit_insn(Insn::Copy { + src_reg: inner_reg, + dst_reg: r, + extra_amount: 0, + }); + } + Expr::Binary(left, op, right) => { + let left_reg = program.alloc_register(); + let right_reg = program.alloc_register(); + build_limit_offset_expr(program, left_reg, left); + build_limit_offset_expr(program, right_reg, right); + + match op { + Operator::Add => { + program.emit_insn(Insn::Add { + lhs: left_reg, + rhs: right_reg, + dest: r, + }); + } + Operator::Subtract => { + program.emit_insn(Insn::Subtract { + lhs: left_reg, + rhs: right_reg, + dest: r, + }); + } + Operator::Multiply => { + program.emit_insn(Insn::Multiply { + lhs: left_reg, + rhs: right_reg, + dest: r, + }); + } + Operator::Divide => { + let zero_reg = program.alloc_register(); + program.emit_int(0, zero_reg); + + let ok_pc = program.allocate_label(); + program.emit_insn(Insn::Ne { + lhs: right_reg, + rhs: zero_reg, + target_pc: ok_pc, + flags: CmpInsFlags::default().jump_if_null(), + collation: None, + }); + + program.emit_insn(Insn::Halt { + err_code: SQLITE_CONSTRAINT, + description: "divide by zero".into(), + }); + + program.resolve_label(ok_pc, program.offset()); + program.emit_insn(Insn::Divide { + lhs: left_reg, + rhs: right_reg, + dest: r, + }); + } + _ => { + program.emit_insn(Insn::Halt { + err_code: SQLITE_CONSTRAINT, + description: "unsupported operator in offset expr".into(), + }); + } + } + } + _ => { + program.emit_insn(Insn::Halt { + err_code: SQLITE_CONSTRAINT, + description: "non-integer expression in offset".into(), + }); + } + } +} + +pub fn try_fold_expr_to_i64(expr: &Expr) -> Option { + match expr { + Expr::Literal(Literal::Numeric(n)) => n.parse::().ok(), + Expr::Unary(UnaryOperator::Negative, inner) => try_fold_expr_to_i64(inner).map(|v| -v), + Expr::Unary(UnaryOperator::Positive, inner) => try_fold_expr_to_i64(inner), + Expr::Binary(left, op, right) => { + let l = try_fold_expr_to_i64(left)?; + let r = try_fold_expr_to_i64(right)?; + match op { + Operator::Add => Some(l.saturating_add(r)), + Operator::Subtract => Some(l.saturating_sub(r)), + Operator::Multiply => Some(l.saturating_mul(r)), + Operator::Divide if r != 0 => Some(l.saturating_div(r)), + _ => None, + } + } + + _ => None, } }