Support OFFSET clause for LIMIT in UPDATE queries

core/translate/delete.rs

@@ -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
 }

core/translate/emitter.rs

@@ -442,25 +442,11 @@ 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(());
     }
-    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,
-    )?;
     if t_ctx.reg_limit.is_none() && plan.limit.is_some() {
         let reg = program.alloc_register();
         t_ctx.reg_limit = Some(reg);
@@ -469,15 +455,43 @@ fn emit_program_for_update(
             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 {
+        program.emit_insn(Insn::Goto {
+            target_pc: after_main_loop_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());
+    emit_update_insns(&plan, &mut t_ctx, program)?;
     close_loop(program, &mut t_ctx, &plan.table_references)?;
 
     program.resolve_label(after_main_loop_label, program.offset());
@@ -491,10 +505,11 @@ fn emit_program_for_update(
 
 fn emit_update_insns(
     plan: &UpdatePlan,
-    t_ctx: &TranslateCtx,
+    t_ctx: &mut TranslateCtx,
     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,24 +523,6 @@ fn emit_update_insns(
         },
         _ => 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,
@@ -537,6 +534,29 @@ fn emit_update_insns(
         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 meta = ConditionMetadata {
+            jump_if_condition_is_true: false,
+            jump_target_when_true: BranchOffset::Placeholder,
+            jump_target_when_false: loop_labels.next,
+        };
+        translate_condition_expr(
+            program,
+            &plan.table_references,
+            &cond.expr,
+            meta,
+            &t_ctx.resolver,
+        )?;
+    }
+    let first_col_reg = program.alloc_registers(table_ref.table.columns().len());
     // 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() {

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 {

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

core/translate/update.rs

@@ -157,11 +157,11 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
         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
-    };
+    let (limit, offset) = body
+        .limit
+        .as_ref()
+        .map(|l| parse_limit(*l.clone()))
+        .unwrap_or(Ok((None, None)))?;
     Ok(Plan::Update(UpdatePlan {
         table_references,
         set_clauses,
@@ -169,6 +169,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,
     }))
 }