evaluate limit or offset expr

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(

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 {

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<isize>,
-    offset: Option<isize>,
+    limit: Option<Expr>,
+    offset: Option<Expr>,
 ) {
-    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,
+                });
+            }
+        }
     }
 }

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<isize>, Option<isize>)> {
-    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::<isize>().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<Expr>, Option<Expr>)> {
+    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::<isize>().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<Expr>) {

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<usize>,
 ) {
-    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::<i64>().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<i64> {
+    match expr {
+        Expr::Literal(Literal::Numeric(n)) => n.parse::<i64>().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,
     }
 }