From e08d23a00853d0eef2c649ab987a8a44b0a9dd7a Mon Sep 17 00:00:00 2001 From: Pekka Enberg Date: Sun, 10 Sep 2023 13:19:04 +0300 Subject: [PATCH] Implement LIMIT clause Note that we handle `LIMIT 0` a bit different from SQLite, which just detects the case and generates an unconditional jump doing nothing. Fixes #3 --- core/translate.rs | 114 ++++++++++++++++++++++++++++------------------ core/vdbe.rs | 22 +++++++++ 2 files changed, 92 insertions(+), 44 deletions(-) diff --git a/core/translate.rs b/core/translate.rs index c8b2b0b82..7e89f4254 100644 --- a/core/translate.rs +++ b/core/translate.rs @@ -2,7 +2,7 @@ use crate::schema::Schema; use crate::vdbe::{Insn, Program, ProgramBuilder}; use anyhow::Result; -use sqlite3_parser::ast::{OneSelect, Select, Stmt}; +use sqlite3_parser::ast::{Expr, OneSelect, Select, Stmt}; pub fn translate(schema: &Schema, stmt: Stmt) -> Result { match stmt { @@ -34,7 +34,13 @@ fn translate_select(schema: &Schema, select: Select) -> Result { let root_page = table.root_page; let mut program = ProgramBuilder::new(); let init_offset = program.emit_placeholder(); - let open_read_offset = program.offset(); + let start_offset = program.offset(); + let limit_reg = if let Some(limit) = select.limit { + assert!(limit.offset.is_none()); + Some(translate_expr(&mut program, &limit.expr)) + } else { + None + }; program.emit_insn(Insn::OpenReadAsync { cursor_id: 0, root_page, @@ -42,6 +48,10 @@ fn translate_select(schema: &Schema, select: Select) -> Result { program.emit_insn(Insn::OpenReadAwait); program.emit_insn(Insn::RewindAsync { cursor_id }); let rewind_await_offset = program.emit_placeholder(); + let limit_decr_insn = match limit_reg { + Some(_) => Some(program.emit_placeholder()), + None => None, + }; let (register_start, register_end) = translate_columns(&mut program, Some(cursor_id), Some(table), columns); program.emit_insn(Insn::ResultRow { @@ -60,6 +70,15 @@ fn translate_select(schema: &Schema, select: Select) -> Result { pc_if_empty: program.offset(), }, ); + if let Some(limit_decr_insn) = limit_decr_insn { + program.fixup_insn( + limit_decr_insn, + Insn::DecrJumpZero { + reg: limit_reg.unwrap(), + target_pc: program.offset(), + }, + ); + } program.emit_insn(Insn::Halt); program.fixup_insn( init_offset, @@ -69,7 +88,7 @@ fn translate_select(schema: &Schema, select: Select) -> Result { ); program.emit_insn(Insn::Transaction); program.emit_insn(Insn::Goto { - target_pc: open_read_offset, + target_pc: start_offset, }); Ok(program.build()) } @@ -112,47 +131,9 @@ fn translate_columns( let register_start = program.next_free_register(); for col in columns { match col { - sqlite3_parser::ast::ResultColumn::Expr(expr, _) => match expr { - sqlite3_parser::ast::Expr::Between { .. } => todo!(), - sqlite3_parser::ast::Expr::Binary(_, _, _) => todo!(), - sqlite3_parser::ast::Expr::Case { .. } => todo!(), - sqlite3_parser::ast::Expr::Cast { .. } => todo!(), - sqlite3_parser::ast::Expr::Collate(_, _) => todo!(), - sqlite3_parser::ast::Expr::DoublyQualified(_, _, _) => todo!(), - sqlite3_parser::ast::Expr::Exists(_) => todo!(), - sqlite3_parser::ast::Expr::FunctionCall { .. } => todo!(), - sqlite3_parser::ast::Expr::FunctionCallStar { .. } => todo!(), - sqlite3_parser::ast::Expr::Id(_) => todo!(), - sqlite3_parser::ast::Expr::InList { .. } => todo!(), - sqlite3_parser::ast::Expr::InSelect { .. } => todo!(), - sqlite3_parser::ast::Expr::InTable { .. } => todo!(), - sqlite3_parser::ast::Expr::IsNull(_) => todo!(), - sqlite3_parser::ast::Expr::Like { .. } => todo!(), - sqlite3_parser::ast::Expr::Literal(lit) => match lit { - sqlite3_parser::ast::Literal::Numeric(val) => { - let dest = program.alloc_register(); - program.emit_insn(Insn::Integer { - value: val.parse().unwrap(), - dest, - }); - } - sqlite3_parser::ast::Literal::String(_) => todo!(), - sqlite3_parser::ast::Literal::Blob(_) => todo!(), - sqlite3_parser::ast::Literal::Keyword(_) => todo!(), - sqlite3_parser::ast::Literal::Null => todo!(), - sqlite3_parser::ast::Literal::CurrentDate => todo!(), - sqlite3_parser::ast::Literal::CurrentTime => todo!(), - sqlite3_parser::ast::Literal::CurrentTimestamp => todo!(), - }, - sqlite3_parser::ast::Expr::Name(_) => todo!(), - sqlite3_parser::ast::Expr::NotNull(_) => todo!(), - sqlite3_parser::ast::Expr::Parenthesized(_) => todo!(), - sqlite3_parser::ast::Expr::Qualified(_, _) => todo!(), - sqlite3_parser::ast::Expr::Raise(_, _) => todo!(), - sqlite3_parser::ast::Expr::Subquery(_) => todo!(), - sqlite3_parser::ast::Expr::Unary(_, _) => todo!(), - sqlite3_parser::ast::Expr::Variable(_) => todo!(), - }, + sqlite3_parser::ast::ResultColumn::Expr(expr, _) => { + let _ = translate_expr(program, &expr); + } sqlite3_parser::ast::ResultColumn::Star => { for (i, col) in table.unwrap().columns.iter().enumerate() { let dest = program.alloc_register(); @@ -176,3 +157,48 @@ fn translate_columns( let register_end = program.next_free_register(); (register_start, register_end) } + +fn translate_expr(program: &mut ProgramBuilder, expr: &Expr) -> usize { + match expr { + sqlite3_parser::ast::Expr::Between { .. } => todo!(), + sqlite3_parser::ast::Expr::Binary(_, _, _) => todo!(), + sqlite3_parser::ast::Expr::Case { .. } => todo!(), + sqlite3_parser::ast::Expr::Cast { .. } => todo!(), + sqlite3_parser::ast::Expr::Collate(_, _) => todo!(), + sqlite3_parser::ast::Expr::DoublyQualified(_, _, _) => todo!(), + sqlite3_parser::ast::Expr::Exists(_) => todo!(), + sqlite3_parser::ast::Expr::FunctionCall { .. } => todo!(), + sqlite3_parser::ast::Expr::FunctionCallStar { .. } => todo!(), + sqlite3_parser::ast::Expr::Id(_) => todo!(), + sqlite3_parser::ast::Expr::InList { .. } => todo!(), + sqlite3_parser::ast::Expr::InSelect { .. } => todo!(), + sqlite3_parser::ast::Expr::InTable { .. } => todo!(), + sqlite3_parser::ast::Expr::IsNull(_) => todo!(), + sqlite3_parser::ast::Expr::Like { .. } => todo!(), + sqlite3_parser::ast::Expr::Literal(lit) => match lit { + sqlite3_parser::ast::Literal::Numeric(val) => { + let dest = program.alloc_register(); + program.emit_insn(Insn::Integer { + value: val.parse().unwrap(), + dest, + }); + dest + } + sqlite3_parser::ast::Literal::String(_) => todo!(), + sqlite3_parser::ast::Literal::Blob(_) => todo!(), + sqlite3_parser::ast::Literal::Keyword(_) => todo!(), + sqlite3_parser::ast::Literal::Null => todo!(), + sqlite3_parser::ast::Literal::CurrentDate => todo!(), + sqlite3_parser::ast::Literal::CurrentTime => todo!(), + sqlite3_parser::ast::Literal::CurrentTimestamp => todo!(), + }, + sqlite3_parser::ast::Expr::Name(_) => todo!(), + sqlite3_parser::ast::Expr::NotNull(_) => todo!(), + sqlite3_parser::ast::Expr::Parenthesized(_) => todo!(), + sqlite3_parser::ast::Expr::Qualified(_, _) => todo!(), + sqlite3_parser::ast::Expr::Raise(_, _) => todo!(), + sqlite3_parser::ast::Expr::Subquery(_) => todo!(), + sqlite3_parser::ast::Expr::Unary(_, _) => todo!(), + sqlite3_parser::ast::Expr::Variable(_) => todo!(), + } +} diff --git a/core/vdbe.rs b/core/vdbe.rs index 87c63a973..5b7b313d8 100644 --- a/core/vdbe.rs +++ b/core/vdbe.rs @@ -84,6 +84,14 @@ pub enum Insn { cursor_id: CursorID, dest: usize, }, + + // Decrement the given register and jump to the given PC if the result is zero. + // + // Unlike in SQLite, if register is already zero, we don't decrement, but take the jump. + DecrJumpZero { + reg: usize, + target_pc: BranchOffset, + }, } pub struct ProgramBuilder { @@ -282,6 +290,17 @@ impl Program { } state.pc += 1; } + Insn::DecrJumpZero { reg, target_pc } => match state.registers[*reg] { + Value::Integer(n) => { + if n > 0 { + state.registers[*reg] = Value::Integer(n - 1); + state.pc += 1; + } else { + state.pc = *target_pc; + } + } + _ => unreachable!("DecrJumpZero on non-integer register"), + }, } } } @@ -370,6 +389,9 @@ fn insn_to_str(addr: usize, insn: &Insn) -> String { ("Integer", *dest, *value as usize, 0, "", 0, "".to_string()) } Insn::RowId { cursor_id, dest } => ("RowId", *cursor_id, *dest, 0, "", 0, "".to_string()), + Insn::DecrJumpZero { reg, target_pc } => { + ("DecrJumpZero", *reg, *target_pc, 0, "", 0, "".to_string()) + } }; format!( "{:<4} {:<13} {:<4} {:<4} {:<4} {:<13} {:<2} {}",