diff --git a/core/schema.rs b/core/schema.rs index a30d59dcf..eceabbcfc 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -748,6 +748,22 @@ impl Affinity { ))), } } + + pub fn to_char_code(&self) -> u8 { + self.aff_mask() as u8 + } + + pub fn from_char_code(code: u8) -> Result { + Self::from_char(code as char) + } + + pub fn is_numeric(&self) -> bool { + matches!(self, Affinity::Integer | Affinity::Real | Affinity::Numeric) + } + + pub fn has_affinity(&self) -> bool { + !matches!(self, Affinity::Blob) + } } impl fmt::Display for Type { diff --git a/core/translate/expr.rs b/core/translate/expr.rs index 3371dbaf9..6b983e49a 100644 --- a/core/translate/expr.rs +++ b/core/translate/expr.rs @@ -1,4 +1,4 @@ -use limbo_sqlite3_parser::ast::{self, UnaryOperator}; +use limbo_sqlite3_parser::ast::{self, Expr, UnaryOperator}; use tracing::{instrument, Level}; use super::emitter::Resolver; @@ -8,7 +8,7 @@ use super::plan::TableReferences; use crate::function::JsonFunc; use crate::function::{Func, FuncCtx, MathFuncArity, ScalarFunc, VectorFunc}; use crate::functions::datetime; -use crate::schema::{Table, Type}; +use crate::schema::{Affinity, Table, Type}; use crate::util::{exprs_are_equivalent, normalize_ident, parse_numeric_literal}; use crate::vdbe::builder::CursorKey; use crate::vdbe::{ @@ -458,11 +458,30 @@ pub fn translate_expr( } ast::Expr::Binary(e1, op, e2) => { // Check if both sides of the expression are equivalent and reuse the same register if so + println!("expr: {:?}", &expr); + println!("e1: {:?}, e2: {:?}, op: {:?}", e1, e2, op); + println!( + "expr affinity: {:?}", + get_expr_affinity(e1, referenced_tables) + ); + println!( + "expr affinity: {:?}", + get_expr_affinity(e2, referenced_tables) + ); if exprs_are_equivalent(e1, e2) { let shared_reg = program.alloc_register(); translate_expr(program, referenced_tables, e1, shared_reg, resolver)?; - emit_binary_insn(program, op, shared_reg, shared_reg, target_register)?; + emit_binary_insn( + program, + op, + shared_reg, + shared_reg, + target_register, + e1, + e2, + referenced_tables, + )?; program.reset_collation(); Ok(target_register) } else { @@ -509,7 +528,16 @@ pub fn translate_expr( }; program.set_collation(collation_ctx); - emit_binary_insn(program, op, e1_reg, e2_reg, target_register)?; + emit_binary_insn( + program, + op, + e1_reg, + e2_reg, + target_register, + e1, + e2, + referenced_tables, + )?; program.reset_collation(); Ok(target_register) } @@ -2201,7 +2229,25 @@ fn emit_binary_insn( lhs: usize, rhs: usize, target_register: usize, + lhs_expr: &Expr, + rhs_expr: &Expr, + referenced_tables: Option<&TableReferences>, ) -> Result<()> { + let mut affinity = Affinity::Blob; + if matches!( + op, + ast::Operator::Equals + | ast::Operator::NotEquals + | ast::Operator::Less + | ast::Operator::LessEquals + | ast::Operator::Greater + | ast::Operator::GreaterEquals + | ast::Operator::Is + | ast::Operator::IsNot + ) { + affinity = comparison_affinity(lhs_expr, rhs_expr, referenced_tables); + } + println!("emit_binary affinity: {:?}", affinity); match op { ast::Operator::NotEquals => { let if_true_label = program.allocate_label(); @@ -2211,7 +2257,7 @@ fn emit_binary_insn( lhs, rhs, target_pc: if_true_label, - flags: CmpInsFlags::default(), + flags: CmpInsFlags::default().with_affinity(affinity), collation: program.curr_collation(), }, target_register, @@ -2228,7 +2274,7 @@ fn emit_binary_insn( lhs, rhs, target_pc: if_true_label, - flags: CmpInsFlags::default(), + flags: CmpInsFlags::default().with_affinity(affinity), collation: program.curr_collation(), }, target_register, @@ -2245,7 +2291,7 @@ fn emit_binary_insn( lhs, rhs, target_pc: if_true_label, - flags: CmpInsFlags::default(), + flags: CmpInsFlags::default().with_affinity(affinity), collation: program.curr_collation(), }, target_register, @@ -2262,7 +2308,7 @@ fn emit_binary_insn( lhs, rhs, target_pc: if_true_label, - flags: CmpInsFlags::default(), + flags: CmpInsFlags::default().with_affinity(affinity), collation: program.curr_collation(), }, target_register, @@ -2279,7 +2325,7 @@ fn emit_binary_insn( lhs, rhs, target_pc: if_true_label, - flags: CmpInsFlags::default(), + flags: CmpInsFlags::default().with_affinity(affinity), collation: program.curr_collation(), }, target_register, @@ -2296,7 +2342,7 @@ fn emit_binary_insn( lhs, rhs, target_pc: if_true_label, - flags: CmpInsFlags::default(), + flags: CmpInsFlags::default().with_affinity(affinity), collation: program.curr_collation(), }, target_register, @@ -3023,3 +3069,75 @@ where Ok(()) } + +pub fn get_expr_affinity( + expr: &ast::Expr, + referenced_tables: Option<&TableReferences>, +) -> Affinity { + match expr { + ast::Expr::Column { table, column, .. } => { + if let Some(tables) = referenced_tables { + if let Some(table_ref) = tables.find_table_by_internal_id(*table) { + if let Some(col) = table_ref.get_column_at(*column) { + return col.affinity(); + } + } + } + Affinity::Blob + } + ast::Expr::Cast { type_name, .. } => { + if let Some(type_name) = type_name { + crate::schema::affinity(&type_name.name) + } else { + Affinity::Blob + } + } + ast::Expr::Collate(expr, _) => get_expr_affinity(expr, referenced_tables), + // Literals have NO affinity in SQLite! + ast::Expr::Literal(_) => Affinity::Blob, // No affinity! + _ => Affinity::Blob, // This may need to change. For now this works. + } +} + +pub fn comparison_affinity( + lhs_expr: &ast::Expr, + rhs_expr: &ast::Expr, + referenced_tables: Option<&TableReferences>, +) -> Affinity { + let mut aff = get_expr_affinity(lhs_expr, referenced_tables); + + aff = compare_affinity(rhs_expr, aff, referenced_tables); + + // If no affinity determined (both operands are literals), default to BLOB + if !aff.has_affinity() { + Affinity::Blob + } else { + aff + } +} + +pub fn compare_affinity( + expr: &ast::Expr, + other_affinity: Affinity, + referenced_tables: Option<&TableReferences>, +) -> Affinity { + let expr_affinity = get_expr_affinity(expr, referenced_tables); + + if expr_affinity.has_affinity() && other_affinity.has_affinity() { + // Both sides have affinity - use numeric if either is numeric + if expr_affinity.is_numeric() || other_affinity.is_numeric() { + Affinity::Numeric + } else { + Affinity::Blob + } + } else { + // One or both sides have no affinity - use the one that does, or Blob if neither + if expr_affinity.has_affinity() { + expr_affinity + } else if other_affinity.has_affinity() { + other_affinity + } else { + Affinity::Blob + } + } +} diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index 23b9da480..c6a13840c 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -6,7 +6,7 @@ use std::{ use super::{execute, AggFunc, BranchOffset, CursorID, FuncCtx, InsnFunction, PageIdx}; use crate::{ - schema::{BTreeTable, Index}, + schema::{Affinity, BTreeTable, Index}, storage::{pager::CreateBTreeFlags, wal::CheckpointMode}, translate::collate::CollationSeq, }; @@ -20,6 +20,7 @@ pub struct CmpInsFlags(usize); impl CmpInsFlags { const NULL_EQ: usize = 0x80; const JUMP_IF_NULL: usize = 0x10; + const AFFINITY_MASK: usize = 0x0F; fn has(&self, flag: usize) -> bool { (self.0 & flag) != 0 @@ -42,6 +43,17 @@ impl CmpInsFlags { pub fn has_nulleq(&self) -> bool { self.has(CmpInsFlags::NULL_EQ) } + + pub fn with_affinity(mut self, affinity: Affinity) -> Self { + let aff_code = affinity.to_char_code() as usize; + self.0 = (self.0 & !Self::AFFINITY_MASK) | aff_code; + self + } + + pub fn get_affinity(&self) -> Affinity { + let aff_code = (self.0 & Self::AFFINITY_MASK) as u8; + Affinity::from_char_code(aff_code).unwrap_or(Affinity::Blob) + } } #[derive(Clone, Copy, Debug, Default)]