Merge 'Fix INSERT UNION ALL' from Duy Dang

Close #3849
Close #3855

Reviewed-by: Jussi Saurio <jussi.saurio@gmail.com>

Closes #3877

core/translate/insert.rs

@@ -927,43 +927,48 @@ fn bind_insert(
                 .collect();
         }
         InsertBody::Select(select, upsert_opt) => {
-            match &mut select.body.select {
+            if select.body.compounds.is_empty() {
-                // TODO see how to avoid clone
+                match &mut select.body.select {
-                OneSelect::Values(values_expr) if values_expr.len() <= 1 => {
+                    // TODO see how to avoid clone
-                    if values_expr.is_empty() {
+                    OneSelect::Values(values_expr) if values_expr.len() <= 1 => {
-                        crate::bail_parse_error!("no values to insert");
+                        if values_expr.is_empty() {
-                    }
+                            crate::bail_parse_error!("no values to insert");
-                    for expr in values_expr.iter_mut().flat_map(|v| v.iter_mut()) {
-                        match expr.as_mut() {
-                            Expr::Id(name) => {
-                                if name.quoted_with('"') {
-                                    *expr = Expr::Literal(ast::Literal::String(name.as_literal()))
-                                        .into();
-                                } else {
-                                    // an INSERT INTO ... VALUES (...) cannot reference columns
-                                    crate::bail_parse_error!("no such column: {name}");
-                                }
-                            }
-                            Expr::Qualified(first_name, second_name) => {
-                                // an INSERT INTO ... VALUES (...) cannot reference columns
-                                crate::bail_parse_error!(
-                                    "no such column: {first_name}.{second_name}"
-                                );
-                            }
-                            _ => {}
                         }
-                        bind_and_rewrite_expr(
+                        for expr in values_expr.iter_mut().flat_map(|v| v.iter_mut()) {
-                            expr,
+                            match expr.as_mut() {
-                            None,
+                                Expr::Id(name) => {
-                            None,
+                                    if name.quoted_with('"') {
-                            connection,
+                                        *expr =
-                            &mut program.param_ctx,
+                                            Expr::Literal(ast::Literal::String(name.as_literal()))
-                            BindingBehavior::ResultColumnsNotAllowed,
+                                                .into();
-                        )?;
+                                    } else {
+                                        // an INSERT INTO ... VALUES (...) cannot reference columns
+                                        crate::bail_parse_error!("no such column: {name}");
+                                    }
+                                }
+                                Expr::Qualified(first_name, second_name) => {
+                                    // an INSERT INTO ... VALUES (...) cannot reference columns
+                                    crate::bail_parse_error!(
+                                        "no such column: {first_name}.{second_name}"
+                                    );
+                                }
+                                _ => {}
+                            }
+                            bind_and_rewrite_expr(
+                                expr,
+                                None,
+                                None,
+                                connection,
+                                &mut program.param_ctx,
+                                BindingBehavior::ResultColumnsNotAllowed,
+                            )?;
+                        }
+                        values = values_expr.pop().unwrap_or_else(Vec::new);
                     }
-                    values = values_expr.pop().unwrap_or_else(Vec::new);
+                    _ => inserting_multiple_rows = true,
                 }
-                _ => inserting_multiple_rows = true,
+            } else {
+                inserting_multiple_rows = true;
             }
             upsert = upsert_opt.take();
         }
@@ -1056,8 +1061,10 @@ fn init_source_emission<'a>(
 ) -> Result<ProgramBuilder> {
     let (num_values, cursor_id) = match body {
         InsertBody::Select(select, _) => {
-            // Simple Common case of INSERT INTO <table> VALUES (...)
+            // Simple common case of INSERT INTO <table> VALUES (...) without compounds.
-            if matches!(&select.body.select, OneSelect::Values(values) if values.len() <= 1) {
+            if select.body.compounds.is_empty()
+                && matches!(&select.body.select, OneSelect::Values(values) if values.len() <= 1)
+            {
                 (
                     values.len(),
                     program.alloc_cursor_id(CursorType::BTreeTable(ctx.table.clone())),

core/translate/values.rs

@@ -1,4 +1,4 @@
-use crate::translate::emitter::{Resolver, TranslateCtx};
+use crate::translate::emitter::TranslateCtx;
 use crate::translate::expr::{translate_expr_no_constant_opt, NoConstantOptReason};
 use crate::translate::plan::{QueryDestination, SelectPlan};
 use crate::translate::result_row::emit_offset;
@@ -10,7 +10,7 @@ use crate::Result;
 pub fn emit_values(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
-    t_ctx: &TranslateCtx,
+    t_ctx: &mut TranslateCtx,
 ) -> Result<usize> {
     if plan.values.len() == 1 {
         let start_reg = emit_values_when_single_row(program, plan, t_ctx)?;
@@ -20,7 +20,7 @@ pub fn emit_values(
     let reg_result_cols_start = match plan.query_destination {
         QueryDestination::ResultRows => emit_toplevel_values(program, plan, t_ctx)?,
         QueryDestination::CoroutineYield { yield_reg, .. } => {
-            emit_values_in_subquery(program, plan, &t_ctx.resolver, yield_reg)?
+            emit_values_in_subquery(program, plan, t_ctx, yield_reg)?
         }
         QueryDestination::EphemeralIndex { .. } => emit_toplevel_values(program, plan, t_ctx)?,
         QueryDestination::EphemeralTable { .. } => unreachable!(),
@@ -42,13 +42,21 @@ pub fn emit_values(
 fn emit_values_when_single_row(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
-    t_ctx: &TranslateCtx,
+    t_ctx: &mut TranslateCtx,
 ) -> Result<usize> {
     let end_label = program.allocate_label();
     emit_offset(program, end_label, t_ctx.reg_offset);
     let first_row = &plan.values[0];
     let row_len = first_row.len();
-    let start_reg = program.alloc_registers(row_len);
+    let start_reg = if let Some(reg) = t_ctx.reg_result_cols_start {
+        program.reg_result_cols_start = Some(reg);
+        reg
+    } else {
+        let reg = program.alloc_registers(row_len);
+        t_ctx.reg_result_cols_start = Some(reg);
+        program.reg_result_cols_start = Some(reg);
+        reg
+    };
     for (i, v) in first_row.iter().enumerate() {
         translate_expr_no_constant_opt(
             program,
@@ -67,7 +75,7 @@ fn emit_values_when_single_row(
 fn emit_toplevel_values(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
-    t_ctx: &TranslateCtx,
+    t_ctx: &mut TranslateCtx,
 ) -> Result<usize> {
     let yield_reg = program.alloc_register();
     let definition_label = program.allocate_label();
@@ -79,7 +87,7 @@ fn emit_toplevel_values(
     });
     program.preassign_label_to_next_insn(start_offset_label);
 
-    let start_reg = emit_values_in_subquery(program, plan, &t_ctx.resolver, yield_reg)?;
+    let start_reg = emit_values_in_subquery(program, plan, t_ctx, yield_reg)?;
 
     program.emit_insn(Insn::EndCoroutine { yield_reg });
     program.preassign_label_to_next_insn(definition_label);
@@ -123,11 +131,19 @@ fn emit_toplevel_values(
 fn emit_values_in_subquery(
     program: &mut ProgramBuilder,
     plan: &SelectPlan,
-    resolver: &Resolver,
+    t_ctx: &mut TranslateCtx,
     yield_reg: usize,
 ) -> Result<usize> {
     let row_len = plan.values[0].len();
-    let start_reg = program.alloc_registers(row_len);
+    let start_reg = if let Some(reg) = t_ctx.reg_result_cols_start {
+        program.reg_result_cols_start = Some(reg);
+        reg
+    } else {
+        let reg = program.alloc_registers(row_len);
+        t_ctx.reg_result_cols_start = Some(reg);
+        program.reg_result_cols_start = Some(reg);
+        reg
+    };
     for value in &plan.values {
         for (i, v) in value.iter().enumerate() {
             translate_expr_no_constant_opt(
@@ -135,7 +151,7 @@ fn emit_values_in_subquery(
                 None,
                 v,
                 start_reg + i,
-                resolver,
+                &t_ctx.resolver,
                 NoConstantOptReason::RegisterReuse,
             )?;
         }

testing/insert.test

@@ -314,6 +314,20 @@ do_execsql_test_on_specific_db {:memory:} insert_from_select_union_all_where {
 6
 8}
 
+do_execsql_test_on_specific_db {:memory:} insert_from_select_union_all_values {
+    CREATE TABLE t(a, b);
+    INSERT INTO t(a, b) SELECT * FROM (VALUES(3, 3)) UNION ALL VALUES(4, 4);
+    SELECT a, b FROM t ORDER BY rowid;
+} {3|3
+4|4}
+
+do_execsql_test_on_specific_db {:memory:} values-union-all {
+    CREATE TABLE test (a INTEGER NOT NULL, b INTEGER NOT NULL);
+    INSERT INTO test(a, b) VALUES(3, 3) UNION ALL VALUES(4, 4);
+    SELECT a, b FROM test ORDER BY rowid;
+} {3|3
+4|4}
+
 do_execsql_test_on_specific_db {:memory:} insert_from_select_same_table {
     CREATE TABLE t (a INTEGER PRIMARY KEY, b);