From aa69b279c39298b72eb49f10cc0a01dab3e8e37a Mon Sep 17 00:00:00 2001
From: meteorgan <meteorite.gan@gmail.com>
Date: Mon, 28 Jul 2025 00:58:20 +0800
Subject: [PATCH] support limit

---
 core/translate/compound_select.rs |  2 +-
 core/translate/emitter.rs         |  2 +-
 core/translate/values.rs          | 27 +++++++++++++++++++++------
 testing/select.test               | 18 ++++++++++++++++++
 4 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/core/translate/compound_select.rs b/core/translate/compound_select.rs
index 426d396ab..d17a28b79 100644
--- a/core/translate/compound_select.rs
+++ b/core/translate/compound_select.rs
@@ -311,7 +311,7 @@ fn create_dedupe_index(
     schema: &Schema,
 ) -> crate::Result<(usize, Arc<Index>)> {
     if !schema.indexes_enabled {
-        crate::bail_parse_error!("UNION OR INTERSECT is not supported without indexes");
+        crate::bail_parse_error!("UNION OR INTERSECT or EXCEPT is not supported without indexes");
     }
 
     let dedupe_index = Arc::new(Index {
diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs
index 3ab881813..32d4035a7 100644
--- a/core/translate/emitter.rs
+++ b/core/translate/emitter.rs
@@ -265,7 +265,7 @@ pub fn emit_query<'a>(
     t_ctx: &mut TranslateCtx<'a>,
 ) -> Result<usize> {
     if !plan.values.is_empty() {
-        let reg_result_cols_start = emit_values(program, plan, &t_ctx.resolver)?;
+        let reg_result_cols_start = emit_values(program, plan, &t_ctx.resolver, t_ctx.limit_ctx)?;
         return Ok(reg_result_cols_start);
     }
 
diff --git a/core/translate/values.rs b/core/translate/values.rs
index efefdf443..73a33d5eb 100644
--- a/core/translate/values.rs
+++ b/core/translate/values.rs
@@ -1,4 +1,4 @@
-use crate::translate::emitter::Resolver;
+use crate::translate::emitter::{LimitCtx, Resolver};
 use crate::translate::expr::{translate_expr_no_constant_opt, NoConstantOptReason};
 use crate::translate::plan::{QueryDestination, SelectPlan};
 use crate::vdbe::builder::ProgramBuilder;
@@ -10,18 +10,21 @@ pub fn emit_values(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
     resolver: &Resolver,
+    limit_ctx: Option<LimitCtx>,
 ) -> Result<usize> {
     if plan.values.len() == 1 {
-        let start_reg = emit_values_when_single_row(program, plan, resolver)?;
+        let start_reg = emit_values_when_single_row(program, plan, resolver, limit_ctx)?;
         return Ok(start_reg);
     }
 
     let reg_result_cols_start = match plan.query_destination {
-        QueryDestination::ResultRows => emit_toplevel_values(program, plan, resolver)?,
+        QueryDestination::ResultRows => emit_toplevel_values(program, plan, resolver, limit_ctx)?,
         QueryDestination::CoroutineYield { yield_reg, .. } => {
             emit_values_in_subquery(program, plan, resolver, yield_reg)?
         }
-        QueryDestination::EphemeralIndex { .. } => emit_toplevel_values(program, plan, resolver)?,
+        QueryDestination::EphemeralIndex { .. } => {
+            emit_toplevel_values(program, plan, resolver, limit_ctx)?
+        }
         QueryDestination::EphemeralTable { .. } => unreachable!(),
     };
     Ok(reg_result_cols_start)
@@ -31,6 +34,7 @@ fn emit_values_when_single_row(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
     resolver: &Resolver,
+    limit_ctx: Option<LimitCtx>,
 ) -> Result<usize> {
     let first_row = &plan.values[0];
     let row_len = first_row.len();
@@ -45,7 +49,9 @@ fn emit_values_when_single_row(
             NoConstantOptReason::RegisterReuse,
         )?;
     }
-    emit_values_to_destination(program, plan, start_reg, row_len);
+    let end_label = program.allocate_label();
+    emit_values_to_destination(program, plan, start_reg, row_len, limit_ctx, end_label);
+    program.preassign_label_to_next_insn(end_label);
     Ok(start_reg)
 }
 
@@ -53,6 +59,7 @@ fn emit_toplevel_values(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
     resolver: &Resolver,
+    limit_ctx: Option<LimitCtx>,
 ) -> Result<usize> {
     let yield_reg = program.alloc_register();
     let definition_label = program.allocate_label();
@@ -91,7 +98,7 @@ fn emit_toplevel_values(
         });
     }
 
-    emit_values_to_destination(program, plan, copy_start_reg, row_len);
+    emit_values_to_destination(program, plan, copy_start_reg, row_len, limit_ctx, end_label);
 
     program.emit_insn(Insn::Goto {
         target_pc: goto_label,
@@ -134,6 +141,8 @@ fn emit_values_to_destination(
     plan: &SelectPlan,
     start_reg: usize,
     row_len: usize,
+    limit_ctx: Option<LimitCtx>,
+    end_label: BranchOffset,
 ) {
     match &plan.query_destination {
         QueryDestination::ResultRows => {
@@ -141,6 +150,12 @@ fn emit_values_to_destination(
                 start_reg,
                 count: row_len,
             });
+            if let Some(limit_ctx) = limit_ctx {
+                program.emit_insn(Insn::DecrJumpZero {
+                    reg: limit_ctx.reg_limit,
+                    target_pc: end_label,
+                });
+            }
         }
         QueryDestination::CoroutineYield { yield_reg, .. } => {
             program.emit_insn(Insn::Yield {
diff --git a/testing/select.test b/testing/select.test
index 988b95f95..ec434b538 100755
--- a/testing/select.test
+++ b/testing/select.test
@@ -606,6 +606,24 @@ if {[info exists ::env(SQLITE_EXEC)] && ($::env(SQLITE_EXEC) eq "scripts/limbo-s
 
       select * from t EXCEPT values('x','x'),('z','y');
     } {y|y}
+
+    do_execsql_test_on_specific_db {:memory:} select-values-union-all-limit {
+    CREATE TABLE t (x TEXT);
+    INSERT INTO t VALUES('x'), ('y'), ('z');
+
+    values('x') UNION ALL select * from t limit 3;
+    } {x
+    x
+    y}
+
+    do_execsql_test_on_specific_db {:memory:} select-values-union-all-limit-2 {
+    CREATE TABLE t (x TEXT);
+    INSERT INTO t VALUES('x'), ('y'), ('z');
+
+    values('a'), ('b') UNION ALL select * from t limit 3;
+    } {a
+    b
+    x}
 }
 
 do_execsql_test_on_specific_db {:memory:} select-no-match-in-leaf-page {