From 1fcc2ddd907d2a6ca468e2f856e8f6ceb204d2d1 Mon Sep 17 00:00:00 2001
From: meteorgan <meteorite.gan@gmail.com>
Date: Mon, 16 Jun 2025 23:03:04 +0800
Subject: [PATCH] support limit

---
 core/translate/compound_select.rs |  8 ++++++++
 testing/select.test               | 34 +++++++++++++++++++------------
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/core/translate/compound_select.rs b/core/translate/compound_select.rs
index 244bd663e..0bb934e12 100644
--- a/core/translate/compound_select.rs
+++ b/core/translate/compound_select.rs
@@ -246,6 +246,7 @@ fn emit_compound_select(
                     &left_index,
                     right_cursor_id,
                     target_cursor_id,
+                    limit_ctx,
                 );
             }
             _ => {
@@ -362,6 +363,7 @@ fn read_intersect_rows(
     index: &Index,
     right_cursor_id: usize,
     target_cursor: Option<usize>,
+    limit_ctx: Option<LimitCtx>,
 ) {
     let label_close = program.allocate_label();
     let label_loop_start = program.allocate_label();
@@ -413,6 +415,12 @@ fn read_intersect_rows(
             count: column_count,
         });
     }
+    if let Some(limit_ctx) = limit_ctx {
+        program.emit_insn(Insn::DecrJumpZero {
+            reg: limit_ctx.reg_limit,
+            target_pc: label_close,
+        });
+    }
     program.preassign_label_to_next_insn(label_next);
     program.emit_insn(Insn::Next {
         cursor_id: left_cursor_id,
diff --git a/testing/select.test b/testing/select.test
index 1d1c5afb4..a6f9c4492 100755
--- a/testing/select.test
+++ b/testing/select.test
@@ -361,12 +361,6 @@ if {[info exists ::env(SQLITE_EXEC)] && ($::env(SQLITE_EXEC) eq "scripts/limbo-s
     y|y}
 }
 
-select * from t UNION select * from u UNION select * from v UNION ALL select * from t;
-} {x|x
-y|y
-x|x
-y|y}
-
 do_execsql_test_on_specific_db {:memory:} select-intersect-1 {
   CREATE TABLE t(x TEXT, y TEXT);
   CREATE TABLE u(x TEXT, y TEXT);
@@ -410,11 +404,25 @@ do_execsql_test_on_specific_db {:memory:} select-union-intersect {
   select * from t UNION select * from u INTERSECT select * from v;
 } {x|x}
 
-#do_execsql_test_on_specific_db {:memory:} select-intersect-with-limit {
-#  CREATE TABLE t(x TEXT, y TEXT);
-#  CREATE TABLE u(x TEXT, y TEXT);
-#  INSERT INTO t VALUES('x','x'),('y','y'), ('z','z');
-#  INSERT INTO u VALUES('x','x'),('y','y'), ('z','z');
+do_execsql_test_on_specific_db {:memory:} select-intersect-with-limit {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y'), ('z','z');
+  INSERT INTO u VALUES('x','x'),('y','y'), ('z','z');
 
-#  select * from t INTERSECT select * from u limit 2;
-#} {x|x}
\ No newline at end of file
+  select * from t INTERSECT select * from u limit 2;
+} {x|x
+y|y}
+
+do_execsql_test_on_specific_db {:memory:} select-intersect-union-with-limit {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  CREATE TABLE v(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y'), ('z','z');
+  INSERT INTO u VALUES('d','d'),('e','e'), ('z','z');
+  INSERT INTO v VALUES('a','a'),('b','b');
+
+  select * from t INTERSECT select * from u UNION select * from v limit 3;
+} {a|a
+b|b
+z|z}
\ No newline at end of file