Repository: tursodatabase/limbo (rust)

Issue:
Support OFFSET modifier in UPDATE query

Pull request description:
Support offset clause in Update queries 
Closes #1245 

Patch:
diff --git a/core/translate/delete.rs b/core/translate/delete.rs
index 1e0d64a9..9652048f 100644
--- a/core/translate/delete.rs
+++ b/core/translate/delete.rs
@@ -68,7 +68,7 @@ pub fn prepare_delete_plan(
     )?;
 
     // Parse the LIMIT/OFFSET clause
-    let (resolved_limit, resolved_offset) = limit.map_or(Ok((None, None)), |l| parse_limit(*l))?;
+    let (resolved_limit, resolved_offset) = limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
 
     let plan = DeletePlan {
         table_references,
@@ -86,7 +86,5 @@ pub fn prepare_delete_plan(
 fn estimate_num_instructions(plan: &DeletePlan) -> usize {
     let base = 20;
 
-    let num_instructions = base + plan.table_references.len() * 10;
-
-    num_instructions
+    base + plan.table_references.len() * 10
 }
diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs
index e4d05bfa..fac46fee 100644
--- a/core/translate/emitter.rs
+++ b/core/translate/emitter.rs
@@ -442,11 +442,36 @@ fn emit_program_for_update(
 
     // Exit on LIMIT 0
     if let Some(0) = plan.limit {
-        epilogue(program, init_label, start_offset, TransactionMode::Read)?;
+        epilogue(program, init_label, start_offset, TransactionMode::None)?;
         program.result_columns = plan.returning.unwrap_or_default();
         program.table_references = plan.table_references;
         return Ok(());
     }
+    if t_ctx.reg_limit.is_none() && plan.limit.is_some() {
+        let reg = program.alloc_register();
+        t_ctx.reg_limit = Some(reg);
+        program.emit_insn(Insn::Integer {
+            value: plan.limit.unwrap() as i64,
+            dest: reg,
+        });
+        program.mark_last_insn_constant();
+        if t_ctx.reg_offset.is_none() && plan.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: plan.offset.unwrap() as i64,
+                dest: reg,
+            });
+            program.mark_last_insn_constant();
+            let combined_reg = program.alloc_register();
+            t_ctx.reg_limit_offset_sum = Some(combined_reg);
+            program.emit_insn(Insn::OffsetLimit {
+                limit_reg: t_ctx.reg_limit.unwrap(),
+                offset_reg: reg,
+                combined_reg,
+            });
+        }
+    }
     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 {
@@ -454,22 +479,12 @@ fn emit_program_for_update(
             target_pc: after_main_loop_label,
         });
     }
-    let skip_label = program.allocate_label();
     init_loop(
         program,
         &mut t_ctx,
         &plan.table_references,
         OperationMode::UPDATE,
     )?;
-    if t_ctx.reg_limit.is_none() && plan.limit.is_some() {
-        let reg = program.alloc_register();
-        t_ctx.reg_limit = Some(reg);
-        program.emit_insn(Insn::Integer {
-            value: plan.limit.unwrap() as i64,
-            dest: reg,
-        });
-        program.mark_last_insn_constant();
-    }
     open_loop(
         program,
         &mut t_ctx,
@@ -477,7 +492,6 @@ fn emit_program_for_update(
         &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());
@@ -495,6 +509,7 @@ fn emit_update_insns(
     program: &mut ProgramBuilder,
 ) -> crate::Result<()> {
     let table_ref = &plan.table_references.first().unwrap();
+    let loop_labels = t_ctx.labels_main_loop.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 {
@@ -508,13 +523,30 @@ fn emit_update_insns(
         },
         _ => return Ok(()),
     };
+    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(),
+    });
+
+    if let Some(offset) = t_ctx.reg_offset {
+        program.emit_insn(Insn::IfPos {
+            reg: offset,
+            target_pc: loop_labels.next,
+            decrement_by: 1,
+        });
+    }
 
     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(),
+            jump_target_when_true: BranchOffset::Placeholder,
+            jump_target_when_false: loop_labels.next,
         };
         translate_condition_expr(
             program,
@@ -523,20 +555,8 @@ fn emit_update_insns(
             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() {
diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs
index e7235e8f..fd1a8ac9 100644
--- a/core/translate/main_loop.rs
+++ b/core/translate/main_loop.rs
@@ -37,11 +37,11 @@ pub struct LeftJoinMetadata {
 #[derive(Debug, Clone, Copy)]
 pub struct LoopLabels {
     /// jump to the start of the loop body
-    loop_start: BranchOffset,
+    pub loop_start: BranchOffset,
     /// jump to the NextAsync instruction (or equivalent)
-    next: BranchOffset,
+    pub next: BranchOffset,
     /// jump to the end of the loop, exiting it
-    loop_end: BranchOffset,
+    pub loop_end: BranchOffset,
 }
 
 impl LoopLabels {
diff --git a/core/translate/plan.rs b/core/translate/plan.rs
index 3d93548d..af14f035 100644
--- a/core/translate/plan.rs
+++ b/core/translate/plan.rs
@@ -176,8 +176,8 @@ pub struct UpdatePlan {
     pub set_clauses: Vec<(usize, ast::Expr)>,
     pub where_clause: Vec<WhereTerm>,
     pub order_by: Option<Vec<(ast::Expr, Direction)>>,
-    // TODO: support OFFSET
     pub limit: Option<isize>,
+    pub offset: Option<isize>,
     // TODO: optional RETURNING clause
     pub returning: Option<Vec<ResultSetColumn>>,
     // whether the WHERE clause is always false
diff --git a/core/translate/planner.rs b/core/translate/planner.rs
index 953a15e5..1b78c954 100644
--- a/core/translate/planner.rs
+++ b/core/translate/planner.rs
@@ -850,30 +850,33 @@ fn parse_join<'a>(
     Ok(())
 }
 
-pub fn parse_limit(limit: Limit) -> Result<(Option<isize>, Option<isize>)> {
-    let offset_val = match limit.offset {
+pub fn parse_limit(limit: &Limit) -> Result<(Option<isize>, Option<isize>)> {
+    let offset_val = match &limit.offset {
         Some(offset_expr) => match offset_expr {
             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) => match *expr {
-                Expr::Literal(ast::Literal::Numeric(n)) => n.parse::<isize>().ok().map(|num| -num),
-                _ => crate::bail_parse_error!("Invalid OFFSET clause"),
-            },
+            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),
     };
 
-    if let Expr::Literal(ast::Literal::Numeric(n)) = limit.expr {
+    if let Expr::Literal(ast::Literal::Numeric(n)) = &limit.expr {
         Ok((n.parse().ok(), offset_val))
-    } else if let Expr::Unary(UnaryOperator::Negative, expr) = limit.expr {
-        if let Expr::Literal(ast::Literal::Numeric(n)) = *expr {
+    } else if let Expr::Unary(UnaryOperator::Negative, expr) = &limit.expr {
+        if let Expr::Literal(ast::Literal::Numeric(n)) = &**expr {
             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 {
+    } else if let Expr::Id(id) = &limit.expr {
         if id.0.eq_ignore_ascii_case("true") {
             Ok((Some(1), offset_val))
         } else if id.0.eq_ignore_ascii_case("false") {
diff --git a/core/translate/select.rs b/core/translate/select.rs
index fe7656f1..bde61880 100644
--- a/core/translate/select.rs
+++ b/core/translate/select.rs
@@ -370,7 +370,7 @@ pub fn prepare_select_plan<'a>(
 
             // Parse the LIMIT/OFFSET clause
             (plan.limit, plan.offset) =
-                select.limit.map_or(Ok((None, None)), |l| parse_limit(*l))?;
+                select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
 
             // Return the unoptimized query plan
             Ok(Plan::Select(plan))
diff --git a/core/translate/update.rs b/core/translate/update.rs
index a9388ee1..f282fff2 100644
--- a/core/translate/update.rs
+++ b/core/translate/update.rs
@@ -151,17 +151,21 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
             })
             .collect()
     });
+    // Parse the WHERE clause
     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
-    };
+
+    // Parse the LIMIT/OFFSET clause
+    let (limit, offset) = body
+        .limit
+        .as_ref()
+        .map(|l| parse_limit(l))
+        .unwrap_or(Ok((None, None)))?;
+
     Ok(Plan::Update(UpdatePlan {
         table_references,
         set_clauses,
@@ -169,6 +173,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
         returning: Some(result_columns),
         order_by,
         limit,
+        offset,
         contains_constant_false_condition: false,
     }))
 }


Test patch:
diff --git a/testing/cli_tests/cli_test_cases.py b/testing/cli_tests/cli_test_cases.py
index 2bce1f16..17035b0b 100755
--- a/testing/cli_tests/cli_test_cases.py
+++ b/testing/cli_tests/cli_test_cases.py
@@ -264,6 +264,32 @@ def test_update_with_limit():
     limbo.quit()
 
 
+
+def test_update_with_limit_and_offset():
+    limbo = TestLimboShell(
+        "CREATE TABLE t (a,b,c); insert into t values (1,2,3), (4,5,6), (7,8,9), (1,2,3),(4,5,6), (7,8,9);"
+    )
+    limbo.run_test("update-limit-offset", "UPDATE t SET a = 10 LIMIT 1 OFFSET 3;", "")
+    limbo.run_test(
+        "update-limit-offset-result", "SELECT COUNT(*) from t WHERE a = 10;", "1"
+    )
+    limbo.run_test("update-limit-result", "SELECT a from t LIMIT 4;", "1\n4\n7\n10")
+    limbo.run_test(
+        "update-limit-offset-zero", "UPDATE t SET a = 100 LIMIT 0 OFFSET 0;", ""
+    )
+    limbo.run_test(
+        "update-limit-zero-result", "SELECT COUNT(*) from t WHERE a = 100;", "0"
+    )
+    limbo.run_test("update-limit-all", "UPDATE t SET a = 100 LIMIT -1 OFFSET 1;", "")
+    limbo.run_test("update-limit-result", "SELECT COUNT(*) from t WHERE a = 100;", "5")
+    limbo.run_test(
+        "udpate-limit-where", "UPDATE t SET a = 333 WHERE b = 5 LIMIT 1 OFFSET 2;", ""
+    )
+    limbo.run_test(
+        "update-limit-where-result", "SELECT COUNT(*) from t WHERE a = 333;", "0"
+    )
+    limbo.quit()
+    
 def test_insert_default_values():
     limbo = TestLimboShell(
         "CREATE TABLE t (a integer default(42),b integer default (43),c integer default(44));"
@@ -292,4 +318,5 @@ if __name__ == "__main__":
     test_import_csv_skip()
     test_table_patterns()
     test_update_with_limit()
+    test_update_with_limit_and_offset()
     print("All tests have passed")