From 1b0aa568a4ff13a9e6a00b210418dfc9b9a0cfa8 Mon Sep 17 00:00:00 2001
From: jussisaurio <jussi.saurio@gmail.com>
Date: Fri, 12 Jul 2024 00:17:47 +0300
Subject: [PATCH] Basic where clause support

---
 core/translate.rs | 102 ++++++++++++++-
 core/vdbe.rs      | 323 +++++++++++++++++++++++++++++++++++++++++++++-
 testing/all.test  |  45 +++++++
 3 files changed, 466 insertions(+), 4 deletions(-)

diff --git a/core/translate.rs b/core/translate.rs
index 09c96feb3..99a3c0a50 100644
--- a/core/translate.rs
+++ b/core/translate.rs
@@ -16,6 +16,7 @@ struct Select {
     src_tables: Vec<SrcTable>, // Tables we use to get data from. This includes "from" and "joins"
     limit: Option<ast::Limit>,
     exist_aggregation: bool,
+    where_clause: Option<ast::Expr>,
     /// Ordered list of opened read table loops
     /// Used for generating a loop that looks like this:
     /// cursor 0 = open table 0
@@ -81,6 +82,7 @@ fn build_select(schema: &Schema, select: ast::Select) -> Result<Select> {
         ast::OneSelect::Select {
             columns,
             from: Some(from),
+            where_clause,
             ..
         } => {
             let table_name = match from.select {
@@ -130,12 +132,14 @@ fn build_select(schema: &Schema, select: ast::Select) -> Result<Select> {
                 src_tables: joins,
                 limit: select.limit.clone(),
                 exist_aggregation,
+                where_clause,
                 loops: Vec::new(),
             })
         }
         ast::OneSelect::Select {
             columns,
             from: None,
+            where_clause,
             ..
         } => {
             let column_info = analyze_columns(&columns, &Vec::new());
@@ -145,6 +149,7 @@ fn build_select(schema: &Schema, select: ast::Select) -> Result<Select> {
                 column_info,
                 src_tables: Vec::new(),
                 limit: select.limit.clone(),
+                where_clause,
                 exist_aggregation,
                 loops: Vec::new(),
             })
@@ -184,6 +189,8 @@ fn translate_select(mut select: Select) -> Result<Program> {
     let limit_insn = if !select.src_tables.is_empty() {
         translate_tables_begin(&mut program, &mut select);
 
+        let where_maybe = insert_where_clause_instructions(&select, &mut program)?;
+
         let (register_start, register_end) = translate_columns(&mut program, &select)?;
 
         let mut limit_insn: Option<usize> = None;
@@ -195,7 +202,17 @@ fn translate_select(mut select: Select) -> Result<Program> {
             limit_insn = limit_reg.map(|_| program.emit_placeholder());
         }
 
-        translate_tables_end(&mut program, &select);
+        let next_offset = translate_tables_end(&mut program, &select);
+
+        if let Some((where_clause_offset, where_reg)) = where_maybe {
+            program.fixup_insn(
+                where_clause_offset,
+                Insn::IfNot {
+                    reg: where_reg,
+                    target_pc: next_offset, // jump to 'next row' instruction if where not matched
+                },
+            );
+        }
 
         if select.exist_aggregation {
             let mut target = register_start;
@@ -218,7 +235,17 @@ fn translate_select(mut select: Select) -> Result<Program> {
         limit_insn
     } else {
         assert!(!select.exist_aggregation);
+        let where_maybe = insert_where_clause_instructions(&select, &mut program)?;
         let (register_start, register_end) = translate_columns(&mut program, &select)?;
+        if let Some((where_clause_offset, where_reg)) = where_maybe {
+            program.fixup_insn(
+                where_clause_offset,
+                Insn::IfNot {
+                    reg: where_reg,
+                    target_pc: program.offset() + 1, // jump directly over the result row if no source tables and where not matched
+                },
+            );
+        }
         program.emit_insn(Insn::ResultRow {
             start_reg: register_start,
             count: register_end - register_start,
@@ -260,6 +287,19 @@ fn translate_select(mut select: Select) -> Result<Program> {
     Ok(program.build())
 }
 
+fn insert_where_clause_instructions(
+    select: &Select,
+    program: &mut ProgramBuilder,
+) -> Result<Option<(usize, usize)>> {
+    if let Some(w) = &select.where_clause {
+        let where_reg = program.alloc_register();
+        let _ = translate_expr(program, &select, w, where_reg)?;
+        Ok(Some((program.emit_placeholder(), where_reg))) // We emit a placeholder because we determine the jump target later (after we know where the 'cursor next' instruction is)
+    } else {
+        Ok(None)
+    }
+}
+
 fn translate_tables_begin(program: &mut ProgramBuilder, select: &mut Select) {
     for join in &select.src_tables {
         let table = &join.table;
@@ -272,11 +312,15 @@ fn translate_tables_begin(program: &mut ProgramBuilder, select: &mut Select) {
     }
 }
 
-fn translate_tables_end(program: &mut ProgramBuilder, select: &Select) {
+fn translate_tables_end(program: &mut ProgramBuilder, select: &Select) -> usize {
     // iterate in reverse order as we open cursors in order
+    let mut next_offset = std::usize::MAX;
     for table_loop in select.loops.iter().rev() {
         let cursor_id = table_loop.open_cursor;
         program.emit_insn(Insn::NextAsync { cursor_id });
+        if next_offset == std::usize::MAX {
+            next_offset = program.offset() - 1;
+        }
         program.emit_insn(Insn::NextAwait {
             cursor_id,
             pc_if_next: table_loop.rewind_offset,
@@ -289,6 +333,8 @@ fn translate_tables_end(program: &mut ProgramBuilder, select: &Select) {
             },
         );
     }
+
+    next_offset
 }
 
 fn translate_table_open_cursor(program: &mut ProgramBuilder, table: &Table) -> LoopInfo {
@@ -468,7 +514,57 @@ fn translate_expr(
 ) -> Result<usize> {
     match expr {
         ast::Expr::Between { .. } => todo!(),
-        ast::Expr::Binary(_, _, _) => todo!(),
+        ast::Expr::Binary(e1, op, e2) => {
+            let e1_reg = program.alloc_register();
+            let e2_reg = program.alloc_register();
+            let _ = translate_expr(program, select, e1, e1_reg)?;
+            let _ = translate_expr(program, select, e2, e2_reg)?;
+            program.emit_insn(match op {
+                ast::Operator::NotEquals => Insn::Ne {
+                    lhs: e1_reg,
+                    rhs: e2_reg,
+                    target_pc: program.offset() + 3, // jump to "emit True" instruction
+                },
+                ast::Operator::Equals => Insn::Eq {
+                    lhs: e1_reg,
+                    rhs: e2_reg,
+                    target_pc: program.offset() + 3,
+                },
+                ast::Operator::Less => Insn::Lt {
+                    lhs: e1_reg,
+                    rhs: e2_reg,
+                    target_pc: program.offset() + 3,
+                },
+                ast::Operator::LessEquals => Insn::Le {
+                    lhs: e1_reg,
+                    rhs: e2_reg,
+                    target_pc: program.offset() + 3,
+                },
+                ast::Operator::Greater => Insn::Gt {
+                    lhs: e1_reg,
+                    rhs: e2_reg,
+                    target_pc: program.offset() + 3,
+                },
+                ast::Operator::GreaterEquals => Insn::Ge {
+                    lhs: e1_reg,
+                    rhs: e2_reg,
+                    target_pc: program.offset() + 3,
+                },
+                _ => todo!(),
+            });
+            program.emit_insn(Insn::Integer {
+                value: 0, // emit False
+                dest: target_register,
+            });
+            program.emit_insn(Insn::Goto {
+                target_pc: program.offset() + 2,
+            });
+            program.emit_insn(Insn::Integer {
+                value: 1, // emit True
+                dest: target_register,
+            });
+            Ok(target_register)
+        }
         ast::Expr::Case { .. } => todo!(),
         ast::Expr::Cast { .. } => todo!(),
         ast::Expr::Collate(_, _) => todo!(),
diff --git a/core/vdbe.rs b/core/vdbe.rs
index 27fb65d95..f22754eb6 100644
--- a/core/vdbe.rs
+++ b/core/vdbe.rs
@@ -20,7 +20,47 @@ pub enum Insn {
     Init {
         target_pc: BranchOffset,
     },
-
+    // Compare two registers and jump to the given PC if they are equal.
+    Eq {
+        lhs: usize,
+        rhs: usize,
+        target_pc: BranchOffset,
+    },
+    // Compare two registers and jump to the given PC if they are not equal.
+    Ne {
+        lhs: usize,
+        rhs: usize,
+        target_pc: BranchOffset,
+    },
+    // Compare two registers and jump to the given PC if the left-hand side is less than the right-hand side.
+    Lt {
+        lhs: usize,
+        rhs: usize,
+        target_pc: BranchOffset,
+    },
+    // Compare two registers and jump to the given PC if the left-hand side is less than or equal to the right-hand side.
+    Le {
+        lhs: usize,
+        rhs: usize,
+        target_pc: BranchOffset,
+    },
+    // Compare two registers and jump to the given PC if the left-hand side is greater than the right-hand side.
+    Gt {
+        lhs: usize,
+        rhs: usize,
+        target_pc: BranchOffset,
+    },
+    // Compare two registers and jump to the given PC if the left-hand side is greater than or equal to the right-hand side.
+    Ge {
+        lhs: usize,
+        rhs: usize,
+        target_pc: BranchOffset,
+    },
+    // Jump to the given PC if the register is zero.
+    IfNot {
+        reg: usize,
+        target_pc: BranchOffset,
+    },
     // Open a cursor for reading.
     OpenReadAsync {
         cursor_id: CursorID,
@@ -294,6 +334,200 @@ impl Program {
                 Insn::Init { target_pc } => {
                     state.pc = *target_pc;
                 }
+                Insn::Eq {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    let lhs = *lhs;
+                    let rhs = *rhs;
+                    let target_pc = *target_pc;
+                    match (&state.registers[lhs], &state.registers[rhs]) {
+                        (OwnedValue::Integer(lhs), OwnedValue::Integer(rhs)) => {
+                            if lhs == rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Float(lhs), OwnedValue::Float(rhs)) => {
+                            if lhs == rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Text(lhs), OwnedValue::Text(rhs)) => {
+                            if lhs == rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        _ => {
+                            todo!();
+                        }
+                    }
+                }
+                Insn::Ne {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    let lhs = *lhs;
+                    let rhs = *rhs;
+                    let target_pc = *target_pc;
+                    match (&state.registers[lhs], &state.registers[rhs]) {
+                        (OwnedValue::Integer(lhs), OwnedValue::Integer(rhs)) => {
+                            if lhs != rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Float(lhs), OwnedValue::Float(rhs)) => {
+                            if lhs != rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Text(lhs), OwnedValue::Text(rhs)) => {
+                            if lhs != rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        _ => {
+                            todo!();
+                        }
+                    }
+                }
+                Insn::Lt {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    let lhs = *lhs;
+                    let rhs = *rhs;
+                    let target_pc = *target_pc;
+                    match (&state.registers[lhs], &state.registers[rhs]) {
+                        (OwnedValue::Integer(lhs), OwnedValue::Integer(rhs)) => {
+                            if lhs < rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Float(lhs), OwnedValue::Float(rhs)) => {
+                            if lhs < rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        _ => {
+                            todo!();
+                        }
+                    }
+                }
+                Insn::Le {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    let lhs = *lhs;
+                    let rhs = *rhs;
+                    let target_pc = *target_pc;
+                    match (&state.registers[lhs], &state.registers[rhs]) {
+                        (OwnedValue::Integer(lhs), OwnedValue::Integer(rhs)) => {
+                            if lhs <= rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Float(lhs), OwnedValue::Float(rhs)) => {
+                            if lhs <= rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        _ => {
+                            todo!();
+                        }
+                    }
+                }
+                Insn::Gt {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    let lhs = *lhs;
+                    let rhs = *rhs;
+                    let target_pc = *target_pc;
+                    match (&state.registers[lhs], &state.registers[rhs]) {
+                        (OwnedValue::Integer(lhs), OwnedValue::Integer(rhs)) => {
+                            if lhs > rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Float(lhs), OwnedValue::Float(rhs)) => {
+                            if lhs > rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        _ => {
+                            todo!();
+                        }
+                    }
+                }
+                Insn::Ge {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    let lhs = *lhs;
+                    let rhs = *rhs;
+                    let target_pc = *target_pc;
+                    match (&state.registers[lhs], &state.registers[rhs]) {
+                        (OwnedValue::Integer(lhs), OwnedValue::Integer(rhs)) => {
+                            if lhs >= rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        (OwnedValue::Float(lhs), OwnedValue::Float(rhs)) => {
+                            if lhs >= rhs {
+                                state.pc = target_pc;
+                            } else {
+                                state.pc += 1;
+                            }
+                        }
+                        _ => {
+                            todo!();
+                        }
+                    }
+                }
+                Insn::IfNot { reg, target_pc } => {
+                    let reg = *reg;
+                    let target_pc = *target_pc;
+                    match &state.registers[reg] {
+                        OwnedValue::Integer(0) => {
+                            state.pc = target_pc;
+                        }
+                        _ => {
+                            state.pc += 1;
+                        }
+                    }
+                }
                 Insn::OpenReadAsync {
                     cursor_id,
                     root_page,
@@ -712,6 +946,93 @@ fn insn_to_str(addr: BranchOffset, insn: &Insn, indent: String) -> String {
                 0,
                 format!("Start at {}", target_pc),
             ),
+            Insn::Eq {
+                lhs,
+                rhs,
+                target_pc,
+            } => (
+                "Eq",
+                *lhs as i32,
+                *rhs as i32,
+                *target_pc as i32,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] == r[{}] -> {}", lhs, rhs, target_pc),
+            ),
+            Insn::Ne {
+                lhs,
+                rhs,
+                target_pc,
+            } => (
+                "Ne",
+                *lhs as i32,
+                *rhs as i32,
+                *target_pc as i32,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] != r[{}] -> {}", lhs, rhs, target_pc),
+            ),
+            Insn::Lt {
+                lhs,
+                rhs,
+                target_pc,
+            } => (
+                "Lt",
+                *lhs as i32,
+                *rhs as i32,
+                *target_pc as i32,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] < r[{}] -> {}", lhs, rhs, target_pc),
+            ),
+            Insn::Le {
+                lhs,
+                rhs,
+                target_pc,
+            } => (
+                "Le",
+                *lhs as i32,
+                *rhs as i32,
+                *target_pc as i32,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] <= r[{}] -> {}", lhs, rhs, target_pc),
+            ),
+            Insn::Gt {
+                lhs,
+                rhs,
+                target_pc,
+            } => (
+                "Gt",
+                *lhs as i32,
+                *rhs as i32,
+                *target_pc as i32,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] > r[{}] -> {}", lhs, rhs, target_pc),
+            ),
+            Insn::Ge {
+                lhs,
+                rhs,
+                target_pc,
+            } => (
+                "Ge",
+                *lhs as i32,
+                *rhs as i32,
+                *target_pc as i32,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] >= r[{}] -> {}", lhs, rhs, target_pc),
+            ),
+            Insn::IfNot { reg, target_pc } => (
+                "IfNot",
+                *reg as i32,
+                *target_pc as i32,
+                0,
+                OwnedValue::Text(Rc::new("".to_string())),
+                0,
+                format!("r[{}] -> {}", reg, target_pc),
+            ),
             Insn::OpenReadAsync {
                 cursor_id,
                 root_page,
diff --git a/testing/all.test b/testing/all.test
index 91fe87ec0..0ab665e1b 100755
--- a/testing/all.test
+++ b/testing/all.test
@@ -81,3 +81,48 @@ do_execsql_test cross-join-specific-columns {
 do_execsql_test realify {
     select price from products limit 1;
 } {79.0}
+
+do_execsql_test where-clause-eq {
+    select last_name from users where id = 2000;
+} {Rodriguez}
+
+do_execsql_test where-clause-eq-string {
+    select count(1) from users where last_name = 'Rodriguez';
+} {61}
+
+do_execsql_test where-clause-ne {
+    select count(1) from users where id != 2000;
+} {9999}
+
+do_execsql_test where-clause-gt {
+    select count(1) from users where id > 2000;
+} {8000}
+
+do_execsql_test where-clause-gte {
+    select count(1) from users where id >= 2000;
+} {8001}
+
+do_execsql_test where-clause-lt {
+    select count(1) from users where id < 2000;
+} {1999}
+
+do_execsql_test where-clause-lte {
+    select count(1) from users where id <= 2000;
+} {2000}
+
+do_execsql_test where-clause-unary-true {
+    select count(1) from users where 1;
+} {10000}
+
+# not correct? should be 0?
+do_execsql_test where-clause-unary-false {
+    select count(1) from users where 0;
+} {NULL}
+
+do_execsql_test where-clause-no-table-unary-true {
+    select 1 where 1;
+} {1}
+
+do_execsql_test where-clause-no-table-unary-false {
+    select 1 where 0;
+} {}