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
This commit is contained in:
Pekka Enberg
2023-09-10 13:19:04 +03:00
parent a2202ed31e
commit e08d23a008
2 changed files with 92 additions and 44 deletions

View File

@@ -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<Program> {
match stmt {
@@ -34,7 +34,13 @@ fn translate_select(schema: &Schema, select: Select) -> Result<Program> {
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> {
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<Program> {
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> {
);
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!(),
}
}

View File

@@ -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} {}",