Added EQ

jitexpr/src/ast/serialize.rs

@@ -119,6 +119,7 @@ fn format_string(value: &str, formatter: &mut fmt::Formatter<'_>) -> fmt::Result
 fn function_name(function: Function) -> &'static str {
     match function {
         Function::Add => "ADD",
+        Function::Eq => "EQ",
         Function::RegexpExtract => "REGEXP_EXTRACT",
     }
 }
@@ -126,6 +127,7 @@ fn function_name(function: Function) -> &'static str {
 fn parse_function(name: &str, offset: usize) -> Result<Function, DeserializeError> {
     match name {
         "ADD" => Ok(Function::Add),
+        "EQ" => Ok(Function::Eq),
         "REGEXP_EXTRACT" => Ok(Function::RegexpExtract),
         _ if !is_function_name(name) => Err(DeserializeError::new(
             offset,
@@ -409,6 +411,15 @@ mod tests {
         assert_eq!(format!("{expr:?}"), "(ADD 1i64 my_col)");
     }
 
+    #[test]
+    fn test_eq_round_trip() {
+        let expr = Function::Eq
+            .call_untyped_expr(vec![UntypedExpr::literal(1u64), UntypedExpr::literal(1i64)]);
+
+        assert_eq!(serialize(&expr), "(EQ 1u64 1i64)");
+        assert_eq!(deserialize("(EQ 1u64 1i64)").unwrap(), expr);
+    }
+
     #[test]
     fn test_serialize_literals() {
         let cases = [

jitexpr/src/functions/eq.rs

@@ -0,0 +1,405 @@
+// EQ compares two values of any type.
+//
+// Values of the same type compare directly. Numerical values also compare across
+// i64, u64, and f64. Values of unrelated types are not equal.
+//
+// In other words:
+// 1u64 == 1i64 ==> true
+// 1f64 == 1u64 ==> true
+// 1.2f64 == 1u64 ==> false
+// 1f64 == "1" ==> false
+// "1" == "1" ==> true
+
+use std::collections::HashMap;
+
+use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, Value, types};
+use cranelift::frontend::FunctionBuilder;
+use cranelift::prelude::{AbiParam, FloatCC, InstBuilder, IntCC};
+use cranelift_jit::{JITBuilder, JITModule};
+use cranelift_module::{Linkage, Module};
+
+use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr};
+use crate::compile::{CompileError, CompileFnBuilder, LoweringContext, TypedExpr, TypedExprAst};
+use crate::functions::{FnCall, FnCallEnum};
+use crate::types::{StringRef, VarType};
+
+const STRING_EQ_SYMBOL: &str = "jitexpr_string_eq";
+
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) struct EqFnCall {
+    pub(crate) args: Box<[TypedExpr]>,
+}
+
+impl FnCall for EqFnCall {
+    fn infer_types<'a>(
+        args: &'a [UntypedExpr],
+        target_type: InferredTypeSet,
+        inferred_types: &mut HashMap<&'a str, InferredTypeSet>,
+    ) -> Result<InferredTypeSet, TypeError> {
+        if target_type.intersect(InferredTypeSet::BOOLEAN).is_none() {
+            return Err(TypeError::WrongFunctionReturnType {
+                function: Function::Eq,
+                expected: target_type,
+                got: InferredTypeSet::BOOLEAN,
+            });
+        }
+        if args.len() != 2 {
+            return Err(TypeError::InvalidNumberOfArguments {
+                function: Function::Eq,
+                expected: 2,
+                got: args.len(),
+            });
+        }
+        // TODO actually we probably want to be stricter here, so that we pick the right column in
+        // the end. thing my_col == 1i64.
+        for arg in args {
+            crate::ast::infer_types_aux(arg, InferredTypeSet::ALL, inferred_types)?;
+        }
+        Ok(InferredTypeSet::BOOLEAN)
+    }
+
+    fn call_with_types(
+        args: &[UntypedExpr],
+        target_type_set: InferredTypeSet,
+        context: &mut CompileFnBuilder<'_, '_>,
+    ) -> Result<TypedExpr, CompileError> {
+        assert_eq!(args.len(), 2, "Expected 2 args for EQ");
+        debug_assert!(target_type_set.contains(VarType::Bool));
+
+        // EQ must retain each literal's declared type. In contrast with ADD, it
+        // does not need to choose one common arithmetic type for its operands.
+        let typed_args = args
+            .iter()
+            .map(|arg| match arg {
+                UntypedExpr::Literal(literal) => {
+                    context.apply_types(arg, InferredTypeSet::singleton(literal.r#type()))
+                }
+                _ => context.apply_types(arg, InferredTypeSet::ALL),
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        Ok(TypedExpr {
+            return_type: VarType::Bool,
+            ast: TypedExprAst::from_call(EqFnCall {
+                args: typed_args.into_boxed_slice(),
+            }),
+        })
+    }
+
+    fn args_mut(&mut self) -> &mut [TypedExpr] {
+        &mut self.args
+    }
+
+    fn emit_cranelift_ir(
+        &self,
+        return_type: VarType,
+        context: &mut LoweringContext<'_>,
+        builder: &mut FunctionBuilder<'_>,
+    ) -> Result<Value, CompileError> {
+        debug_assert_eq!(return_type, VarType::Bool);
+        let lhs_type = self.args[0].return_type;
+        let rhs_type = self.args[1].return_type;
+
+        if lhs_type == VarType::None || rhs_type == VarType::None {
+            let equal = lhs_type == VarType::None && rhs_type == VarType::None;
+            return Ok(builder.ins().iconst(types::I8, i64::from(equal)));
+        }
+
+        if lhs_type == VarType::Str && rhs_type == VarType::Str {
+            let lhs = context.compile_expr(&self.args[0], builder)?;
+            let rhs = context.compile_expr(&self.args[1], builder)?;
+            let call = builder
+                .ins()
+                .call(context.native_functions().string_eq(), &[lhs, rhs]);
+            return Ok(builder.inst_results(call)[0]);
+        }
+
+        if !types_are_comparable(lhs_type, rhs_type) {
+            return Ok(builder.ins().iconst(types::I8, 0));
+        }
+
+        let lhs = context.compile_expr(&self.args[0], builder)?;
+        let rhs = context.compile_expr(&self.args[1], builder)?;
+        let equal = match (lhs_type, rhs_type) {
+            (VarType::Bool, VarType::Bool)
+            | (VarType::I64, VarType::I64)
+            | (VarType::U64, VarType::U64) => builder.ins().icmp(IntCC::Equal, lhs, rhs),
+            (VarType::F64, VarType::F64) => builder.ins().fcmp(FloatCC::Equal, lhs, rhs),
+            (VarType::I64, VarType::U64) => emit_signed_unsigned_eq(lhs, rhs, builder),
+            (VarType::U64, VarType::I64) => emit_signed_unsigned_eq(rhs, lhs, builder),
+            (VarType::F64, VarType::I64) => emit_float_integer_eq(lhs, rhs, VarType::I64, builder),
+            (VarType::I64, VarType::F64) => emit_float_integer_eq(rhs, lhs, VarType::I64, builder),
+            (VarType::F64, VarType::U64) => emit_float_integer_eq(lhs, rhs, VarType::U64, builder),
+            (VarType::U64, VarType::F64) => emit_float_integer_eq(rhs, lhs, VarType::U64, builder),
+            _ => unreachable!("the operand types were checked above"),
+        };
+        Ok(equal)
+    }
+}
+
+fn types_are_comparable(lhs: VarType, rhs: VarType) -> bool {
+    lhs == rhs || (is_numerical(lhs) && is_numerical(rhs))
+}
+
+fn is_numerical(var_type: VarType) -> bool {
+    matches!(var_type, VarType::I64 | VarType::U64 | VarType::F64)
+}
+
+fn emit_signed_unsigned_eq(
+    signed: Value,
+    unsigned: Value,
+    builder: &mut FunctionBuilder<'_>,
+) -> Value {
+    let nonnegative = builder
+        .ins()
+        .icmp_imm_s(IntCC::SignedGreaterThanOrEqual, signed, 0);
+    let same_bits = builder.ins().icmp(IntCC::Equal, signed, unsigned);
+    builder.ins().band(nonnegative, same_bits)
+}
+
+fn emit_float_integer_eq(
+    float: Value,
+    integer: Value,
+    integer_type: VarType,
+    builder: &mut FunctionBuilder<'_>,
+) -> Value {
+    let (lower_bound, upper_bound) = match integer_type {
+        VarType::I64 => (i64::MIN as f64, -(i64::MIN as f64)),
+        VarType::U64 => (0.0, (u64::MAX as f64)),
+        _ => unreachable!("EQ only compares f64 to i64 or u64 here"),
+    };
+    let lower_bound = builder.ins().f64const(lower_bound);
+    let upper_bound = builder.ins().f64const(upper_bound);
+    let above_lower = builder
+        .ins()
+        .fcmp(FloatCC::GreaterThanOrEqual, float, lower_bound);
+    let below_upper = builder.ins().fcmp(FloatCC::LessThan, float, upper_bound);
+    let in_range = builder.ins().band(above_lower, below_upper);
+
+    let converted = match integer_type {
+        VarType::I64 => builder.ins().fcvt_to_sint_sat(types::I64, float),
+        VarType::U64 => builder.ins().fcvt_to_uint_sat(types::I64, float),
+        _ => unreachable!("EQ only compares f64 to i64 or u64 here"),
+    };
+    let same_integer = builder.ins().icmp(IntCC::Equal, converted, integer);
+    let round_trip = match integer_type {
+        VarType::I64 => builder.ins().fcvt_from_sint(types::F64, converted),
+        VarType::U64 => builder.ins().fcvt_from_uint(types::F64, converted),
+        _ => unreachable!("EQ only compares f64 to i64 or u64 here"),
+    };
+    let is_integral = builder.ins().fcmp(FloatCC::Equal, float, round_trip);
+    let equal = builder.ins().band(in_range, same_integer);
+    builder.ins().band(equal, is_integral)
+}
+
+pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) {
+    jit_builder.symbol(STRING_EQ_SYMBOL, string_eq as *const u8);
+}
+
+pub(super) fn declare_native_function(
+    module: &mut JITModule,
+    function: &mut CraneliftFunction,
+    pointer_type: Type,
+) -> Result<FuncRef, CompileError> {
+    let mut signature = module.make_signature();
+    signature
+        .params
+        .extend([AbiParam::new(pointer_type), AbiParam::new(pointer_type)]);
+    signature.returns.push(AbiParam::new(types::I8));
+    let function_id = module.declare_function(STRING_EQ_SYMBOL, Linkage::Import, &signature)?;
+    Ok(module.declare_func_in_func(function_id, function))
+}
+
+unsafe extern "C" fn string_eq(lhs: *const StringRef, rhs: *const StringRef) -> u8 {
+    match (unsafe { lhs.as_ref() }, unsafe { rhs.as_ref() }) {
+        (None, None) => 1,
+        (Some(lhs), Some(rhs)) => u8::from(unsafe { lhs.as_str() } == unsafe { rhs.as_str() }),
+        _ => 0,
+    }
+}
+
+impl From<EqFnCall> for FnCallEnum {
+    fn from(call: EqFnCall) -> Self {
+        FnCallEnum::Eq(call)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ast::{self, infer_types};
+    use crate::compile::compile;
+    use crate::types::VariableValue;
+
+    fn eval(expression: &str) -> bool {
+        let expression = ast::deserialize(expression).unwrap();
+        let compiled = compile(&expression, &HashMap::new()).unwrap();
+        let mut output = VariableValue { boolean: false };
+
+        unsafe { compiled.call(&[], &mut output) };
+
+        unsafe { output.boolean }
+    }
+
+    #[test]
+    fn test_infer_types_accepts_any_operand_types() {
+        let expression = ast::deserialize(r#"(EQ value "hello")"#).unwrap();
+
+        let inferred_types = infer_types(&expression).unwrap();
+
+        assert_eq!(inferred_types.get("value"), Some(&InferredTypeSet::ALL));
+    }
+
+    #[test]
+    fn test_infer_types_requires_two_arguments() {
+        let expression = Function::Eq.call_untyped_expr(vec![UntypedExpr::literal(1i64)]);
+
+        let error = infer_types(&expression).unwrap_err();
+
+        assert!(matches!(
+            error,
+            TypeError::InvalidNumberOfArguments {
+                function: Function::Eq,
+                expected: 2,
+                got: 1,
+            }
+        ));
+    }
+
+    #[test]
+    fn test_compile_numeric_examples() {
+        assert!(eval("(EQ 1u64 1i64)"));
+        assert!(eval("(EQ 1f64 1i64)"));
+        assert!(!eval("(EQ 1.2f64 1i64)"));
+    }
+
+    #[test]
+    fn test_compile_different_types_are_not_equal() {
+        assert!(!eval(r#"(EQ 1i64 "1")"#));
+        assert!(!eval("(EQ true 1u64)"));
+    }
+
+    #[test]
+    fn test_compile_signed_unsigned_comparison() {
+        let expression = ast::deserialize("(EQ signed unsigned)").unwrap();
+        let variable_types = HashMap::from([("signed", VarType::I64), ("unsigned", VarType::U64)]);
+        let compiled = compile(&expression, &variable_types).unwrap();
+
+        for (signed, unsigned, expected) in [(7i64, 7u64, true), (-1, u64::MAX, false)] {
+            let input = [
+                VariableValue { int_i64: signed },
+                VariableValue { int_u64: unsigned },
+            ];
+            let mut output = VariableValue { boolean: false };
+            unsafe { compiled.call(&input, &mut output) };
+            assert_eq!(unsafe { output.boolean }, expected);
+        }
+    }
+
+    #[test]
+    fn test_compile_float_integer_comparison_is_exact() {
+        let expression = ast::deserialize("(EQ float integer)").unwrap();
+        let variable_types = HashMap::from([("float", VarType::F64), ("integer", VarType::I64)]);
+        let compiled = compile(&expression, &variable_types).unwrap();
+        let cases = [
+            (1.0, 1, true),
+            (1.2, 1, false),
+            ((1u64 << 53) as f64, (1i64 << 53) + 1, false),
+            (i64::MIN as f64, i64::MIN, true),
+            (2f64.powi(63), i64::MAX, false),
+        ];
+
+        for (float, integer, expected) in cases {
+            let input = [VariableValue { float }, VariableValue { int_i64: integer }];
+            let mut output = VariableValue { boolean: false };
+            unsafe { compiled.call(&input, &mut output) };
+            assert_eq!(unsafe { output.boolean }, expected);
+        }
+    }
+
+    #[test]
+    fn test_compile_float_unsigned_comparison_is_exact() {
+        let expression = ast::deserialize("(EQ float integer)").unwrap();
+        let variable_types = HashMap::from([("float", VarType::F64), ("integer", VarType::U64)]);
+        let compiled = compile(&expression, &variable_types).unwrap();
+        let cases = [
+            (1.0, 1, true),
+            (1.2, 1, false),
+            (2f64.powi(63), 1u64 << 63, true),
+            (u64::MAX as f64, u64::MAX, false),
+            (f64::NAN, 0, false),
+        ];
+
+        for (float, integer, expected) in cases {
+            let input = [VariableValue { float }, VariableValue { int_u64: integer }];
+            let mut output = VariableValue { boolean: false };
+            unsafe { compiled.call(&input, &mut output) };
+            assert_eq!(unsafe { output.boolean }, expected);
+        }
+    }
+
+    #[test]
+    fn test_compile_string_equality_compares_contents() {
+        let expression = ast::deserialize("(EQ left right)").unwrap();
+        let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]);
+        let compiled = compile(&expression, &variable_types).unwrap();
+        let left_value = String::from("same contents");
+        let right_value = String::from("same contents");
+        let mut left = StringRef::new(&left_value);
+        let mut right = StringRef::new(&right_value);
+        let input = [
+            VariableValue { string: &mut left },
+            VariableValue { string: &mut right },
+        ];
+        let mut output = VariableValue { boolean: false };
+
+        unsafe { compiled.call(&input, &mut output) };
+
+        assert!(unsafe { output.boolean });
+
+        let different_value = String::from("different contents");
+        let mut different = StringRef::new(&different_value);
+        let input = [
+            VariableValue { string: &mut left },
+            VariableValue {
+                string: &mut different,
+            },
+        ];
+        unsafe { compiled.call(&input, &mut output) };
+        assert!(!unsafe { output.boolean });
+
+        let null_input = [
+            VariableValue {
+                string: std::ptr::null_mut(),
+            },
+            VariableValue {
+                string: std::ptr::null_mut(),
+            },
+        ];
+        unsafe { compiled.call(&null_input, &mut output) };
+        assert!(unsafe { output.boolean });
+    }
+
+    #[test]
+    fn test_compile_none_equality() {
+        assert!(eval("(EQ none none)"));
+        assert!(!eval("(EQ none false)"));
+    }
+
+    #[test]
+    fn test_call_with_types_preserves_literal_types() {
+        let typed_expr = crate::typed_expr_from_str("(EQ 1u64 1f64)", &HashMap::new());
+        let TypedExprAst::FnCall(FnCallEnum::Eq(call)) = typed_expr.ast else {
+            panic!("expected an EQ call");
+        };
+
+        assert_eq!(call.args[0].return_type, VarType::U64);
+        assert_eq!(call.args[1].return_type, VarType::F64);
+    }
+
+    #[test]
+    fn test_compile_boolean_equality() {
+        assert!(eval("(EQ true true)"));
+        assert!(!eval("(EQ true false)"));
+    }
+}

jitexpr/src/functions/mod.rs

@@ -1,4 +1,5 @@
 mod add;
+mod eq;
 mod native_function;
 mod regexp_extract;
 
@@ -7,6 +8,7 @@ use std::collections::HashMap;
 use cranelift::frontend::FunctionBuilder;
 
 pub(crate) use self::add::AddFnCall;
+pub(crate) use self::eq::EqFnCall;
 pub(crate) use self::native_function::{
     NativeFunctions, declare_native_functions, register_jit_symbols,
 };
@@ -20,6 +22,8 @@ use crate::types::VarType;
 pub enum Function {
     /// Adds zero or more numerical expressions.
     Add,
+    /// Compares two expressions for value equality.
+    Eq,
     /// Extracts a capture group from a string using a constant regular expression.
     RegexpExtract,
 }
@@ -33,6 +37,7 @@ impl Function {
     ) -> Result<TypedExpr, CompileError> {
         match self {
             Function::Add => <AddFnCall as FnCall>::call_with_types(args, target_type_set, context),
+            Function::Eq => <EqFnCall as FnCall>::call_with_types(args, target_type_set, context),
             Function::RegexpExtract => {
                 <RegexpExtractFnCall as FnCall>::call_with_types(args, target_type_set, context)
             }
@@ -47,6 +52,7 @@ impl Function {
     ) -> Result<InferredTypeSet, TypeError> {
         match self {
             Function::Add => <AddFnCall as FnCall>::infer_types(args, target_type, inferred_types),
+            Function::Eq => <EqFnCall as FnCall>::infer_types(args, target_type, inferred_types),
             Function::RegexpExtract => {
                 <RegexpExtractFnCall as FnCall>::infer_types(args, target_type, inferred_types)
             }
@@ -64,6 +70,7 @@ impl Function {
 #[derive(Clone, Debug, PartialEq)]
 pub(crate) enum FnCallEnum {
     Add(AddFnCall),
+    Eq(EqFnCall),
     RegexpExtract(RegexpExtractFnCall),
 }
 
@@ -71,6 +78,7 @@ impl FnCallEnum {
     pub(crate) fn args_mut(&mut self) -> &mut [TypedExpr] {
         match self {
             FnCallEnum::Add(call) => call.args_mut(),
+            FnCallEnum::Eq(call) => call.args_mut(),
             FnCallEnum::RegexpExtract(call) => call.args_mut(),
         }
     }
@@ -84,6 +92,7 @@ impl FnCallEnum {
     ) -> Result<cranelift::codegen::ir::Value, CompileError> {
         match self {
             FnCallEnum::Add(call) => call.emit_cranelift_ir(return_type, context, builder),
+            FnCallEnum::Eq(call) => call.emit_cranelift_ir(return_type, context, builder),
             FnCallEnum::RegexpExtract(call) => {
                 call.emit_cranelift_ir(return_type, context, builder)
             }

jitexpr/src/functions/native_function.rs

@@ -1,15 +1,20 @@
 use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type};
 use cranelift_jit::{JITBuilder, JITModule};
 
-use super::regexp_extract;
+use super::{eq, regexp_extract};
 use crate::compile::CompileError;
 
 /// References to native functions imported into the current Cranelift function.
 pub(crate) struct NativeFunctions {
+    string_eq: FuncRef,
     regexp_extract: FuncRef,
 }
 
 impl NativeFunctions {
+    pub(crate) fn string_eq(&self) -> FuncRef {
+        self.string_eq
+    }
+
     pub(crate) fn regexp_extract(&self) -> FuncRef {
         self.regexp_extract
     }
@@ -17,6 +22,7 @@ impl NativeFunctions {
 
 /// Registers the process symbols that native calls may reference from generated code.
 pub(crate) fn register_jit_symbols(jit_builder: &mut JITBuilder) {
+    eq::register_jit_symbol(jit_builder);
     regexp_extract::register_jit_symbol(jit_builder);
 }
 
@@ -27,6 +33,7 @@ pub(crate) fn declare_native_functions(
     pointer_type: Type,
 ) -> Result<NativeFunctions, CompileError> {
     Ok(NativeFunctions {
+        string_eq: eq::declare_native_function(module, function, pointer_type)?,
         regexp_extract: regexp_extract::declare_native_function(module, function, pointer_type)?,
     })
 }