diff --git a/jitexpr/src/ast/serialize.rs b/jitexpr/src/ast/serialize.rs index 01fdfc30a..f4794c7fd 100644 --- a/jitexpr/src/ast/serialize.rs +++ b/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 { 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 = [ diff --git a/jitexpr/src/functions/eq.rs b/jitexpr/src/functions/eq.rs new file mode 100644 index 000000000..bd41c4a1c --- /dev/null +++ b/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 { + 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 { + 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::, _>>()?; + + 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 { + 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 { + 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 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)")); + } +} diff --git a/jitexpr/src/functions/mod.rs b/jitexpr/src/functions/mod.rs index 8dfb0ce55..50d9a0657 100644 --- a/jitexpr/src/functions/mod.rs +++ b/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 { match self { Function::Add => ::call_with_types(args, target_type_set, context), + Function::Eq => ::call_with_types(args, target_type_set, context), Function::RegexpExtract => { ::call_with_types(args, target_type_set, context) } @@ -47,6 +52,7 @@ impl Function { ) -> Result { match self { Function::Add => ::infer_types(args, target_type, inferred_types), + Function::Eq => ::infer_types(args, target_type, inferred_types), Function::RegexpExtract => { ::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 { 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) } diff --git a/jitexpr/src/functions/native_function.rs b/jitexpr/src/functions/native_function.rs index 42afc6829..a85fd0d50 100644 --- a/jitexpr/src/functions/native_function.rs +++ b/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 { Ok(NativeFunctions { + string_eq: eq::declare_native_function(module, function, pointer_type)?, regexp_extract: regexp_extract::declare_native_function(module, function, pointer_type)?, }) }