Added affinity inference and conversion for comparison ops.

Added affinity helper function for `CmpInsFlags`

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, LimboError> {
+        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 {

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
+        }
+    }
+}

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)]