diff --git a/Cargo.toml b/Cargo.toml index cb2e870756..0098e1ab59 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -164,6 +164,7 @@ members = [ "sstable", "tokenizer-api", "columnar", + "jitexpr", ] # Following the "fail" crate best practises, we isolate diff --git a/jitexpr/Cargo.toml b/jitexpr/Cargo.toml new file mode 100644 index 0000000000..3ed15ea37c --- /dev/null +++ b/jitexpr/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "jitexpr" +version = "0.1.0" +edition = "2024" + +[dependencies] +cranelift = "0.134.3" +cranelift-jit = "0.134.3" +cranelift-module = "0.134.3" +cranelift-native = "0.134.3" +regex = "1" +thiserror = "2.0.1" diff --git a/jitexpr/README.md b/jitexpr/README.md new file mode 100644 index 0000000000..6fbd154f9d --- /dev/null +++ b/jitexpr/README.md @@ -0,0 +1,12 @@ +This is an expression compiler relying on Cranelift. + + UntypedExpr + ↓ injecting variable types, and type checking + TypedExpr + ↓ lowering + Cranelift IR + ↓ Cranelift code generation + Machine code (or assembly) + +The project does not rely on cranelifts function call abstraction. +Instead it just manipulates expression, so everything is always inlined. diff --git a/jitexpr/examples/basic.rs b/jitexpr/examples/basic.rs new file mode 100644 index 0000000000..41923b4be1 --- /dev/null +++ b/jitexpr/examples/basic.rs @@ -0,0 +1,43 @@ +use std::collections::HashMap; +use std::error::Error; +use std::sync::Arc; + +use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, infer_types}; +use jitexpr::compile::{CompiledFn, CompiledFnCtx, compile}; +use jitexpr::types::{VarType, VariableValue}; + +fn main() -> Result<(), Box> { + // A simple expression that goes: + // my_col + 1 + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("my_col"), + UntypedExpr::literal(1.0f64), + ]); + + // Infer types does not return specific types, but instead a set of acceptable + // types for each variables. + let inferred_types: HashMap<&str, InferredTypeSet> = infer_types(&untyped_expr)?; + assert_eq!( + inferred_types.get("my_col").unwrap(), + &InferredTypeSet::NUMERICAL + ); + + // This is then up to us to decide the actual type for each variable. + // In tantivy, this means picking the first column with a type in inferred_types. + // + // If none match then we should use the VarType::None. + let variable_types: HashMap<&str, VarType> = + std::iter::once(("my_col", VarType::F64)).collect(); + + let compiled_fn: Arc = compile(&untyped_expr, &variable_types)?; + let mut compiled_fn_ctx = CompiledFnCtx::new(compiled_fn); + + // We use a nullable wrapper around the value union to pass typed variables. + // For present values, it is up to us to populate the correct union member. + // Not doing so is UB. + let input: Box<[VariableValue]> = vec![VariableValue::from(1.2f64)].into_boxed_slice(); + let output = unsafe { compiled_fn_ctx.call(&input[..]) }; + assert_eq!(unsafe { output.as_f64() }, Some(1.2f64 + 1.0f64)); + + Ok(()) +} diff --git a/jitexpr/src/ast/infer_types.rs b/jitexpr/src/ast/infer_types.rs new file mode 100644 index 0000000000..9d5ea56472 --- /dev/null +++ b/jitexpr/src/ast/infer_types.rs @@ -0,0 +1,296 @@ +use std::collections::HashMap; +use std::collections::hash_map::Entry; + +use crate::ast::{Function, Literal, UntypedExpr}; +use crate::functions::InvalidFunctionCall; +use crate::types::VarType; + +#[derive(Default, Copy, Clone, Debug, Eq, PartialEq)] +pub struct InferredTypeSet { + pub string: bool, + pub i64: bool, + pub u64: bool, + pub f64: bool, + pub boolean: bool, +} + +impl InferredTypeSet { + pub const NONE: InferredTypeSet = InferredTypeSet { + string: false, + i64: false, + u64: false, + f64: false, + boolean: false, + }; + + pub const ALL: InferredTypeSet = InferredTypeSet { + string: true, + i64: true, + u64: true, + f64: true, + boolean: true, + }; + + pub const NUMERICAL: InferredTypeSet = InferredTypeSet { + i64: true, + u64: true, + f64: true, + boolean: false, + string: false, + }; + + pub const I64: InferredTypeSet = InferredTypeSet { + i64: true, + ..Self::NONE + }; + + pub const U64: InferredTypeSet = InferredTypeSet { + u64: true, + ..Self::NONE + }; + + pub const F64: InferredTypeSet = InferredTypeSet { + f64: true, + ..Self::NONE + }; + + pub const STRING: InferredTypeSet = InferredTypeSet { + string: true, + ..Self::NONE + }; + + pub const BOOLEAN: InferredTypeSet = InferredTypeSet { + boolean: true, + ..Self::NONE + }; + + pub(crate) fn is_none(self) -> bool { + self == Self::NONE + } + + pub fn singleton(var_type: VarType) -> InferredTypeSet { + match var_type { + VarType::Bool => Self::BOOLEAN, + VarType::F64 => Self::F64, + VarType::U64 => Self::U64, + VarType::I64 => Self::I64, + VarType::Str => Self::STRING, + VarType::None => Self::NONE, + } + } + + pub(crate) fn intersect(self, target_inferred_type: InferredTypeSet) -> InferredTypeSet { + InferredTypeSet { + string: self.string && target_inferred_type.string, + i64: self.i64 && target_inferred_type.i64, + u64: self.u64 && target_inferred_type.u64, + f64: self.f64 && target_inferred_type.f64, + boolean: self.boolean && target_inferred_type.boolean, + } + } + + pub fn contains(&self, var_type: VarType) -> bool { + match var_type { + VarType::Bool => self.boolean, + VarType::F64 => self.f64, + VarType::U64 => self.u64, + VarType::I64 => self.i64, + VarType::Str => self.string, + VarType::None => self.is_none(), + } + } +} + +impl From for InferredTypeSet { + fn from(var_type: VarType) -> Self { + match var_type { + VarType::Bool => Self::BOOLEAN, + VarType::F64 => Self::F64, + VarType::U64 => Self::U64, + VarType::I64 => Self::I64, + VarType::Str => Self::STRING, + VarType::None => Self::NONE, + } + } +} + +impl std::fmt::Display for InferredTypeSet { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let mut types = Vec::new(); + if self.string { + types.push("string"); + } + if self.i64 { + types.push("i64"); + } + if self.u64 { + types.push("u64"); + } + if self.f64 { + types.push("f64"); + } + if self.boolean { + types.push("boolean"); + } + write!(f, "{{{}}}", types.join(", ")) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum TypeError { + #[error(transparent)] + InvalidFunctionCall(#[from] InvalidFunctionCall), + #[error("function `{function:?}` returns `{got}`, expected `{expected}`")] + WrongFunctionReturnType { + function: Function, + expected: InferredTypeSet, + got: InferredTypeSet, + }, + #[error("function `{function:?}` expects `{expected}` args, was passed `{got}`")] + InvalidNumberOfArguments { + function: Function, + expected: usize, + got: usize, + }, + #[error("expected `{expected}` , got `{literal:?}`")] + InvalidLiteralType { + literal: Literal, + expected: InferredTypeSet, + }, +} + +/// Infer the accepted types for the different variables present in the formula. +pub fn infer_types(expr: &UntypedExpr) -> Result, TypeError> { + infer_types_with_target(expr, InferredTypeSet::ALL) +} + +/// Infer the accepted variable types while constraining the expression's result type. +pub fn infer_types_with_target( + expr: &UntypedExpr, + target_type: InferredTypeSet, +) -> Result, TypeError> { + let mut inferred_type_res = HashMap::default(); + infer_types_aux(expr, target_type, &mut inferred_type_res)?; + Ok(inferred_type_res) +} + +pub(crate) fn infer_types_aux<'a>( + expr: &'a UntypedExpr, + target_inferred_type: InferredTypeSet, + inferred_types_res: &mut HashMap<&'a str, InferredTypeSet>, +) -> Result { + match expr { + UntypedExpr::Literal(literal) => { + let literal_type: InferredTypeSet = target_inferred_type.intersect(literal.types()); + if literal_type.is_none() { + return Err(TypeError::InvalidLiteralType { + literal: literal.clone(), + expected: target_inferred_type, + }); + } + Ok(literal_type) + } + UntypedExpr::Variable(variable_name) => { + match inferred_types_res.entry(variable_name.as_ref()) { + Entry::Occupied(mut occupied_entry) => { + let inferred_types = occupied_entry.get().intersect(target_inferred_type); + occupied_entry.insert(inferred_types); + Ok(inferred_types) + } + Entry::Vacant(vacant_entry) => { + vacant_entry.insert_entry(target_inferred_type); + Ok(target_inferred_type) + } + } + } + UntypedExpr::Call { function, args } => { + function.infer_types(args, target_inferred_type, inferred_types_res) + } + } +} + +pub(crate) fn infer_type_with_variable_types( + expr: &UntypedExpr, + target_inferred_type: InferredTypeSet, + variable_types: &HashMap<&str, VarType>, +) -> Result { + let mut inferred_types = HashMap::new(); + seed_variable_types(expr, variable_types, &mut inferred_types); + infer_types_aux(expr, target_inferred_type, &mut inferred_types) +} + +fn seed_variable_types<'a>( + expr: &'a UntypedExpr, + variable_types: &HashMap<&str, VarType>, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, +) { + match expr { + UntypedExpr::Literal(_) => {} + UntypedExpr::Variable(variable_name) => { + let inferred_type = variable_types + .get(variable_name.as_ref()) + .copied() + .map(InferredTypeSet::from) + .unwrap_or(InferredTypeSet::NONE); + inferred_types.insert(variable_name.as_ref(), inferred_type); + } + UntypedExpr::Call { args, .. } => { + for arg in args { + seed_variable_types(arg, variable_types, inferred_types); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_infer_types_bare_variable_accepts_all() { + // A lone variable should accept all types. + let expr = UntypedExpr::variable("a"); + let inferred_types = infer_types(&expr).unwrap(); + let a_types = inferred_types.get("a").unwrap(); + assert_eq!(a_types, &InferredTypeSet::ALL); + } + + #[test] + fn test_infer_type_uses_concrete_variable_types() { + let expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("my_col"), + UntypedExpr::literal(1i64), + ]); + let variable_types = HashMap::from([("my_col", VarType::U64)]); + + let inferred_type = + infer_type_with_variable_types(&expr, InferredTypeSet::NUMERICAL, &variable_types) + .unwrap(); + + assert_eq!(inferred_type, InferredTypeSet::U64); + } + + #[test] + fn test_infer_type_falls_back_to_f64_for_disjoint_numeric_types() { + let expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("unsigned"), + UntypedExpr::variable("signed"), + ]); + let variable_types = HashMap::from([("unsigned", VarType::U64), ("signed", VarType::I64)]); + + let inferred_type = + infer_type_with_variable_types(&expr, InferredTypeSet::NUMERICAL, &variable_types) + .unwrap(); + + assert_eq!(inferred_type, InferredTypeSet::F64); + } + + #[test] + fn test_inferred_type_set_display_lists_concrete_numeric_types() { + assert_eq!( + InferredTypeSet::ALL.to_string(), + "{string, i64, u64, f64, boolean}" + ); + assert_eq!(InferredTypeSet::NUMERICAL.to_string(), "{i64, u64, f64}"); + } +} diff --git a/jitexpr/src/ast/literal.rs b/jitexpr/src/ast/literal.rs new file mode 100644 index 0000000000..d6643cd63e --- /dev/null +++ b/jitexpr/src/ast/literal.rs @@ -0,0 +1,182 @@ +use std::sync::Arc; + +use crate::ast::InferredTypeSet; +use crate::types::VarType; + +/// A literal supported by the first expression-language milestone. +#[derive(Clone, Debug, PartialEq)] +pub enum Literal { + None, + Bool(bool), + U64(u64), + I64(i64), + F64(f64), + String(Arc), +} + +impl Literal { + pub fn is_none(&self) -> bool { + matches!(self, Literal::None) + } + + pub fn types(&self) -> InferredTypeSet { + match self { + Literal::None => InferredTypeSet::ALL, + Literal::Bool(_) => InferredTypeSet::BOOLEAN, + // A literal number represents a "real number". It can sometime be represented by a i64, + // a u64 or a f64. The choice of this representation is rather arbitrary. It + // can be the result of an implementation detail of serde_json for instance. + // + // Here we want to return the set of possible representation for the associated number. + Literal::I64(value) => InferredTypeSet { + i64: true, + u64: *value >= 0, // Any non-negative i64 can be represented as u64. + f64: true, // We always accept f64. + ..InferredTypeSet::NONE + }, + Literal::U64(value) => InferredTypeSet { + i64: *value <= i64::MAX as u64, // any u64 below i64::MAX can be represented as a + // i64. + u64: true, + f64: true, // We always accept f64 + ..InferredTypeSet::NONE + }, + Literal::F64(value) => { + let is_integral = value.is_finite() && value.fract() == 0.0; + InferredTypeSet { + i64: is_integral && *value >= i64::MIN as f64 && *value < -(i64::MIN as f64), + u64: is_integral && *value >= 0.0 && *value < u64::MAX as f64, + f64: true, + ..InferredTypeSet::NONE + } + } + Literal::String(_) => InferredTypeSet::STRING, + } + } + + // TODO let's remove it + pub fn r#type(&self) -> VarType { + match self { + Literal::None => VarType::None, + Literal::Bool(_) => VarType::Bool, + Literal::U64(_) => VarType::U64, + Literal::I64(_) => VarType::I64, + Literal::F64(_) => VarType::F64, + Literal::String(_) => VarType::Str, + } + } +} + +impl From for Literal { + fn from(value: bool) -> Self { + Literal::Bool(value) + } +} + +impl From for Literal { + fn from(value: u64) -> Self { + Literal::U64(value) + } +} + +impl From for Literal { + fn from(value: i64) -> Self { + Literal::I64(value) + } +} + +impl From for Literal { + fn from(value: f64) -> Self { + Literal::F64(value) + } +} + +impl From for Literal { + fn from(value: String) -> Self { + Literal::String(Arc::from(value)) + } +} + +impl From<&str> for Literal { + fn from(value: &str) -> Self { + Literal::String(Arc::from(value.to_string())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_literal_types_depend_on_representable_value() { + let i64_f64 = InferredTypeSet { + i64: true, + f64: true, + ..InferredTypeSet::NONE + }; + let u64_f64 = InferredTypeSet { + u64: true, + f64: true, + ..InferredTypeSet::NONE + }; + + assert_eq!(Literal::U64(1).types(), InferredTypeSet::NUMERICAL); + assert_eq!(Literal::I64(1).types(), InferredTypeSet::NUMERICAL); + assert_eq!(Literal::I64(-1).types(), i64_f64); + assert_eq!(Literal::U64(1 << 63).types(), u64_f64); + assert_eq!(Literal::F64(1.2).types(), InferredTypeSet::F64); + assert_eq!(Literal::F64(1.0).types(), InferredTypeSet::NUMERICAL); + } + + #[test] + fn test_literal_types_accept_lossless_float_representation() { + assert_eq!( + Literal::I64((1 << 53) + 1).types(), + InferredTypeSet::NUMERICAL + ); + + // even though i64::MAX - 1i64 cannot be represented as f64 in a lossless manner... + assert_ne!(((i64::MAX - 1i64) as f64) as i64, (i64::MAX - 1)); + // ... we list f64 as a valid inferred type. + assert_eq!( + Literal::I64(i64::MAX - 1).types(), + InferredTypeSet::NUMERICAL + ); + assert_eq!( + Literal::U64(u64::MAX).types(), + InferredTypeSet { + u64: true, + f64: true, + ..InferredTypeSet::NONE + } + ); + assert_eq!( + Literal::I64(i64::MIN).types(), + InferredTypeSet { + i64: true, + f64: true, + ..InferredTypeSet::NONE + } + ); + } + + #[test] + fn test_f64_literal_types_handle_integer_boundaries_and_special_values() { + assert_eq!( + Literal::F64(2f64.powi(63)).types(), + InferredTypeSet { + u64: true, + f64: true, + ..InferredTypeSet::NONE + } + ); + assert_eq!(Literal::F64(2f64.powi(64)).types(), InferredTypeSet::F64); + assert_eq!(Literal::F64(-0.0).types(), InferredTypeSet::NUMERICAL); + assert_eq!(Literal::F64(f64::NAN).types(), InferredTypeSet::F64); + assert_eq!(Literal::F64(f64::INFINITY).types(), InferredTypeSet::F64); + assert_eq!( + Literal::F64(f64::NEG_INFINITY).types(), + InferredTypeSet::F64 + ); + } +} diff --git a/jitexpr/src/ast/mod.rs b/jitexpr/src/ast/mod.rs new file mode 100644 index 0000000000..e7c3a396e0 --- /dev/null +++ b/jitexpr/src/ast/mod.rs @@ -0,0 +1,12 @@ +mod infer_types; +mod literal; +mod serialize; +mod untyped_expr; + +pub use infer_types::{InferredTypeSet, TypeError, infer_types, infer_types_with_target}; +pub(crate) use infer_types::{infer_type_with_variable_types, infer_types_aux}; +pub use literal::Literal; +pub use serialize::{DeserializeError, deserialize, serialize}; +pub use untyped_expr::UntypedExpr; + +pub use crate::functions::{Function, InvalidFunctionCall}; diff --git a/jitexpr/src/ast/serialize.rs b/jitexpr/src/ast/serialize.rs new file mode 100644 index 0000000000..1708356ca0 --- /dev/null +++ b/jitexpr/src/ast/serialize.rs @@ -0,0 +1,589 @@ +//! Serialization for [`UntypedExpr`] using a small Lisp-like syntax. +//! +//! Calls are lists whose first item is an uppercase function name, while +//! lowercase identifiers name variables. For example: +//! +//! ```text +//! (ADD 1i64 my_col) +//! ``` +//! +//! Numerical literals always carry a type suffix. The other literals are +//! `none`, `true`, `false`, and quoted strings. Strings use backslash escapes. + +use std::fmt; +use std::sync::Arc; + +use crate::ast::{Function, Literal, UntypedExpr}; + +/// Serializes an untyped expression into its canonical Lisp-like form. +pub fn serialize(expr: &UntypedExpr) -> String { + expr.to_string() +} + +/// Deserializes an untyped expression from its Lisp-like form. +pub fn deserialize(input: &str) -> Result { + Parser::new(input).parse() +} + +/// An error encountered while deserializing an [`UntypedExpr`]. +#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)] +#[error("failed to deserialize expression at byte {offset}: {message}")] +pub struct DeserializeError { + offset: usize, + message: String, +} + +impl DeserializeError { + fn new(offset: usize, message: impl Into) -> Self { + Self { + offset, + message: message.into(), + } + } + + /// Returns the byte offset at which parsing failed. + pub fn offset(&self) -> usize { + self.offset + } + + /// Returns a description of the parsing failure. + pub fn message(&self) -> &str { + &self.message + } +} + +impl fmt::Display for UntypedExpr { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + format_expr(self, formatter) + } +} + +impl fmt::Debug for UntypedExpr { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + format_expr(self, formatter) + } +} + +impl std::str::FromStr for UntypedExpr { + type Err = DeserializeError; + + fn from_str(input: &str) -> Result { + deserialize(input) + } +} + +fn format_expr(expr: &UntypedExpr, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match expr { + UntypedExpr::Literal(literal) => format_literal(literal, formatter), + UntypedExpr::Variable(variable_name) => formatter.write_str(variable_name), + UntypedExpr::Call { function, args } => { + write!(formatter, "({}", function_name(*function))?; + for arg in args { + write!(formatter, " {arg}")?; + } + formatter.write_str(")") + } + } +} + +fn format_literal(literal: &Literal, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match literal { + Literal::None => formatter.write_str("none"), + Literal::Bool(value) => write!(formatter, "{value}"), + Literal::U64(value) => write!(formatter, "{value}u64"), + Literal::I64(value) => write!(formatter, "{value}i64"), + Literal::F64(value) => write!(formatter, "{value}f64"), + Literal::String(value) => format_string(value, formatter), + } +} + +fn format_string(value: &str, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("\"")?; + for character in value.chars() { + match character { + '\"' => formatter.write_str("\\\""), + '\\' => formatter.write_str("\\\\"), + '\n' => formatter.write_str("\\n"), + '\r' => formatter.write_str("\\r"), + '\t' => formatter.write_str("\\t"), + '\0' => formatter.write_str("\\0"), + character if character.is_control() => { + write!(formatter, "{}", character.escape_unicode()) + } + character => write!(formatter, "{character}"), + }?; + } + formatter.write_str("\"") +} + +fn function_name(function: Function) -> &'static str { + match function { + Function::Abs => "ABS", + Function::And => "AND", + Function::Ceil => "CEIL", + Function::Concat => "CONCAT", + Function::Add => "ADD", + Function::Divide => "DIVIDE", + Function::Eq => "EQ", + Function::Floor => "FLOOR", + Function::Gt => "GT", + Function::GtEq => "GT_EQ", + Function::If => "IF", + Function::IntMod => "INT_MOD", + Function::Left => "LEFT", + Function::Lt => "LT", + Function::LtEq => "LT_EQ", + Function::IsNull => "IS_NULL", + Function::IsNotNull => "IS_NOT_NULL", + Function::Lower => "LOWER", + Function::Max => "MAX", + Function::Min => "MIN", + Function::Multiply => "MULTIPLY", + Function::Neq => "NEQ", + Function::Not => "NOT", + Function::Or => "OR", + Function::Pow => "POW", + Function::Sqrt => "SQRT", + Function::RegexpExtract => "REGEXP_EXTRACT", + Function::RegexpLike => "REGEXP_LIKE", + Function::Right => "RIGHT", + Function::Round => "ROUND", + Function::SplitAfter => "SPLIT_AFTER", + Function::SplitBefore => "SPLIT_BEFORE", + Function::Subtract => "SUBTRACT", + Function::Substring => "SUBSTRING", + Function::SubstringCount => "SUBSTRING_COUNT", + Function::TextJoin => "TEXT_JOIN", + Function::Trim => "TRIM", + Function::Upper => "UPPER", + } +} + +fn parse_function(name: &str, offset: usize) -> Result { + match name { + "ABS" => Ok(Function::Abs), + "AND" => Ok(Function::And), + "CEIL" => Ok(Function::Ceil), + "CONCAT" => Ok(Function::Concat), + "ADD" => Ok(Function::Add), + "DIVIDE" => Ok(Function::Divide), + "EQ" => Ok(Function::Eq), + "FLOOR" => Ok(Function::Floor), + "GT" => Ok(Function::Gt), + "GT_EQ" => Ok(Function::GtEq), + "IF" => Ok(Function::If), + "INT_MOD" => Ok(Function::IntMod), + "LEFT" => Ok(Function::Left), + "LT" => Ok(Function::Lt), + "LT_EQ" => Ok(Function::LtEq), + "IS_NULL" => Ok(Function::IsNull), + "IS_NOT_NULL" => Ok(Function::IsNotNull), + "LOWER" => Ok(Function::Lower), + "MAX" => Ok(Function::Max), + "MIN" => Ok(Function::Min), + "MULTIPLY" => Ok(Function::Multiply), + "NEQ" => Ok(Function::Neq), + "NOT" => Ok(Function::Not), + "OR" => Ok(Function::Or), + "POW" => Ok(Function::Pow), + "SQRT" => Ok(Function::Sqrt), + "REGEXP_EXTRACT" => Ok(Function::RegexpExtract), + "REGEXP_LIKE" => Ok(Function::RegexpLike), + "RIGHT" => Ok(Function::Right), + "ROUND" => Ok(Function::Round), + "SPLIT_AFTER" => Ok(Function::SplitAfter), + "SPLIT_BEFORE" => Ok(Function::SplitBefore), + "SUBTRACT" => Ok(Function::Subtract), + "SUBSTRING" => Ok(Function::Substring), + "SUBSTRING_COUNT" => Ok(Function::SubstringCount), + "TEXT_JOIN" => Ok(Function::TextJoin), + "TRIM" => Ok(Function::Trim), + "UPPER" => Ok(Function::Upper), + _ if !is_function_name(name) => Err(DeserializeError::new( + offset, + format!("function name `{name}` must be uppercase"), + )), + _ => Err(DeserializeError::new( + offset, + format!("unknown function `{name}`"), + )), + } +} + +fn is_function_name(name: &str) -> bool { + let mut chars = name.chars(); + matches!(chars.next(), Some(first) if first.is_ascii_uppercase()) + && chars.all(|character| { + character.is_ascii_uppercase() || character.is_ascii_digit() || character == '_' + }) +} + +fn is_variable_name(name: &str) -> bool { + let mut chars = name.chars(); + matches!(chars.next(), Some(first) if first.is_ascii_lowercase()) + && chars.all(|character| { + character.is_ascii_lowercase() || character.is_ascii_digit() || character == '_' + }) +} + +struct Parser<'a> { + input: &'a str, + offset: usize, +} + +impl<'a> Parser<'a> { + fn new(input: &'a str) -> Self { + Self { input, offset: 0 } + } + + fn parse(mut self) -> Result { + self.skip_whitespace(); + let expr = self.parse_expr()?; + self.skip_whitespace(); + if self.peek().is_some() { + return Err(DeserializeError::new( + self.offset, + "unexpected characters after expression", + )); + } + Ok(expr) + } + + fn parse_expr(&mut self) -> Result { + self.skip_whitespace(); + match self.peek() { + Some('(') => self.parse_call(), + Some('"') => self + .parse_string() + .map(|value| UntypedExpr::Literal(Literal::String(Arc::from(value)))), + Some(')') => Err(DeserializeError::new( + self.offset, + "unexpected closing parenthesis", + )), + Some(_) => self.parse_atom(), + None => Err(DeserializeError::new(self.offset, "expected an expression")), + } + } + + fn parse_call(&mut self) -> Result { + let call_offset = self.offset; + self.advance(); + self.skip_whitespace(); + + if self.peek().is_none() { + return Err(DeserializeError::new( + call_offset, + "unterminated function call", + )); + } + if self.peek() == Some(')') { + return Err(DeserializeError::new( + self.offset, + "expected a function name", + )); + } + + let function_offset = self.offset; + let function_name = self.take_atom(); + if function_name.is_empty() { + return Err(DeserializeError::new( + function_offset, + "expected an uppercase function name", + )); + } + let function = parse_function(function_name, function_offset)?; + + let mut args = Vec::new(); + loop { + self.skip_whitespace(); + match self.peek() { + Some(')') => { + self.advance(); + return Ok(UntypedExpr::Call { function, args }); + } + Some(_) => args.push(self.parse_expr()?), + None => { + return Err(DeserializeError::new( + call_offset, + "unterminated function call", + )); + } + } + } + } + + fn parse_atom(&mut self) -> Result { + let atom_offset = self.offset; + let atom = self.take_atom(); + match atom { + "none" => Ok(UntypedExpr::Literal(Literal::None)), + "true" => Ok(UntypedExpr::Literal(Literal::Bool(true))), + "false" => Ok(UntypedExpr::Literal(Literal::Bool(false))), + _ => self.parse_number_or_variable(atom, atom_offset), + } + } + + fn parse_number_or_variable( + &self, + atom: &str, + atom_offset: usize, + ) -> Result { + if let Some(value) = atom.strip_suffix("u64") + && let Ok(value) = value.parse::() + { + return Ok(UntypedExpr::Literal(Literal::U64(value))); + } + if let Some(value) = atom.strip_suffix("i64") + && let Ok(value) = value.parse::() + { + return Ok(UntypedExpr::Literal(Literal::I64(value))); + } + if let Some(value) = atom.strip_suffix("f64") + && let Ok(value) = value.parse::() + { + return Ok(UntypedExpr::Literal(Literal::F64(value))); + } + + if is_variable_name(atom) { + return Ok(UntypedExpr::Variable(Arc::from(atom))); + } + + if is_function_name(atom) { + return Err(DeserializeError::new( + atom_offset, + format!("function `{atom}` must be the first item in a list"), + )); + } + + Err(DeserializeError::new( + atom_offset, + format!("invalid literal or identifier `{atom}`"), + )) + } + + fn parse_string(&mut self) -> Result { + let string_offset = self.offset; + self.advance(); + let mut value = String::new(); + + loop { + let character_offset = self.offset; + let Some(character) = self.advance() else { + return Err(DeserializeError::new( + string_offset, + "unterminated string literal", + )); + }; + match character { + '"' => return Ok(value), + '\\' => value.push(self.parse_escape(character_offset)?), + character if character.is_control() => { + return Err(DeserializeError::new( + character_offset, + "unescaped control character in string literal", + )); + } + character => value.push(character), + } + } + } + + fn parse_escape(&mut self, escape_offset: usize) -> Result { + let Some(escaped) = self.advance() else { + return Err(DeserializeError::new( + escape_offset, + "unterminated string escape", + )); + }; + match escaped { + '"' => Ok('"'), + '\\' => Ok('\\'), + 'n' => Ok('\n'), + 'r' => Ok('\r'), + 't' => Ok('\t'), + '0' => Ok('\0'), + 'u' => self.parse_unicode_escape(escape_offset), + _ => Err(DeserializeError::new( + escape_offset, + format!("unsupported string escape `\\{escaped}`"), + )), + } + } + + fn parse_unicode_escape(&mut self, escape_offset: usize) -> Result { + if self.advance() != Some('{') { + return Err(DeserializeError::new( + escape_offset, + "Unicode escape must start with `\\u{`", + )); + } + + let digits_offset = self.offset; + while matches!(self.peek(), Some(character) if character.is_ascii_hexdigit()) { + self.advance(); + } + let digits = &self.input[digits_offset..self.offset]; + if digits.is_empty() || self.advance() != Some('}') { + return Err(DeserializeError::new( + escape_offset, + "invalid Unicode escape", + )); + } + + let codepoint = u32::from_str_radix(digits, 16).ok(); + codepoint + .and_then(char::from_u32) + .ok_or_else(|| DeserializeError::new(escape_offset, "invalid Unicode scalar value")) + } + + fn take_atom(&mut self) -> &'a str { + let start = self.offset; + while matches!(self.peek(), Some(character) if !is_delimiter(character)) { + self.advance(); + } + &self.input[start..self.offset] + } + + fn skip_whitespace(&mut self) { + while matches!(self.peek(), Some(character) if character.is_whitespace()) { + self.advance(); + } + } + + fn peek(&self) -> Option { + self.input[self.offset..].chars().next() + } + + fn advance(&mut self) -> Option { + let character = self.peek()?; + self.offset += character.len_utf8(); + Some(character) + } +} + +fn is_delimiter(character: char) -> bool { + character.is_whitespace() || matches!(character, '(' | ')' | '"') +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_serialize_example() { + let expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::literal(1i64), + UntypedExpr::variable("my_col"), + ]); + + assert_eq!(serialize(&expr), "(ADD 1i64 my_col)"); + assert_eq!(format!("{expr}"), "(ADD 1i64 my_col)"); + 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 = [ + (UntypedExpr::Literal(Literal::None), "none"), + (UntypedExpr::literal(true), "true"), + (UntypedExpr::literal(false), "false"), + (UntypedExpr::literal(u64::MAX), "18446744073709551615u64"), + (UntypedExpr::literal(i64::MIN), "-9223372036854775808i64"), + (UntypedExpr::literal(1.5f64), "1.5f64"), + (UntypedExpr::literal(1.0f64), "1f64"), + ]; + + for (expr, expected) in cases { + assert_eq!(serialize(&expr), expected); + assert_eq!(deserialize(expected).unwrap(), expr); + } + } + + #[test] + fn test_nested_call_and_escaped_string_round_trip() { + let string = "quoted: \"hello\"\\world\n\t\0\u{7} café"; + let regexp_extract = Function::RegexpExtract.call_untyped_expr(vec![ + UntypedExpr::variable("message"), + UntypedExpr::literal(string), + UntypedExpr::literal(1u64), + ]); + let expr = + Function::Add.call_untyped_expr(vec![regexp_extract, UntypedExpr::literal(2i64)]); + + let serialized = serialize(&expr); + assert_eq!( + serialized, + "(ADD (REGEXP_EXTRACT message \"quoted: \\\"hello\\\"\\\\world\\n\\t\\0\\u{7} café\" \ + 1u64) 2i64)" + ); + assert_eq!(deserialize(&serialized).unwrap(), expr); + } + + #[test] + fn test_deserialize_accepts_whitespace() { + let parsed = deserialize(" \n ( ADD\t1i64\nmy_col ) \r").unwrap(); + let expected = Function::Add.call_untyped_expr(vec![ + UntypedExpr::literal(1i64), + UntypedExpr::variable("my_col"), + ]); + assert_eq!(parsed, expected); + } + + #[test] + fn test_float_special_values_round_trip() { + for value in [f64::INFINITY, f64::NEG_INFINITY, -0.0] { + let serialized = serialize(&UntypedExpr::literal(value)); + let UntypedExpr::Literal(Literal::F64(parsed)) = deserialize(&serialized).unwrap() + else { + panic!("expected an f64 literal"); + }; + assert_eq!(parsed.to_bits(), value.to_bits()); + } + + let serialized = serialize(&UntypedExpr::literal(f64::NAN)); + let UntypedExpr::Literal(Literal::F64(parsed)) = deserialize(&serialized).unwrap() else { + panic!("expected an f64 literal"); + }; + assert!(parsed.is_nan()); + } + + #[test] + fn test_from_str() { + let parsed: UntypedExpr = "(ADD 3u64 value)".parse().unwrap(); + assert_eq!(serialize(&parsed), "(ADD 3u64 value)"); + } + + #[test] + fn test_deserialize_errors() { + let cases = [ + ("", 0, "expected an expression"), + ("()", 1, "expected a function name"), + ("(add 1i64)", 1, "must be uppercase"), + ("(UNKNOWN 1i64)", 1, "unknown function"), + ("ADD", 0, "must be the first item in a list"), + ("1i32", 0, "invalid literal or identifier"), + ("\"unterminated", 0, "unterminated string literal"), + ("\"bad\\x\"", 4, "unsupported string escape"), + ("(ADD 1i64", 0, "unterminated function call"), + ("value other", 6, "unexpected characters after expression"), + ]; + + for (input, offset, expected_message) in cases { + let error = deserialize(input).unwrap_err(); + assert_eq!(error.offset(), offset, "input: {input}"); + assert!( + error.message().contains(expected_message), + "input: {input}; error: {error}" + ); + } + } +} diff --git a/jitexpr/src/ast/untyped_expr.rs b/jitexpr/src/ast/untyped_expr.rs new file mode 100644 index 0000000000..c19933ff68 --- /dev/null +++ b/jitexpr/src/ast/untyped_expr.rs @@ -0,0 +1,43 @@ +use std::sync::Arc; + +use crate::ast::{Function, Literal}; +use crate::functions::InvalidFunctionCall; + +/// An expression AST. +/// +/// The expression at this point is untyped and not necessarily valid. +#[derive(Clone, PartialEq)] +pub enum UntypedExpr { + Literal(Literal), + Variable(Arc), + Call { + function: Function, + args: Vec, + }, +} + +impl UntypedExpr { + pub fn literal(val: impl Into) -> UntypedExpr { + UntypedExpr::Literal(val.into()) + } + + pub fn variable(variable_name: impl ToString) -> UntypedExpr { + UntypedExpr::Variable(Arc::from(variable_name.to_string())) + } + + /// Creates an untyped expression that is a function over different arguments. + /// + /// This call will validate the arguments and + pub fn call( + function: Function, + args: Vec, + ) -> Result { + function.call(args) + } +} + +impl From for UntypedExpr { + fn from(literal: Literal) -> Self { + UntypedExpr::Literal(literal) + } +} diff --git a/jitexpr/src/bin/jitexpr-asm.rs b/jitexpr/src/bin/jitexpr-asm.rs new file mode 100644 index 0000000000..57f4b01342 --- /dev/null +++ b/jitexpr/src/bin/jitexpr-asm.rs @@ -0,0 +1,361 @@ +//! Prints typed expressions and native assembly for serialized `UntypedExpr` values read from +//! stdin. +//! +//! The serialization does not attach concrete types to variables. This tool +//! therefore uses `Str` for string variables, `Bool` for boolean variables, +//! and `F64` for numerical or otherwise unconstrained variables. +//! ANSI colors are enabled when stdout is a terminal and can be disabled with `NO_COLOR`. + +use std::collections::HashMap; +use std::io::{self, BufRead, IsTerminal, Write}; +use std::process::ExitCode; +use std::time::{Duration, Instant}; + +use jitexpr::ast::{ + DeserializeError, InferredTypeSet, TypeError, deserialize, infer_types, + serialize as serialize_untyped, +}; +use jitexpr::compile::{CompileError, compile, compile_to_assembly, serialize as serialize_typed}; +use jitexpr::types::VarType; + +const RESET: &str = "\x1b[0m"; +const BOLD_CYAN: &str = "\x1b[1;36m"; +const BOLD_MAGENTA: &str = "\x1b[1;35m"; +const BLUE: &str = "\x1b[34m"; +const CYAN: &str = "\x1b[36m"; +const GREEN: &str = "\x1b[32m"; +const YELLOW: &str = "\x1b[33m"; +const DIM: &str = "\x1b[2m"; + +#[derive(Debug, thiserror::Error)] +enum ExpressionError { + #[error(transparent)] + Deserialize(#[from] DeserializeError), + #[error(transparent)] + Type(#[from] TypeError), + #[error(transparent)] + Compile(#[from] CompileError), +} + +struct LineOutput { + input_expression: String, + typed_expression: String, + assembly: String, + codegen_duration: Duration, +} + +fn main() -> ExitCode { + let stdin = io::stdin(); + let stdout = io::stdout(); + let stderr = io::stderr(); + let color = stdout.is_terminal() && std::env::var_os("NO_COLOR").is_none(); + match process_lines(stdin.lock(), stdout.lock(), stderr.lock(), color) { + Ok(true) => ExitCode::SUCCESS, + Ok(false) => ExitCode::FAILURE, + Err(error) => { + eprintln!("I/O error: {error}"); + ExitCode::FAILURE + } + } +} + +/// Returns whether every non-empty input line compiled successfully. +fn process_lines( + input: impl BufRead, + mut output: impl Write, + mut errors: impl Write, + color: bool, +) -> io::Result { + let mut all_succeeded = true; + let mut wrote_output = false; + + for (line_index, line) in input.lines().enumerate() { + let line = line?; + if line.trim().is_empty() { + continue; + } + + match compile_line(&line) { + Ok(line_output) => { + if wrote_output { + writeln!(output)?; + write_styled(&mut output, DIM, &"-".repeat(80), color)?; + writeln!(output, "\n")?; + } + write_header(&mut output, "Input expression", color)?; + write_expression(&mut output, &line_output.input_expression, color)?; + writeln!(output, "\n")?; + + write_header(&mut output, "Typed expression", color)?; + write_expression(&mut output, &line_output.typed_expression, color)?; + writeln!(output, "\n")?; + + write_styled(&mut output, BOLD_CYAN, "Code generation time:", color)?; + write!(output, " ")?; + write_styled( + &mut output, + YELLOW, + &format!("{:?}", line_output.codegen_duration), + color, + )?; + writeln!(output, "\n")?; + + write_header(&mut output, "Assembly", color)?; + write_assembly(&mut output, &line_output.assembly, color)?; + wrote_output = true; + } + Err(error) => { + writeln!(errors, "line {}: {error}", line_index + 1)?; + all_succeeded = false; + } + } + } + + Ok(all_succeeded) +} + +fn write_header(output: &mut impl Write, header: &str, color: bool) -> io::Result<()> { + write_styled(output, BOLD_CYAN, header, color)?; + writeln!(output, ":") +} + +fn write_styled(output: &mut impl Write, style: &str, text: &str, color: bool) -> io::Result<()> { + if color { + write!(output, "{style}{text}{RESET}") + } else { + output.write_all(text.as_bytes()) + } +} + +fn write_expression(output: &mut impl Write, expression: &str, color: bool) -> io::Result<()> { + if !color { + return output.write_all(expression.as_bytes()); + } + + let mut offset = 0; + while offset < expression.len() { + let character = expression[offset..] + .chars() + .next() + .expect("offset is before end of expression"); + if character == '"' { + let end = quoted_literal_end(expression, offset); + write_styled(output, GREEN, &expression[offset..end], true)?; + offset = end; + } else if character.is_whitespace() { + write!(output, "{character}")?; + offset += character.len_utf8(); + } else if matches!(character, '(' | ')' | '[' | ']' | ':') { + write_styled(output, DIM, &expression[offset..offset + 1], true)?; + offset += 1; + } else { + let end = expression[offset..] + .char_indices() + .find_map(|(relative_offset, character)| { + (relative_offset > 0 + && (character.is_whitespace() + || matches!(character, '(' | ')' | '[' | ']' | ':' | '"'))) + .then_some(offset + relative_offset) + }) + .unwrap_or(expression.len()); + let token = &expression[offset..end]; + let style = expression_token_style(token, &expression[end..]); + write_styled(output, style, token, true)?; + offset = end; + } + } + Ok(()) +} + +fn quoted_literal_end(expression: &str, quote_offset: usize) -> usize { + let mut escaped = false; + for (relative_offset, character) in expression[quote_offset + 1..].char_indices() { + if escaped { + escaped = false; + } else if character == '\\' { + escaped = true; + } else if character == '"' { + return quote_offset + 1 + relative_offset + 1; + } + } + expression.len() +} + +fn expression_token_style(token: &str, suffix: &str) -> &'static str { + if suffix.starts_with(':') + && matches!( + token, + "boolean" | "float64" | "uint64" | "int64" | "string" | "none" + ) + { + CYAN + } else if token + .chars() + .next() + .is_some_and(|character| character.is_ascii_uppercase()) + && token.chars().all(|character| { + character.is_ascii_uppercase() || character.is_ascii_digit() || character == '_' + }) + { + BOLD_MAGENTA + } else if matches!(token, "true" | "false" | "none") + || token.ends_with("i64") + || token.ends_with("u64") + || token.ends_with("f64") + { + YELLOW + } else { + BLUE + } +} + +fn write_assembly(output: &mut impl Write, assembly: &str, color: bool) -> io::Result<()> { + if !color { + output.write_all(assembly.as_bytes())?; + if !assembly.ends_with('\n') { + writeln!(output)?; + } + return Ok(()); + } + + for line in assembly.lines() { + let instruction = line.trim_start(); + let indentation = &line[..line.len() - instruction.len()]; + output.write_all(indentation.as_bytes())?; + if instruction.starts_with("block") && instruction.ends_with(':') { + write_styled(output, BOLD_MAGENTA, instruction, true)?; + } else if let Some(opcode_end) = instruction.find(char::is_whitespace) { + write_styled(output, CYAN, &instruction[..opcode_end], true)?; + output.write_all(&instruction.as_bytes()[opcode_end..])?; + } else { + write_styled(output, CYAN, instruction, true)?; + } + writeln!(output)?; + } + Ok(()) +} + +fn compile_line(line: &str) -> Result { + let expression = deserialize(line)?; + let inferred_types = infer_types(&expression)?; + let variable_types = inferred_types + .into_iter() + .map(|(name, inferred_type)| (name, concrete_type(inferred_type))) + .collect::>(); + + let codegen_start = Instant::now(); + let compiled_fn = compile(&expression, &variable_types)?; + let codegen_duration = codegen_start.elapsed(); + drop(compiled_fn); + + let input_expression = serialize_untyped(&expression); + let typed_expression = serialize_typed(&expression, &variable_types)?; + let assembly = compile_to_assembly(&expression, &variable_types)?; + Ok(LineOutput { + input_expression, + typed_expression, + assembly, + codegen_duration, + }) +} + +fn concrete_type(inferred_type: InferredTypeSet) -> VarType { + if inferred_type == InferredTypeSet::STRING { + VarType::Str + } else if inferred_type == InferredTypeSet::BOOLEAN { + VarType::Bool + } else if inferred_type.i64 { + VarType::I64 + } else if inferred_type.u64 { + VarType::U64 + } else if inferred_type.f64 { + VarType::F64 + } else if inferred_type.string { + VarType::Str + } else if inferred_type.boolean { + VarType::Bool + } else { + VarType::None + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compile_line_infers_variable_type() { + let line_output = compile_line("(ADD 1i64 my_col)").unwrap(); + + assert_eq!(line_output.input_expression, "(ADD 1i64 my_col)"); + assert_eq!( + line_output.typed_expression, + "[int64: ADD 1i64 [int64: my_col]]" + ); + assert!(line_output.assembly.contains("block0:")); + assert!(!line_output.assembly.trim().is_empty()); + } + + #[test] + fn test_concrete_type_selects_an_inferred_numeric_type() { + assert_eq!(concrete_type(InferredTypeSet::ALL), VarType::I64); + assert_eq!(concrete_type(InferredTypeSet::I64), VarType::I64); + assert_eq!(concrete_type(InferredTypeSet::U64), VarType::U64); + assert_eq!(concrete_type(InferredTypeSet::F64), VarType::F64); + assert_eq!(concrete_type(InferredTypeSet::NONE), VarType::None); + } + + #[test] + fn test_process_lines_continues_after_an_error() { + let input = b"1i64\nnot-valid!\n2u64\n".as_slice(); + let mut output = Vec::new(); + let mut errors = Vec::new(); + + let all_succeeded = process_lines(input, &mut output, &mut errors, false).unwrap(); + + assert!(!all_succeeded); + let output = String::from_utf8(output).unwrap(); + assert_eq!(output.matches("Code generation time:").count(), 2); + assert_eq!(output.matches("Typed expression:").count(), 2); + assert_eq!(output.matches("Input expression:").count(), 2); + assert_eq!(output.matches(&"-".repeat(80)).count(), 1); + assert_eq!(output.matches("block0:").count(), 2); + assert!(!output.contains("\x1b[")); + assert!(String::from_utf8(errors).unwrap().contains("line 2:")); + } + + #[test] + fn test_process_lines_ignores_empty_lines() { + let mut output = Vec::new(); + let mut errors = Vec::new(); + + let all_succeeded = + process_lines(b" \n\t\n".as_slice(), &mut output, &mut errors, false).unwrap(); + + assert!(all_succeeded); + assert!(output.is_empty()); + assert!(errors.is_empty()); + } + + #[test] + fn test_process_lines_colors_terminal_output() { + let mut output = Vec::new(); + let mut errors = Vec::new(); + + let all_succeeded = process_lines( + b"(ADD 1i64 my_col)\n".as_slice(), + &mut output, + &mut errors, + true, + ) + .unwrap(); + + assert!(all_succeeded); + let output = String::from_utf8(output).unwrap(); + assert!(output.contains("\x1b[1;36mInput expression\x1b[0m:")); + assert!(output.contains("\x1b[1;35mADD\x1b[0m")); + assert!(output.contains("\x1b[36mint64\x1b[0m")); + assert!(output.contains("\x1b[33m1i64\x1b[0m")); + assert!(errors.is_empty()); + } +} diff --git a/jitexpr/src/compile/compile_fn_builder.rs b/jitexpr/src/compile/compile_fn_builder.rs new file mode 100644 index 0000000000..c039d0bb35 --- /dev/null +++ b/jitexpr/src/compile/compile_fn_builder.rs @@ -0,0 +1,502 @@ +use std::collections::HashMap; +use std::mem; +use std::sync::Arc; + +use cranelift::codegen::Context as CodegenContext; +use cranelift::codegen::control::ControlPlane; +use cranelift::codegen::ir::{MemFlagsData, UserFuncName}; +use cranelift::prelude::*; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{FuncId, Module, ModuleError, default_libcall_names}; + +use super::compiled_fn::JitEntry; +use super::{ + CompileError, CompiledFn, LoweringContext, TypedExpr, TypedExprAst, TypedLiteral, TypedVariable, +}; +use crate::ast::{InferredTypeSet, Literal, UntypedExpr}; +use crate::functions::{declare_native_functions, register_jit_symbols}; +use crate::types::VarType; + +pub(crate) struct CompileFnBuilder<'types, 'names> { + variable_types: &'types HashMap<&'names str, VarType>, + input_vars: Vec, +} + +struct LoweredFunction { + module: JITModule, + context: CodegenContext, + function_id: FuncId, + input_vars: Vec, + expression: Box, +} + +fn make_jit_builder() -> Result { + let mut shared_flags = settings::builder(); + shared_flags + .set("opt_level", "speed") + .map_err(ModuleError::from)?; + shared_flags + .set("use_colocated_libcalls", "false") + .map_err(ModuleError::from)?; + shared_flags + .set("is_pic", "false") + .map_err(ModuleError::from)?; + + let isa_builder = cranelift_native::builder().unwrap_or_else(|message| { + panic!("host machine is not supported: {message}"); + }); + #[cfg(target_arch = "aarch64")] + let isa_builder = { + let mut isa_builder = isa_builder; + isa_builder + .set("sign_return_address", "false") + .map_err(ModuleError::from)?; + isa_builder + .set("sign_return_address_all", "false") + .map_err(ModuleError::from)?; + isa_builder + }; + let isa = isa_builder + .finish(settings::Flags::new(shared_flags)) + .map_err(ModuleError::from)?; + + Ok(JITBuilder::with_isa(isa, default_libcall_names())) +} + +impl<'types, 'names> CompileFnBuilder<'types, 'names> { + pub(crate) fn new(variable_types: &'types HashMap<&'names str, VarType>) -> Self { + CompileFnBuilder { + variable_types, + input_vars: Vec::new(), + } + } + + pub(crate) fn variable_types(&self) -> &HashMap<&'names str, VarType> { + self.variable_types + } + + /// If a variable is missing from `variable_types`, it is treated as `None`. + pub(crate) fn build_typed_expr( + &mut self, + untyped_expr: &UntypedExpr, + ) -> Result { + let mut typed_expr = self.apply_types(untyped_expr, InferredTypeSet::ALL)?; + self.assign_variable_ids(&mut typed_expr); + Ok(typed_expr) + } + + pub(crate) fn assign_variable_ids(&mut self, typed_expr: &mut TypedExpr) { + self.input_vars = assign_variable_ids(typed_expr); + } + + fn apply_literal_type( + &mut self, + literal: &Literal, + target_type_set: InferredTypeSet, + ) -> TypedLiteral { + if literal.is_none() { + return TypedLiteral::None; + } + let inferred_type_set = literal.types(); + let intersection = inferred_type_set.intersect(target_type_set); + + if intersection.contains(VarType::Bool) { + match literal { + Literal::Bool(value) => TypedLiteral::Bool(*value), + _ => panic!("cannot coerce literal {literal:?} to bool"), + } + } else if intersection.contains(VarType::I64) { + match literal { + Literal::U64(value) => TypedLiteral::I64(*value as i64), + Literal::I64(value) => TypedLiteral::I64(*value), + Literal::F64(value) if f64_to_i64_lossless(*value).is_some() => { + TypedLiteral::I64(f64_to_i64_lossless(*value).unwrap()) + } + _ => panic!("cannot coerce literal {literal:?} to i64"), + } + } else if intersection.contains(VarType::U64) { + match literal { + Literal::U64(value) => TypedLiteral::U64(*value), + Literal::I64(value) => TypedLiteral::U64(*value as u64), + Literal::F64(value) if f64_to_u64_lossless(*value).is_some() => { + TypedLiteral::U64(f64_to_u64_lossless(*value).unwrap()) + } + _ => panic!("cannot coerce literal {literal:?} to u64"), + } + } else if intersection.contains(VarType::F64) { + match literal { + Literal::U64(value) => TypedLiteral::F64(*value as f64), + Literal::I64(value) => TypedLiteral::F64(*value as f64), + Literal::F64(value) => TypedLiteral::F64(*value), + _ => panic!("cannot coerce literal {literal:?} to f64"), + } + } else if intersection.contains(VarType::Str) { + match literal { + Literal::String(value) => TypedLiteral::String(value.clone()), + _ => panic!("cannot coerce literal {literal:?} to string"), + } + } else if intersection.contains(VarType::None) { + match literal { + Literal::None => TypedLiteral::None, + _ => panic!("cannot coerce literal {literal:?} to none"), + } + } else { + panic!( + "no compatible type for literal {literal:?} with target type set \ + {target_type_set:?}" + ) + } + } + + pub(crate) fn apply_types( + &mut self, + untyped_expr: &UntypedExpr, + target_type_set: InferredTypeSet, + ) -> Result { + match untyped_expr { + UntypedExpr::Literal(literal) => { + let typed_literal = self.apply_literal_type(literal, target_type_set); + let return_type = typed_literal.r#type(); + Ok(TypedExpr { + return_type, + ast: TypedExprAst::Literal(typed_literal), + }) + } + UntypedExpr::Variable(variable_name) => { + let variable_type = self + .variable_types + .get(variable_name.as_ref()) + .copied() + .unwrap_or(VarType::None); + if variable_type == VarType::None { + Ok(TypedExpr { + return_type: VarType::None, + ast: TypedExprAst::Literal(TypedLiteral::None), + }) + } else { + let typed_expr = TypedExpr { + return_type: variable_type, + ast: TypedExprAst::variable(variable_name, variable_type), + }; + if target_type_set.contains(variable_type) { + Ok(typed_expr) + } else if let Some(target_type) = preferred_numerical_type(target_type_set) + && is_numerical(variable_type) + { + Ok(typed_expr.coerce(target_type)) + } else { + Ok(typed_expr) + } + } + } + UntypedExpr::Call { function, args } => { + function.call_with_types(args, target_type_set, self) + } + } + } + + pub(super) fn compile_typed_expr( + self, + expression: TypedExpr, + ) -> Result { + self.lower_typed_expr(expression)?.into_compiled_fn() + } + + pub(super) fn compile_typed_expr_to_assembly( + self, + expression: TypedExpr, + ) -> Result { + self.lower_typed_expr(expression)?.into_assembly() + } + + fn lower_typed_expr(self, expression: TypedExpr) -> Result { + let CompileFnBuilder { input_vars, .. } = self; + let expression = Box::new(expression); + + let mut jit_builder = make_jit_builder()?; + register_jit_symbols(&mut jit_builder); + let mut module = JITModule::new(jit_builder); + let target_config = module.target_config(); + let pointer_type = target_config.pointer_type(); + + // The native entry point mirrors JitEntry: its arguments point to the + // input slots and call-scoped string arena. VariableValue is returned as + // two integer-class values according to the native C ABI. + let mut signature = module.make_signature(); + signature.params.push(AbiParam::new(pointer_type)); + signature.params.push(AbiParam::new(pointer_type)); + signature.returns.push(AbiParam::new(types::I64)); + signature.returns.push(AbiParam::new(types::I64)); + let function_id = module.declare_anonymous_function(&signature)?; + + let mut context = module.make_context(); + context.func.signature = signature; + context.func.name = UserFuncName::user(0, function_id.as_u32()); + let native_functions = + declare_native_functions(&mut module, &mut context.func, pointer_type)?; + + let mut function_builder_context = FunctionBuilderContext::new(); + { + let mut builder = + FunctionBuilder::new(&mut context.func, &mut function_builder_context); + let entry_block = builder.create_block(); + builder.append_block_params_for_function_params(entry_block); + builder.switch_to_block(entry_block); + builder.seal_block(entry_block); + + let args_ptr = builder.block_params(entry_block)[0]; + let string_arena_ptr = builder.block_params(entry_block)[1]; + let mut lowering_context = LoweringContext { + args_ptr, + string_arena_ptr, + string_arena_was_reset: false, + pointer_type, + native_functions: &native_functions, + }; + let lowered = lowering_context.compile_expr(&expression, &mut builder)?; + let value_bits = match expression.return_type { + VarType::Bool => builder.ins().uextend(types::I64, lowered.value), + VarType::F64 => { + builder + .ins() + .bitcast(types::I64, MemFlagsData::new(), lowered.value) + } + VarType::U64 | VarType::I64 | VarType::Str | VarType::None => lowered.value, + }; + let second_word = if expression.return_type == VarType::Str { + lowered.string_len + } else { + builder.ins().uextend(types::I64, lowered.is_present) + }; + builder.ins().return_(&[value_bits, second_word]); + builder.finalize(target_config); + } + + Ok(LoweredFunction { + module, + context, + function_id, + input_vars, + expression, + }) + } +} + +fn preferred_numerical_type(inferred_types: InferredTypeSet) -> Option { + if inferred_types.i64 { + Some(VarType::I64) + } else if inferred_types.u64 { + Some(VarType::U64) + } else if inferred_types.f64 { + Some(VarType::F64) + } else { + None + } +} + +fn is_numerical(var_type: VarType) -> bool { + matches!(var_type, VarType::I64 | VarType::U64 | VarType::F64) +} + +impl LoweredFunction { + fn into_compiled_fn(self) -> Result { + let LoweredFunction { + mut module, + mut context, + function_id, + input_vars, + expression, + } = self; + + module.define_function(function_id, &mut context)?; + module.finalize_definitions()?; + + let code = module.get_finalized_function(function_id); + // SAFETY: `code` is the finalized entry point for the function whose ABI + // was built above to exactly match `JitEntry`. The module is retained by + // `CompiledFn`, so its executable allocation outlives `entry`. + let entry = unsafe { mem::transmute::<*const u8, JitEntry>(code) }; + Ok(CompiledFn { + entry, + _module: module, + inputs: input_vars, + _typed_expr: expression, + }) + } + + fn into_assembly(mut self) -> Result { + self.context.set_disasm(true); + let compiled_code = self + .context + .compile(self.module.isa(), &mut ControlPlane::default()) + .map_err(ModuleError::from)?; + Ok(compiled_code + .vcode + .clone() + .expect("Cranelift assembly was requested before compilation")) + } +} + +/// Converts an `f64` to an `i64` only when the value can be represented exactly. +fn f64_to_i64_lossless(value: f64) -> Option { + let is_integral = value.is_finite() && value.fract() == 0.0; + if is_integral && value >= i64::MIN as f64 && value < -(i64::MIN as f64) { + Some(value as i64) + } else { + None + } +} + +/// Converts an `f64` to a `u64` only when the value can be represented exactly. +fn f64_to_u64_lossless(value: f64) -> Option { + let is_integral = value.is_finite() && value.fract() == 0.0; + if is_integral && value >= 0.0 && value < u64::MAX as f64 { + Some(value as u64) + } else { + None + } +} + +fn assign_variable_ids(expr: &mut TypedExpr) -> Vec { + let mut name_to_vars: HashMap, TypedVariable> = HashMap::new(); + assign_variable_ids_aux(&mut expr.ast, &mut name_to_vars); + let mut input_vars: Vec = name_to_vars.into_values().collect(); + input_vars.sort_by_key(|var| var.variable_id); + input_vars +} + +fn assign_variable_ids_aux( + ast: &mut TypedExprAst, + name_to_vars: &mut HashMap, TypedVariable>, +) { + match ast { + TypedExprAst::Literal(_) => {} + TypedExprAst::Variable(var) => { + if let Some(typed_var) = name_to_vars.get(&var.variable_name) { + assert_eq!( + typed_var.r#type, var.r#type, + "variable `{}` appears with two different types (`{:?}` and `{:?}`); a typed \ + expr AST must be built with a single explicit type per variable", + var.variable_name, typed_var.r#type, var.r#type, + ); + var.variable_id = typed_var.variable_id; + } else { + var.variable_id = name_to_vars.len(); + name_to_vars.insert(var.variable_name.clone(), var.clone()); + }; + } + TypedExprAst::Coerce { expr, .. } => { + assign_variable_ids_aux(&mut expr.ast, name_to_vars); + } + TypedExprAst::FnCall(fn_call) => { + for arg in fn_call.args_mut() { + assign_variable_ids_aux(&mut arg.ast, name_to_vars); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::ast::Function; + use crate::functions::{AddFnCall, FnCallEnum}; + + #[test] + fn test_apply_types_to_literal() { + let untyped_expr = UntypedExpr::literal("hello"); + let variable_types = HashMap::new(); + let mut builder = CompileFnBuilder::new(&variable_types); + let typed_expr = builder.build_typed_expr(&untyped_expr).unwrap(); + + assert_eq!(typed_expr.return_type, VarType::Str); + let TypedExprAst::Literal(TypedLiteral::String(value)) = typed_expr.ast else { + panic!("expected a typed string literal"); + }; + assert_eq!(value.as_ref(), "hello"); + } + + #[test] + fn test_assign_variable_ids_two_variables_different_types() { + let untyped_expr = Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("x"), UntypedExpr::variable("y")]); + let variable_types = HashMap::from([("x", VarType::U64), ("y", VarType::F64)]); + let mut builder = CompileFnBuilder::new(&variable_types); + + let _typed_expr = builder.build_typed_expr(&untyped_expr).unwrap(); + let var_args = &builder.input_vars; + + assert_eq!(var_args.len(), 2); + assert_eq!(var_args[0].variable_name.as_ref(), "x"); + assert_eq!(var_args[0].r#type, VarType::U64); + assert_eq!(var_args[0].variable_id, 0); + assert_eq!(var_args[1].variable_name.as_ref(), "y"); + assert_eq!(var_args[1].r#type, VarType::F64); + assert_eq!(var_args[1].variable_id, 1); + } + + #[test] + #[should_panic(expected = "appears with two different types")] + fn test_assign_variable_ids_panics_on_inconsistent_types() { + let mut typed_expr = TypedExpr { + return_type: VarType::F64, + ast: TypedExprAst::FnCall(FnCallEnum::Add(AddFnCall { + args: vec![ + TypedExprAst::variable("x", VarType::U64).with_type(VarType::U64), + TypedExprAst::variable("x", VarType::F64).with_type(VarType::F64), + ] + .into_boxed_slice(), + })), + }; + + assign_variable_ids(&mut typed_expr); + } + + #[test] + fn test_assign_variable_ids_dedups_repeated_variable() { + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("x"), + Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("y"), UntypedExpr::variable("x")]), + ]); + let variable_types: HashMap<&str, VarType> = + HashMap::from([("x", VarType::U64), ("y", VarType::U64)]); + let mut builder = CompileFnBuilder::new(&variable_types); + + let _typed_expr = builder.build_typed_expr(&untyped_expr).unwrap(); + let var_args = &builder.input_vars; + + assert_eq!(var_args.len(), 2); + assert_eq!(var_args[0].variable_name.as_ref(), "x"); + assert_eq!(var_args[0].r#type, VarType::U64); + assert_eq!(var_args[0].variable_id, 0); + assert_eq!(var_args[1].variable_name.as_ref(), "y"); + assert_eq!(var_args[1].r#type, VarType::U64); + assert_eq!(var_args[1].variable_id, 1); + } + + #[test] + fn test_jit_entry_returns_variable_value_as_two_abi_words() { + let variable_types = HashMap::new(); + let mut builder = CompileFnBuilder::new(&variable_types); + let expression = builder + .build_typed_expr(&UntypedExpr::literal(1u64)) + .unwrap(); + + let lowered = builder.lower_typed_expr(expression).unwrap(); + let signature = &lowered.context.func.signature; + + assert_eq!(signature.params.len(), 2); + assert!( + signature + .params + .iter() + .all(|param| param.value_type == types::I64) + ); + assert_eq!(signature.returns.len(), 2); + assert_eq!(signature.returns[0].value_type, types::I64); + assert_eq!(signature.returns[1].value_type, types::I64); + } +} diff --git a/jitexpr/src/compile/compiled_fn.rs b/jitexpr/src/compile/compiled_fn.rs new file mode 100644 index 0000000000..2f0561110f --- /dev/null +++ b/jitexpr/src/compile/compiled_fn.rs @@ -0,0 +1,154 @@ +use std::ops::Deref; +use std::sync::Arc; + +use cranelift_jit::JITModule; + +use super::{StringArena, TypedExpr, TypedVariable}; +use crate::types::{VarType, VariableValue}; + +#[cfg(not(any( + all(target_arch = "x86_64", not(target_os = "windows")), + target_arch = "aarch64" +)))] +// Windows is not supported because apparently returning more than one 64 bits word throught +// registers is not supported by its ABI. +compile_error!( + "the direct VariableValue JIT return ABI is only implemented for x86-64 System V and AArch64" +); + +// On the supported targets, VariableValue's two eightbytes are returned in two +// integer registers by the platform C ABI. The lifetime is selected by +// CompiledFn::call so that it is bounded by all possible sources of strings. +// This is a Rust-to-JIT boundary whose VariableValue layout is asserted in +// types.rs, not an interface intended for C callers. +#[allow(improper_ctypes_definitions)] +pub(crate) type JitEntry = + for<'a> unsafe extern "C" fn(*const VariableValue<'a>, *mut StringArena) -> VariableValue<'a>; + +/// An expression compiled to native machine code. +/// +/// This object owns the JIT module containing its executable memory and every +/// resource referenced by the generated code. +pub struct CompiledFn { + pub(crate) entry: JitEntry, + pub(crate) _module: JITModule, + /// Input slots in the exact order expected by [`CompiledFn::call`]. + pub inputs: Vec, + // This AST owns the Arc-backed literals and regexes embedded in generated code. + pub(crate) _typed_expr: Box, +} + +// `JITModule` is not `Sync` because it supports lazily looking up symbols through +// interior mutability. A `CompiledFn` only retains a finalized module to keep its +// executable allocation alive and never invokes those mutable APIs. Its entry +// point and the immutable resources referenced by the generated code can be +// called concurrently when each caller supplies a distinct `StringArena`. +unsafe impl Sync for CompiledFn {} + +impl CompiledFn { + /// Returns the concrete result type selected during compilation. + pub fn result_type(&self) -> VarType { + self._typed_expr.return_type + } + + /// Creates an evaluation context with a private string arena. + pub fn context(self: &Arc) -> CompiledFnCtx { + CompiledFnCtx::new(Arc::clone(self)) + } + + /// Evaluates the compiled expression using the supplied string arena. + /// + /// The mutable arena borrow prevents another evaluation from clearing the + /// arena while an arena-backed result from this call is still live. + /// + /// # Safety + /// + /// `args` must follow [`CompiledFn::inputs`] exactly: every present slot must + /// contain the union member corresponding to that variable's type. Absent + /// slots must use [`VariableValue::none`], and any borrowed strings must + /// remain alive for the duration of this call. + /// + /// The result cannot outlive the compiled function, the string arena, or + /// the passed arguments' lifetime. + #[inline(always)] + pub unsafe fn call<'args, 'compiled, 'arena, 'output>( + &'compiled self, + args: &[VariableValue<'args>], + string_arena: &'arena mut StringArena, + ) -> VariableValue<'output> + where + 'args: 'output, + 'compiled: 'output, + 'arena: 'output, + { + debug_assert_eq!(args.len(), self.inputs.len()); + let args: &[VariableValue<'output>] = args; + let string_arena = &raw mut *string_arena; + // SAFETY: Guaranteed by the caller. The input, compiled-function, and + // arena lifetimes outlive the lifetime selected for the returned value. + unsafe { (self.entry)(args.as_ptr(), string_arena) } + } +} + +/// Per-caller mutable state used to evaluate a shared [`CompiledFn`]. +pub struct CompiledFnCtx { + compiled_fn: Arc, + pub(crate) string_arena: StringArena, +} + +impl CompiledFnCtx { + /// Creates an evaluation context for `compiled_fn`. + pub fn new(compiled_fn: Arc) -> Self { + Self { + compiled_fn, + string_arena: StringArena::new(), + } + } + + /// Evaluates the compiled expression using this context's string arena. + /// + /// The mutable borrow prevents another evaluation from clearing the string + /// arena while an arena-backed result from this call is still live. + /// + /// # Safety + /// + /// `args` must follow [`CompiledFn::inputs`] exactly: every present slot must + /// contain the union member corresponding to that variable's type. Absent + /// slots must use [`VariableValue::none`], and any borrowed strings must + /// remain alive for the duration of this call. + /// + /// The result cannot outlive this context nor the passed arguments' + /// lifetime. + #[inline(always)] + pub unsafe fn call<'args, 'ctx, 'output>( + &'ctx mut self, + args: &[VariableValue<'args>], + ) -> VariableValue<'output> + where + 'args: 'output, + 'ctx: 'output, + { + // SAFETY: Guaranteed by the caller. The context owns both the compiled + // function and arena for the lifetime selected for the returned value. + unsafe { self.compiled_fn.call(args, &mut self.string_arena) } + } + + /// Returns the shared compiled expression owned by this context. + pub fn compiled_fn(&self) -> &Arc { + &self.compiled_fn + } +} + +impl From> for CompiledFnCtx { + fn from(compiled_fn: Arc) -> Self { + Self::new(compiled_fn) + } +} + +impl Deref for CompiledFnCtx { + type Target = CompiledFn; + + fn deref(&self) -> &Self::Target { + &self.compiled_fn + } +} diff --git a/jitexpr/src/compile/error.rs b/jitexpr/src/compile/error.rs new file mode 100644 index 0000000000..c63d2da592 --- /dev/null +++ b/jitexpr/src/compile/error.rs @@ -0,0 +1,31 @@ +use crate::ast::{Function, InvalidFunctionCall, TypeError}; +use crate::types::VarType; + +#[derive(Debug, thiserror::Error)] +pub enum CompileError { + #[error("type inference failed: {0}")] + TypeInference(#[from] TypeError), + #[error("JIT compilation failed: {0}")] + Module(#[source] Box), + #[error("cannot coerce an expression from {from_type:?} to {target:?}")] + UnsupportedCoercion { from_type: VarType, target: VarType }, + #[error("cannot compile {function:?} with result type {return_type:?}")] + UnsupportedFunctionType { + function: Function, + return_type: VarType, + }, + #[error("invalid regular expression `{pattern}`: {source}")] + InvalidRegex { + pattern: String, + #[source] + source: regex::Error, + }, + #[error("arguments do not match the function {0}")] + InvalidArguments(#[from] InvalidFunctionCall), +} + +impl From for CompileError { + fn from(error: cranelift_module::ModuleError) -> Self { + CompileError::Module(Box::new(error)) + } +} diff --git a/jitexpr/src/compile/mod.rs b/jitexpr/src/compile/mod.rs new file mode 100644 index 0000000000..de4e8cc441 --- /dev/null +++ b/jitexpr/src/compile/mod.rs @@ -0,0 +1,362 @@ +mod compile_fn_builder; +mod compiled_fn; +mod error; +mod string_arena; +mod typed_expr; +mod typed_expr_serialize; + +use std::collections::HashMap; +use std::sync::Arc; + +pub(crate) use compile_fn_builder::CompileFnBuilder; +pub use compiled_fn::{CompiledFn, CompiledFnCtx}; +use cranelift::codegen::ir::{ + InstBuilder as _, MemFlagsData, Type, Value as CraneliftValue, types as cranelift_types, +}; +use cranelift::frontend::FunctionBuilder; +pub use error::CompileError; +#[cfg(test)] +pub(crate) use string_arena::STRING_ARENA_CAPACITY; +pub use string_arena::StringArena; +pub use typed_expr::TypedVariable; +pub(crate) use typed_expr::{TypedExpr, TypedExprAst, TypedLiteral}; +pub(crate) use typed_expr_serialize::{format_function_call, format_string_literal}; + +use crate::ast::UntypedExpr; +use crate::functions::NativeFunctions; +use crate::types::{VarType, VariablePrimitiveOpt, VariableValue}; + +/// Compiles an expression into an immutable, shareable native function. +pub fn compile( + untyped_expr: &UntypedExpr, + var_types: &HashMap<&str, VarType>, +) -> Result, CompileError> { + let mut builder = CompileFnBuilder::new(var_types); + let typed_expr = builder.build_typed_expr(untyped_expr)?; + builder.compile_typed_expr(typed_expr).map(Arc::new) +} + +/// Applies concrete variable types and serializes the resulting typed expression. +pub fn serialize( + untyped_expr: &UntypedExpr, + var_types: &HashMap<&str, VarType>, +) -> Result { + let mut builder = CompileFnBuilder::new(var_types); + let typed_expr = builder.build_typed_expr(untyped_expr)?; + Ok(typed_expr_serialize::serialize(&typed_expr)) +} + +/// Compiles an expression and returns Cranelift's assembly listing for the host target. +pub fn compile_to_assembly( + untyped_expr: &UntypedExpr, + var_types: &HashMap<&str, VarType>, +) -> Result { + let mut builder = CompileFnBuilder::new(var_types); + let typed_expr = builder.build_typed_expr(untyped_expr)?; + builder.compile_typed_expr_to_assembly(typed_expr) +} + +pub(crate) struct LoweringContext<'a> { + args_ptr: CraneliftValue, + string_arena_ptr: CraneliftValue, + string_arena_was_reset: bool, + pointer_type: Type, + native_functions: &'a NativeFunctions, +} + +/// The two SSA values used to represent a nullable expression result. +#[derive(Clone, Copy)] +pub(crate) struct LoweredValue { + pub(crate) value: CraneliftValue, + pub(crate) is_present: CraneliftValue, + pub(crate) string_len: CraneliftValue, +} + +impl LoweringContext<'_> { + pub(crate) fn compile_expr( + &mut self, + expression: &TypedExpr, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + match &expression.ast { + TypedExprAst::Literal(literal) => Ok(lower_literal(literal, self, builder)), + TypedExprAst::Variable(variable) => { + let slot_offset = variable.variable_id * std::mem::size_of::(); + let value_offset = slot_offset as i32; + let second_word_offset = + (slot_offset + std::mem::offset_of!(VariablePrimitiveOpt, is_present)) as i32; + let value = builder.ins().load( + cranelift_type(variable.r#type, self.pointer_type), + MemFlagsData::trusted(), + self.args_ptr, + value_offset, + ); + let (is_present, string_len) = if variable.r#type == VarType::Str { + let string_len = builder.ins().load( + cranelift_types::I64, + MemFlagsData::trusted(), + self.args_ptr, + second_word_offset, + ); + let is_present = + builder + .ins() + .icmp_imm_u(cranelift::prelude::IntCC::NotEqual, value, 0); + (is_present, string_len) + } else { + let is_present = builder.ins().load( + cranelift_types::I8, + MemFlagsData::trusted(), + self.args_ptr, + second_word_offset, + ); + let string_len = builder.ins().iconst(cranelift_types::I64, 0); + (is_present, string_len) + }; + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } + TypedExprAst::Coerce { target_type, expr } => { + let source_type = expr.return_type; + let lowered = self.compile_expr(expr, builder)?; + let value = lower_coercion(lowered.value, source_type, *target_type, builder)?; + Ok(LoweredValue { + value, + is_present: lowered.is_present, + string_len: lowered.string_len, + }) + } + TypedExprAst::FnCall(fn_call) => fn_call.lower(expression.return_type, self, builder), + } + } + + pub(crate) fn pointer_type(&self) -> Type { + self.pointer_type + } + + pub(crate) fn string_arena_ptr(&mut self, builder: &mut FunctionBuilder<'_>) -> CraneliftValue { + if !self.string_arena_was_reset { + let zero = builder.ins().iconst(cranelift_types::I64, 0); + builder.ins().store( + MemFlagsData::trusted(), + zero, + self.string_arena_ptr, + StringArena::CURSOR_OFFSET, + ); + self.string_arena_was_reset = true; + } + self.string_arena_ptr + } + + pub(crate) fn native_functions(&self) -> &NativeFunctions { + self.native_functions + } +} + +fn lower_literal( + literal: &TypedLiteral, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, +) -> LoweredValue { + let value = match literal { + TypedLiteral::None => builder.ins().iconst(cranelift_types::I64, 0), + TypedLiteral::Bool(value) => builder.ins().iconst(cranelift_types::I8, i64::from(*value)), + TypedLiteral::U64(value) => builder.ins().iconst(cranelift_types::I64, *value as i64), + TypedLiteral::I64(value) => builder.ins().iconst(cranelift_types::I64, *value), + TypedLiteral::F64(value) => { + builder + .ins() + .f64const(cranelift::codegen::ir::immediates::Ieee64::with_bits( + value.to_bits(), + )) + } + TypedLiteral::String(value) => { + let string_ptr = value.as_ptr() as usize; + builder + .ins() + .iconst(context.pointer_type, string_ptr as i64) + } + }; + let is_present = builder.ins().iconst( + cranelift_types::I8, + i64::from(!matches!(literal, TypedLiteral::None)), + ); + let string_len = match literal { + TypedLiteral::String(value) => builder + .ins() + .iconst(cranelift_types::I64, value.len() as i64), + _ => builder.ins().iconst(cranelift_types::I64, 0), + }; + LoweredValue { + value, + is_present, + string_len, + } +} + +fn lower_coercion( + value: CraneliftValue, + source: VarType, + target: VarType, + builder: &mut FunctionBuilder<'_>, +) -> Result { + let coerced = match (source, target) { + (source, target) if source == target => value, + (VarType::U64, VarType::F64) => builder.ins().fcvt_from_uint(cranelift_types::F64, value), + (VarType::I64, VarType::F64) => builder.ins().fcvt_from_sint(cranelift_types::F64, value), + // Cranelift integers do not carry signedness. These two coercions have + // the same machine representation and therefore need no instruction. + (VarType::U64, VarType::I64) | (VarType::I64, VarType::U64) => value, + _ => { + return Err(CompileError::UnsupportedCoercion { + from_type: source, + target, + }); + } + }; + Ok(coerced) +} + +fn cranelift_type(var_type: VarType, pointer_type: Type) -> Type { + match var_type { + VarType::Bool => cranelift_types::I8, + VarType::F64 => cranelift_types::F64, + VarType::U64 | VarType::I64 | VarType::None => cranelift_types::I64, + VarType::Str => pointer_type, + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::ast::{Function, UntypedExpr}; + use crate::types::{VarType, VariableValue}; + + #[test] + fn test_compile_bool_variable() { + let untyped_expr = UntypedExpr::variable("flag"); + let variable_types = HashMap::from([("flag", VarType::Bool)]); + let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); + let mut string_arena = StringArena::new(); + assert!(string_arena.allocate(1).is_some()); + let input = [VariableValue::some(true)]; + let output = unsafe { compiled_fn.call(&input, &mut string_arena) }; + + assert_eq!(unsafe { output.as_bool() }, Some(true)); + assert_eq!(string_arena.used_bytes(), 1); + } + + #[test] + fn test_compiled_fn_is_send_and_sync() { + fn assert_send_and_sync() {} + + assert_send_and_sync::(); + } + + #[test] + fn test_contexts_have_independent_string_arenas() { + let untyped_expr = Function::Lower.call_untyped_expr(vec![UntypedExpr::variable("value")]); + let variable_types = HashMap::from([("value", VarType::Str)]); + let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); + let mut first_ctx = compiled_fn.context(); + let mut second_ctx = CompiledFnCtx::from(compiled_fn); + + let first = unsafe { first_ctx.call(&[VariableValue::some("FIRST")]) }; + assert_eq!(unsafe { first.as_str() }, Some("first")); + assert_eq!(first_ctx.string_arena.used_bytes(), 5); + + let second = unsafe { second_ctx.call(&[VariableValue::some("SECOND")]) }; + assert_eq!(unsafe { second.as_str() }, Some("second")); + assert_eq!(second_ctx.string_arena.used_bytes(), 6); + assert_eq!(first_ctx.string_arena.used_bytes(), 5); + } + + #[test] + fn test_compile_none_variable() { + let untyped_expr = UntypedExpr::variable("value"); + let variable_types = HashMap::from([("value", VarType::U64)]); + let mut compiled_fn = compile(&untyped_expr, &variable_types).unwrap().context(); + let input = [VariableValue::none()]; + assert_eq!(compiled_fn.result_type(), VarType::U64); + let output = unsafe { compiled_fn.call(&input) }; + + assert_eq!(unsafe { output.as_u64() }, None); + } + + #[test] + fn test_compile_string_literal_keeps_backing_data_alive() { + let untyped_expr = UntypedExpr::literal("hello"); + let mut compiled_fn = compile(&untyped_expr, &HashMap::new()).unwrap().context(); + drop(untyped_expr); + let output = unsafe { compiled_fn.call(&[]) }; + + assert_eq!(unsafe { output.as_str() }, Some("hello")); + } + + #[test] + fn test_compile_returns_borrowed_string_variable_directly() { + let untyped_expr = UntypedExpr::variable("value"); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled_fn = compile(&untyped_expr, &variable_types).unwrap().context(); + let value = String::from("hello from an input"); + let input = [VariableValue::some(value.as_str())]; + let output = unsafe { compiled_fn.call(&input) }; + + let output = unsafe { output.as_str() }.unwrap(); + assert_eq!(output, value); + assert_eq!(output.as_ptr(), value.as_ptr()); + } + + #[test] + fn test_compile_none_literal_returns_absent_value() { + let untyped_expr = UntypedExpr::literal(crate::ast::Literal::None); + let mut compiled_fn = compile(&untyped_expr, &HashMap::new()).unwrap().context(); + assert_eq!(compiled_fn.result_type(), VarType::None); + let output = unsafe { compiled_fn.call(&[]) }; + + assert_eq!(unsafe { output.as_u64() }, None); + } + + #[test] + fn test_compile_to_assembly() { + let untyped_expr = UntypedExpr::variable("value"); + let variable_types = HashMap::from([("value", VarType::F64)]); + + let assembly = compile_to_assembly(&untyped_expr, &variable_types).unwrap(); + + assert!(assembly.contains("block0:")); + assert!(!assembly.trim().is_empty()); + } + + #[test] + fn test_compile_native_call_to_assembly() { + let untyped_expr = Function::RegexpExtract.call_untyped_expr(vec![ + UntypedExpr::variable("message"), + UntypedExpr::literal("([a-z]+)"), + UntypedExpr::literal(0u64), + ]); + let variable_types = HashMap::from([("message", VarType::Str)]); + + let assembly = compile_to_assembly(&untyped_expr, &variable_types).unwrap(); + + assert!(assembly.contains("block0:")); + assert!(!assembly.trim().is_empty()); + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn test_compile_to_assembly_does_not_sign_return_address() { + let untyped_expr = Function::Lower.call_untyped_expr(vec![UntypedExpr::variable("value")]); + let variable_types = HashMap::from([("value", VarType::Str)]); + + let assembly = compile_to_assembly(&untyped_expr, &variable_types).unwrap(); + + assert!(!assembly.contains("pacibsp")); + assert!(!assembly.contains("retabsp")); + } +} diff --git a/jitexpr/src/compile/string_arena.rs b/jitexpr/src/compile/string_arena.rs new file mode 100644 index 0000000000..bab3621d80 --- /dev/null +++ b/jitexpr/src/compile/string_arena.rs @@ -0,0 +1,43 @@ +pub(crate) const STRING_ARENA_CAPACITY: usize = 262_144; + +/// Fixed-capacity storage for strings constructed while evaluating an expression. +pub struct StringArena { + buffer: Box<[u8; STRING_ARENA_CAPACITY]>, + cursor: usize, +} + +impl StringArena { + pub(crate) const CURSOR_OFFSET: i32 = std::mem::offset_of!(StringArena, cursor) as i32; + + /// Creates an empty string arena. + pub fn new() -> Self { + let buffer = vec![0; STRING_ARENA_CAPACITY].into_boxed_slice(); + let buffer = buffer + .try_into() + .unwrap_or_else(|_| unreachable!("the arena buffer has the requested capacity")); + Self { buffer, cursor: 0 } + } + + /// Reserves `len` contiguous bytes without growing the backing allocation. + pub(crate) fn allocate(&mut self, len: usize) -> Option<*mut u8> { + let end = self.cursor.checked_add(len)?; + if end > STRING_ARENA_CAPACITY { + return None; + } + // SAFETY: `cursor <= end <= STRING_ARENA_CAPACITY`. + let allocation = unsafe { self.buffer.as_mut_ptr().add(self.cursor) }; + self.cursor = end; + Some(allocation) + } + + #[cfg(test)] + pub(crate) fn used_bytes(&self) -> usize { + self.cursor + } +} + +impl Default for StringArena { + fn default() -> Self { + Self::new() + } +} diff --git a/jitexpr/src/compile/typed_expr.rs b/jitexpr/src/compile/typed_expr.rs new file mode 100644 index 0000000000..376d2816a2 --- /dev/null +++ b/jitexpr/src/compile/typed_expr.rs @@ -0,0 +1,177 @@ +use std::sync::Arc; + +#[cfg(test)] +use crate::ast::Literal; +use crate::functions::FnCallEnum; +use crate::types::VarType; + +#[derive(Clone, PartialEq)] +pub struct TypedVariable { + /// The source-level variable name. + pub variable_name: Arc, + /// The concrete type expected in this input slot. + pub r#type: VarType, + /// The position in the compiled input array. + pub variable_id: usize, +} + +#[derive(Clone, PartialEq)] +pub(crate) struct TypedExpr { + pub(crate) return_type: VarType, + pub(crate) ast: TypedExprAst, +} + +impl TypedExpr { + pub(crate) fn coerce(self, target_type: VarType) -> TypedExpr { + if target_type == self.return_type { + self + } else { + let TypedExpr { return_type, ast } = self; + let ast = match (ast, target_type) { + (TypedExprAst::Literal(TypedLiteral::U64(value)), VarType::I64) + if value <= i64::MAX as u64 => + { + TypedExprAst::Literal(TypedLiteral::I64(value as i64)) + } + (TypedExprAst::Literal(TypedLiteral::U64(value)), VarType::F64) => { + TypedExprAst::Literal(TypedLiteral::F64(value as f64)) + } + (TypedExprAst::Literal(TypedLiteral::I64(value)), VarType::U64) if value >= 0 => { + TypedExprAst::Literal(TypedLiteral::U64(value as u64)) + } + (TypedExprAst::Literal(TypedLiteral::I64(value)), VarType::F64) => { + TypedExprAst::Literal(TypedLiteral::F64(value as f64)) + } + (TypedExprAst::Literal(TypedLiteral::F64(value)), VarType::U64) + if value.is_finite() + && value.fract() == 0.0 + && value >= 0.0 + && value < u64::MAX as f64 => + { + TypedExprAst::Literal(TypedLiteral::U64(value as u64)) + } + (TypedExprAst::Literal(TypedLiteral::F64(value)), VarType::I64) + if value.is_finite() + && value.fract() == 0.0 + && value >= i64::MIN as f64 + && value < -(i64::MIN as f64) => + { + TypedExprAst::Literal(TypedLiteral::I64(value as i64)) + } + (ast, target_type) => TypedExprAst::Coerce { + target_type, + expr: Box::new(TypedExpr { return_type, ast }), + }, + }; + TypedExpr { + return_type: target_type, + ast, + } + } + } + + pub(crate) fn none() -> TypedExpr { + TypedExpr { + return_type: VarType::None, + ast: TypedExprAst::Literal(TypedLiteral::None), + } + } + + #[cfg(test)] + pub(crate) fn literal(val: impl Into) -> TypedExpr { + let literal: Literal = val.into(); + let r#type = literal.r#type(); + let literal = match literal { + Literal::None => TypedLiteral::None, + Literal::Bool(value) => TypedLiteral::Bool(value), + Literal::U64(value) => TypedLiteral::U64(value), + Literal::I64(value) => TypedLiteral::I64(value), + Literal::F64(value) => TypedLiteral::F64(value), + Literal::String(_) => panic!("typed string literals require registered backing data"), + }; + TypedExprAst::Literal(literal).with_type(r#type) + } +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum TypedLiteral { + None, + Bool(bool), + U64(u64), + I64(i64), + F64(f64), + String(Arc), +} + +impl TypedLiteral { + pub(crate) fn r#type(&self) -> VarType { + match self { + TypedLiteral::None => VarType::None, + TypedLiteral::Bool(_) => VarType::Bool, + TypedLiteral::U64(_) => VarType::U64, + TypedLiteral::I64(_) => VarType::I64, + TypedLiteral::F64(_) => VarType::F64, + TypedLiteral::String(_) => VarType::Str, + } + } +} + +#[derive(Clone, PartialEq)] +pub(crate) enum TypedExprAst { + Literal(TypedLiteral), + Variable(TypedVariable), + Coerce { + target_type: VarType, + expr: Box, + }, + FnCall(FnCallEnum), +} + +impl TypedExprAst { + #[cfg(test)] + pub(crate) fn with_type(self, return_type: VarType) -> TypedExpr { + TypedExpr { + return_type, + ast: self, + } + } + + pub(crate) fn variable(variable_name: impl ToString, r#type: VarType) -> TypedExprAst { + TypedExprAst::Variable(TypedVariable { + variable_name: Arc::from(variable_name.to_string()), + r#type, + variable_id: 0, + }) + } + + pub(crate) fn from_call(fn_call: impl Into) -> TypedExprAst { + TypedExprAst::FnCall(fn_call.into()) + } +} + +impl std::fmt::Debug for TypedExpr { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "({:?} : {:?})", self.ast, self.return_type) + } +} + +impl std::fmt::Debug for TypedExprAst { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + TypedExprAst::Literal(literal) => write!(f, "{:?}", literal), + TypedExprAst::Variable(variable) => write!(f, "{:?}", variable), + TypedExprAst::Coerce { target_type, expr } => { + write!(f, "coerce({:?} as {:?})", expr, target_type) + } + TypedExprAst::FnCall(fn_call) => { + write!(f, "{fn_call:?}") + } + } + } +} + +impl std::fmt::Debug for TypedVariable { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{{{}:{:?}}}", self.variable_name, self.r#type) + } +} diff --git a/jitexpr/src/compile/typed_expr_serialize.rs b/jitexpr/src/compile/typed_expr_serialize.rs new file mode 100644 index 0000000000..2774c996cb --- /dev/null +++ b/jitexpr/src/compile/typed_expr_serialize.rs @@ -0,0 +1,214 @@ +//! Serialization for the normalized typed expression tree. +//! +//! Calls, variables, and coercions use `[type: expression]`, while literals retain the canonical +//! untyped literal syntax because numerical suffixes and literal spellings already identify their +//! types. For example: +//! +//! ```text +//! [int64: ADD 1i64 [int64: my_col]] +//! ``` + +use std::fmt; + +use super::{TypedExpr, TypedExprAst, TypedLiteral}; +use crate::types::VarType; + +pub(super) fn serialize(expr: &TypedExpr) -> String { + expr.to_string() +} + +impl fmt::Display for TypedExpr { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + format_expr(self, formatter) + } +} + +fn format_expr(expr: &TypedExpr, formatter: &mut fmt::Formatter) -> fmt::Result { + if let TypedExprAst::Literal(literal) = &expr.ast { + return format_literal(literal, formatter); + } + + write!(formatter, "[{}: ", type_name(expr.return_type))?; + match &expr.ast { + TypedExprAst::Literal(_) => unreachable!(), + TypedExprAst::Variable(variable) => formatter.write_str(&variable.variable_name)?, + TypedExprAst::Coerce { expr, .. } => write!(formatter, "COERCE {expr}")?, + TypedExprAst::FnCall(fn_call) => fn_call.serialize(formatter)?, + } + formatter.write_str("]") +} + +fn format_literal(literal: &TypedLiteral, formatter: &mut fmt::Formatter) -> fmt::Result { + match literal { + TypedLiteral::None => formatter.write_str("none"), + TypedLiteral::Bool(value) => write!(formatter, "{value}"), + TypedLiteral::U64(value) => write!(formatter, "{value}u64"), + TypedLiteral::I64(value) => write!(formatter, "{value}i64"), + TypedLiteral::F64(value) => write!(formatter, "{value}f64"), + TypedLiteral::String(value) => format_string_literal(value, formatter), + } +} + +pub(crate) fn format_function_call<'a>( + name: &str, + args: impl IntoIterator, + formatter: &mut fmt::Formatter, +) -> fmt::Result { + formatter.write_str(name)?; + for arg in args { + write!(formatter, " {arg}")?; + } + Ok(()) +} + +pub(crate) fn format_string_literal(value: &str, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("\"")?; + for character in value.chars() { + match character { + '"' => formatter.write_str("\\\""), + '\\' => formatter.write_str("\\\\"), + '\n' => formatter.write_str("\\n"), + '\r' => formatter.write_str("\\r"), + '\t' => formatter.write_str("\\t"), + '\0' => formatter.write_str("\\0"), + character if character.is_control() => { + write!(formatter, "{}", character.escape_unicode()) + } + character => write!(formatter, "{character}"), + }?; + } + formatter.write_str("\"") +} + +fn type_name(var_type: VarType) -> &'static str { + match var_type { + VarType::Bool => "boolean", + VarType::F64 => "float64", + VarType::U64 => "uint64", + VarType::I64 => "int64", + VarType::Str => "string", + VarType::None => "none", + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use crate::types::VarType; + use crate::{ast, compile}; + + fn serialize(expression: &str, variable_types: &HashMap<&str, VarType>) -> String { + let expression = ast::deserialize(expression).unwrap(); + compile::serialize(&expression, variable_types).unwrap() + } + + #[test] + fn test_serializes_typed_add_expression() { + let variable_types = HashMap::from([("my_col", VarType::I64)]); + + assert_eq!( + serialize("(ADD 1i64 my_col)", &variable_types), + "[int64: ADD 1i64 [int64: my_col]]" + ); + } + + #[test] + fn test_serializes_explicit_coercion() { + let variable_types = HashMap::from([("my_col", VarType::I64)]); + + assert_eq!( + serialize("(ADD 1.5f64 my_col)", &variable_types), + "[float64: ADD 1.5f64 [float64: COERCE [int64: my_col]]]" + ); + } + + #[test] + fn test_serializes_all_variable_types_and_literals() { + let cases = [ + (VarType::Bool, "[boolean: value]"), + (VarType::F64, "[float64: value]"), + (VarType::U64, "[uint64: value]"), + (VarType::I64, "[int64: value]"), + (VarType::Str, "[string: value]"), + ]; + for (var_type, expected) in cases { + assert_eq!( + serialize("value", &HashMap::from([("value", var_type)])), + expected + ); + } + + let literals = [ + ("none", "none"), + ("true", "true"), + ("1u64", "1i64"), + ("18446744073709551615u64", "18446744073709551615u64"), + ("-2i64", "-2i64"), + ("1.5f64", "1.5f64"), + ( + r#""quoted: \"hello\"\\world\n\t\0\u{7} café""#, + r#""quoted: \"hello\"\\world\n\t\0\u{7} café""#, + ), + ]; + for (literal, expected) in literals { + assert_eq!(serialize(literal, &HashMap::new()), expected); + } + } + + #[test] + fn test_serializes_normalized_compile_time_arguments() { + let variable_types = HashMap::from([ + ("message", VarType::Str), + ("number", VarType::F64), + ("other", VarType::Str), + ]); + let cases = [ + ( + "(CONCAT \" / \" \"TRUE\" message other)", + "[string: CONCAT \" / \" \"true\" [string: message] [string: other]]", + ), + ( + "(LEFT message 2i64)", + "[string: LEFT [string: message] 2u64]", + ), + ( + r#"(REGEXP_EXTRACT message "([a-z]+)")"#, + r#"[string: REGEXP_EXTRACT [string: message] "([a-z]+)" 0u64]"#, + ), + ( + r#"(REGEXP_LIKE message "[a-z]+")"#, + r#"[boolean: REGEXP_LIKE [string: message] "[a-z]+"]"#, + ), + ( + "(RIGHT message 2i64)", + "[string: RIGHT [string: message] 2u64]", + ), + ("(ROUND number)", "[int64: ROUND [float64: number] 0i64]"), + ( + "(SPLIT_AFTER message \".\")", + "[string: SPLIT_AFTER [string: message] \".\" 0u64]", + ), + ( + "(SPLIT_BEFORE message \".\")", + "[string: SPLIT_BEFORE [string: message] \".\" 0u64]", + ), + ( + "(SUBSTRING message 1i64 2i64)", + "[string: SUBSTRING [string: message] 1u64 2u64]", + ), + ( + "(TEXT_JOIN \" / \" \"FALSE\" message other)", + "[string: TEXT_JOIN \" / \" \"false\" [string: message] [string: other]]", + ), + ( + "(TRIM message \"x\" \"BOTH\")", + "[string: TRIM [string: message] \"x\" \"both\"]", + ), + ]; + + for (expression, expected) in cases { + assert_eq!(serialize(expression, &variable_types), expected); + } + } +} diff --git a/jitexpr/src/functions/abs.rs b/jitexpr/src/functions/abs.rs new file mode 100644 index 0000000000..a8917fbc47 --- /dev/null +++ b/jitexpr/src/functions/abs.rs @@ -0,0 +1,230 @@ +//! `ABS` computes the absolute value of one numeric argument. +//! +//! It accepts exactly one number and preserves its selected input type. For signed integers, +//! negative values are negated with 64-bit wrapping, so `ABS(i64::MIN)` remains `i64::MIN`. +//! Unsigned values are unchanged. Floating-point values are negated only when `< 0.0`; +//! this preserves negative zero and leaves NaN (including its payload/sign bits) untouched. +//! +//! Null input returns null. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{FloatCC, InstBuilder, IntCC, types}; + +use super::add::{select_return_type, with_float_fallback}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct AbsFnCall { + arg: Box, +} + +impl FnCall for AbsFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::NUMERICAL).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Abs, + expected: target_type, + got: InferredTypeSet::NUMERICAL, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Abs, + expected: 1, + got: args.len(), + }); + } + let arg_types = + crate::ast::infer_types_aux(&args[0], InferredTypeSet::NUMERICAL, inferred_types)?; + let return_types = arg_types.intersect(target_type); + if return_types.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Abs, + expected: target_type, + got: arg_types, + }); + } + Ok(return_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let target_types = match &args[0] { + UntypedExpr::Literal(literal) => { + let declared = InferredTypeSet::singleton(literal.r#type()); + let constrained = declared.intersect(target_type_set); + if constrained.is_none() { + InferredTypeSet::NUMERICAL.intersect(target_type_set) + } else { + constrained + } + } + UntypedExpr::Variable(variable) => context + .variable_types() + .get(variable.as_ref()) + .copied() + .map(InferredTypeSet::singleton) + .unwrap_or(InferredTypeSet::NONE) + .intersect(target_type_set), + _ => InferredTypeSet::NUMERICAL.intersect(target_type_set), + }; + let return_type = select_return_type(with_float_fallback(target_types)); + let arg = context.apply_types(&args[0], InferredTypeSet::singleton(return_type))?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(AbsFnCall { arg: Box::new(arg) }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("ABS", std::iter::once(self.arg.as_ref()), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let arg = context.compile_expr(&self.arg, builder)?; + let value = match return_type { + VarType::I64 => { + let is_negative = builder + .ins() + .icmp_imm_s(IntCC::SignedLessThan, arg.value, 0); + let negated = builder.ins().ineg(arg.value); + builder.ins().select(is_negative, negated, arg.value) + } + VarType::U64 => arg.value, + VarType::F64 => { + let zero = builder.ins().f64const(0.0); + let is_negative = builder.ins().fcmp(FloatCC::LessThan, arg.value, zero); + let negated = builder.ins().fneg(arg.value); + builder.ins().select(is_negative, negated, arg.value) + } + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Abs, + return_type, + }); + } + }; + Ok(LoweredValue { + value, + is_present: arg.is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: AbsFnCall) -> Self { + FnCallEnum::Abs(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_requires_one_numeric_argument() { + let expression = deserialize("(ABS value)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!( + inferred_types.get("value"), + Some(&InferredTypeSet::NUMERICAL) + ); + + for expression in ["(ABS)", "(ABS one two)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Abs, + expected: 1, + .. + }) + )); + } + } + + #[test] + fn test_preserves_declared_numeric_types() { + let cases = [ + ("(ABS -7i64)", VarType::I64), + ("(ABS 7u64)", VarType::U64), + ("(ABS -7f64)", VarType::F64), + ]; + for (expression, expected_type) in cases { + let expression = deserialize(expression).unwrap(); + let compiled = compile(&expression, &HashMap::new()).unwrap(); + assert_eq!(compiled.result_type(), expected_type); + } + } + + #[test] + fn test_integer_edges_and_runtime_null() { + let expression = deserialize("(ABS value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::I64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(-7i64)]).as_i64() }, + Some(7) + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(i64::MIN)]).as_i64() }, + Some(i64::MIN) + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_i64() }, + None + ); + } + + #[test] + fn test_float_preserves_negative_zero_and_nan_bits() { + let expression = deserialize("(ABS value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::F64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + let negative_zero = unsafe { + compiled + .call(&[VariableValue::some(-0.0f64)]) + .as_f64() + .unwrap() + }; + assert_eq!(negative_zero.to_bits(), (-0.0f64).to_bits()); + + let nan = f64::from_bits(0xfff8_0000_0000_0042); + let output = unsafe { compiled.call(&[VariableValue::some(nan)]).as_f64().unwrap() }; + assert_eq!(output.to_bits(), nan.to_bits()); + } +} diff --git a/jitexpr/src/functions/add.rs b/jitexpr/src/functions/add.rs new file mode 100644 index 0000000000..5a97825aed --- /dev/null +++ b/jitexpr/src/functions/add.rs @@ -0,0 +1,482 @@ +// Adds takes an arbitrary number of arguments and adds them. +// +// The type of the addition is rather complex. +// We consider the possible types of all arguments, make an intersection of those, and +// pick the first available type with the order of priority i64, u64, f64. +// +// For instance (ADD mycol 1f64) where mycol is i64 will coerce +// 1f64 to 1i64 at compile time (because we have detected that the conversion was lossless), and the +// operation will run over integer. +// +// On the other hand, (ADD mycol 1.2f64) where mycol is i64 will coerce +// mycol to float dynamically (because 1.2f64 cannot be converted to u64 with loss). +// +// If any of the values of the arguments is NaN, the function return NaN. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct AddFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for AddFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Any; + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::NUMERICAL).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Add, + expected: target_type, + got: InferredTypeSet::NUMERICAL, + }); + } + let mut return_types = InferredTypeSet::NUMERICAL; + for arg in args { + let arg_types = + crate::ast::infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types)?; + return_types = return_types.intersect(arg_types); + } + return_types = with_float_fallback(return_types); + let constrained_return_types = return_types.intersect(target_type); + if constrained_return_types.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Add, + expected: target_type, + got: return_types, + }); + } + Ok(constrained_return_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let mut return_types = InferredTypeSet::NUMERICAL.intersect(target_type_set); + for arg in args { + let arg_types = crate::ast::infer_type_with_variable_types( + arg, + InferredTypeSet::NUMERICAL, + context.variable_types(), + )?; + return_types = return_types.intersect(arg_types); + } + let return_type = select_return_type(with_float_fallback(return_types)); + let typed_args: Vec = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::singleton(return_type))) + .collect::>()?; + if typed_args + .iter() + .any(|typed_arg| !is_numerical(typed_arg.return_type)) + { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(AddFnCall { + args: typed_args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("ADD", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let mut sum = match return_type { + VarType::U64 | VarType::I64 => builder.ins().iconst(types::I64, 0), + VarType::F64 => builder.ins().f64const(0.0), + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Add, + return_type, + }); + } + }; + let mut is_present = builder.ins().iconst(types::I8, 1); + + for arg in &self.args { + let lowered = context.compile_expr(arg, builder)?; + sum = match return_type { + VarType::U64 | VarType::I64 => builder.ins().iadd(sum, lowered.value), + VarType::F64 => builder.ins().fadd(sum, lowered.value), + _ => unreachable!("the return type was checked above"), + }; + is_present = builder.ins().band(is_present, lowered.is_present); + } + Ok(LoweredValue { + value: sum, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +pub(super) fn with_float_fallback(inferred_types: InferredTypeSet) -> InferredTypeSet { + if inferred_types.is_none() { + InferredTypeSet::F64 + } else { + inferred_types + } +} + +pub(super) fn select_return_type(inferred_types: InferredTypeSet) -> VarType { + if inferred_types.i64 { + VarType::I64 + } else if inferred_types.u64 { + VarType::U64 + } else { + debug_assert!(inferred_types.f64); + VarType::F64 + } +} + +pub(super) fn is_numerical(var_type: VarType) -> bool { + matches!(var_type, VarType::I64 | VarType::U64 | VarType::F64) +} + +impl From for FnCallEnum { + fn from(call: AddFnCall) -> Self { + FnCallEnum::Add(call) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::ast::{Literal, infer_types}; + use crate::compile::{TypedExprAst, compile}; + use crate::types::VariableValue; + + #[test] + fn test_infer_types_rejects_string_argument() { + let expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::literal(1.0), + UntypedExpr::literal("hello"), + ]); + let error = infer_types(&expr).unwrap_err(); + assert!(matches!( + error, + TypeError::InvalidLiteralType { + literal: Literal::String(_), + expected: InferredTypeSet::NUMERICAL, + } + )); + } + + #[test] + fn test_infer_types_constrains_variables_to_numerical() { + let expr = Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("a"), UntypedExpr::variable("b")]); + let inferred_types = infer_types(&expr).unwrap(); + assert_eq!(inferred_types.get("a"), Some(&InferredTypeSet::NUMERICAL)); + assert_eq!(inferred_types.get("b"), Some(&InferredTypeSet::NUMERICAL)); + } + + #[test] + fn test_call_with_types_preserves_u64() { + let variable_types = HashMap::from([("present", VarType::U64)]); + let typed_expr = crate::typed_expr_from_str("(ADD present 1u64)", &variable_types); + assert_eq!(typed_expr.return_type, VarType::U64); + assert_eq!( + typed_expr, + TypedExpr { + return_type: VarType::U64, + ast: TypedExprAst::from_call(AddFnCall { + args: vec![ + TypedExprAst::variable("present", VarType::U64).with_type(VarType::U64), + TypedExpr::literal(1u64), + ] + .into_boxed_slice() + }), + } + ); + } + + #[test] + fn test_call_with_types_rematerializes_compatible_literal_as_u64() { + let variable_types = HashMap::from([("present", VarType::U64)]); + let typed_expr = crate::typed_expr_from_str("(ADD present 1i64)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::U64); + assert_eq!( + typed_expr, + TypedExpr { + return_type: VarType::U64, + ast: TypedExprAst::from_call(AddFnCall { + args: vec![ + TypedExprAst::variable("present", VarType::U64).with_type(VarType::U64), + TypedExpr::literal(1u64), + ] + .into_boxed_slice() + }), + } + ); + } + + #[test] + fn test_call_with_types_rematerializes_integral_f64_literal_as_u64() { + let variable_types = HashMap::from([("present", VarType::U64)]); + let typed_expr = crate::typed_expr_from_str("(ADD 1.0f64 present)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::U64); + let TypedExprAst::FnCall(FnCallEnum::Add(call)) = typed_expr.ast else { + panic!("expected an ADD call"); + }; + assert_eq!(call.args[0], TypedExpr::literal(1u64)); + } + + #[test] + fn test_call_with_types_rematerializes_compatible_literal_as_i64() { + let variable_types = HashMap::from([("present", VarType::I64)]); + let typed_expr = crate::typed_expr_from_str("(ADD present 1u64)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::I64); + let TypedExprAst::FnCall(FnCallEnum::Add(call)) = typed_expr.ast else { + panic!("expected an ADD call"); + }; + assert_eq!(call.args[1], TypedExpr::literal(1i64)); + } + + #[test] + fn test_call_with_types_prefers_i64_for_compatible_literals() { + let variable_types = HashMap::new(); + let typed_expr = crate::typed_expr_from_str("(ADD 1u64 2.0f64)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::I64); + let TypedExprAst::FnCall(FnCallEnum::Add(call)) = typed_expr.ast else { + panic!("expected an ADD call"); + }; + assert_eq!(call.args[0], TypedExpr::literal(1i64)); + assert_eq!(call.args[1], TypedExpr::literal(2i64)); + } + + #[test] + fn test_call_with_types_uses_u64_when_i64_is_not_possible() { + let variable_types = HashMap::new(); + let expression = crate::ast::deserialize("(ADD 9223372036854775808u64)").unwrap(); + let typed_expr = + crate::typed_expr_from_str("(ADD 9223372036854775808u64)", &variable_types); + assert_eq!(typed_expr.return_type, VarType::U64); + assert_eq!( + typed_expr, + TypedExpr { + return_type: VarType::U64, + ast: TypedExprAst::from_call(AddFnCall { + args: vec![TypedExpr::literal(9223372036854775808u64)].into_boxed_slice(), + }), + } + ); + + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let output = unsafe { compiled.call(&[]) }; + + assert_eq!(unsafe { output.as_u64() }, Some(9223372036854775808u64)); + } + + #[test] + fn test_call_with_types_coerces_mixed_numbers_to_f64() { + let variable_types = HashMap::from([("present", VarType::U64)]); + let typed_expr = crate::typed_expr_from_str("(ADD present 1.2f64)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::F64); + assert!(matches!(typed_expr.ast, TypedExprAst::FnCall(_))); + } + + #[test] + fn test_call_with_types_propagates_missing_variable() { + let variable_types = HashMap::from([("present", VarType::U64)]); + let typed_expr = + crate::typed_expr_from_str("(ADD present (ADD 1u64 missing))", &variable_types); + + assert_eq!(typed_expr, TypedExpr::none()); + } + + #[test] + fn test_compile_signed_add() { + let expression = Function::Add.call_untyped_expr(vec![ + UntypedExpr::literal(-4i64), + UntypedExpr::variable("myfield"), + ]); + let variable_types = HashMap::from([("myfield", VarType::I64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(-8i64)]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_i64() }, Some(-12)); + } + + #[test] + fn test_compile_adds_i64_literal_to_u64_variable_without_float_coercion() { + let variable_types = HashMap::from([("myfield", VarType::U64)]); + let argument_orders = [ + vec![UntypedExpr::variable("myfield"), UntypedExpr::literal(1i64)], + vec![UntypedExpr::literal(1i64), UntypedExpr::variable("myfield")], + ]; + + for args in argument_orders { + let expression = Function::Add.call_untyped_expr(args); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(41u64)]; + let output = unsafe { compiled.call(&input) }; + assert_eq!(unsafe { output.as_u64() }, Some(42)); + } + } + + #[test] + fn test_compile_coerces_compatible_nested_add_to_u64() { + let nested_literals = Function::Add + .call_untyped_expr(vec![UntypedExpr::literal(1i64), UntypedExpr::literal(2u64)]); + let expression = Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("myfield"), nested_literals]); + let variable_types = HashMap::from([("myfield", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(39u64)]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_u64() }, Some(42)); + } + + #[test] + fn test_compile_coerces_integers_to_float() { + let expression = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("myfield"), + UntypedExpr::literal(-2i64), + UntypedExpr::literal(0.5f64), + ]); + let variable_types = HashMap::from([("myfield", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(10u64)]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_f64() }, Some(8.5)); + } + + #[test] + fn test_compile_loads_multiple_variable_slots() { + let expression = Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("x"), UntypedExpr::variable("y")]); + let variable_types = HashMap::from([("x", VarType::U64), ("y", VarType::F64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(10u64), VariableValue::some(0.5f64)]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_f64() }, Some(10.5)); + } + + #[test] + fn test_compile_add_propagates_none_input() { + let expression = crate::ast::deserialize(r#"(ADD x 0.5f64)"#).unwrap(); + let variable_types = HashMap::from([("x", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::none()]; + let output = unsafe { compiled.call(&input) }; + assert_eq!(unsafe { output.as_f64() }, None); + } + + #[test] + fn test_compile_u64_to_float_coercion_is_unsigned() { + let expression = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("x"), + UntypedExpr::literal(0.5f64), + ]); + let variable_types = HashMap::from([("x", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(u64::MAX)]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_f64() }, Some(u64::MAX as f64 + 0.5)); + } + + #[test] + fn test_compile_reuses_repeated_variable_slot() { + let expression = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("x"), + UntypedExpr::variable("x"), + UntypedExpr::literal(1u64), + ]); + let variable_types = HashMap::from([("x", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(4i64)]; + assert_eq!(compiled.inputs.len(), 1); + assert_eq!(compiled.inputs[0].variable_name.as_ref(), "x"); + assert_eq!(compiled.inputs[0].r#type, VarType::U64); + assert_eq!(compiled.inputs[0].variable_id, 0); + assert_eq!(compiled.result_type(), VarType::U64); + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_u64() }, Some(9)); + } + + #[test] + fn test_compile_can_coerce_variable_when_necessary() { + let expression = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("x"), + UntypedExpr::literal(1.2f64), + ]); + let variable_types = HashMap::from([("x", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some(4u64)]; + assert_eq!(compiled.inputs.len(), 1); + assert_eq!(compiled.inputs[0].variable_name.as_ref(), "x"); + assert_eq!(compiled.inputs[0].r#type, VarType::U64); + assert_eq!(compiled.inputs[0].variable_id, 0); + assert_eq!(compiled.result_type(), VarType::F64); + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_f64() }, Some(5.2f64)); + } + + #[test] + fn test_compile_empty_add_uses_zero_identity() { + let variable_types = HashMap::new(); + let typed_expr = crate::typed_expr_from_str("(ADD)", &variable_types); + assert_eq!(typed_expr.return_type, VarType::I64); + + let expression = Function::Add.call_untyped_expr(Vec::new()); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + let output = unsafe { compiled.call(&[]) }; + assert_eq!(unsafe { output.as_i64() }, Some(0)); + } + + #[test] + fn test_no_variable_works() { + let args = vec![UntypedExpr::literal(1.2f64), UntypedExpr::literal(1u64)]; + let variable_types = HashMap::new(); + let typed_expr = crate::typed_expr_from_str("(ADD 1.2f64 1u64)", &variable_types); + assert_eq!(typed_expr.return_type, VarType::F64); + let expression = Function::Add.call_untyped_expr(args); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + let output = unsafe { compiled.call(&[]) }; + assert_eq!(unsafe { output.as_f64() }, Some(2.2f64)); + } +} diff --git a/jitexpr/src/functions/and.rs b/jitexpr/src/functions/and.rs new file mode 100644 index 0000000000..801f975625 --- /dev/null +++ b/jitexpr/src/functions/and.rs @@ -0,0 +1,188 @@ +//! `AND` combines one or more nullable boolean expressions. +//! +//! All arguments must be boolean. With present inputs it is ordinary conjunction. Null handling is +//! deliberately stricter than SQL: if any argument is absent, the result is absent even when +//! another argument is already `false`. Thus `FALSE AND NULL = NULL`. Null results carry a false +//! payload. One or more operands are required, and every child expression is evaluated. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct AndFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for AndFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::AtLeast(1); + + 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::And, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.is_empty() { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::And, + expected: 1, + got: 0, + }); + } + + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::BOOLEAN, inferred_types)?; + } + Ok(InferredTypeSet::BOOLEAN) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::BOOLEAN)) + .collect::, _>>()?; + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(AndFnCall { + args: args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("AND", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + let mut value = builder.ins().iconst(types::I8, 1); + let mut is_present = value; + + for arg in &self.args { + let lowered = context.compile_expr(arg, builder)?; + is_present = builder.ins().band(is_present, lowered.is_present); + if arg.return_type == VarType::None { + value = builder.ins().iconst(types::I8, 0); + } else { + value = builder.ins().band(value, lowered.value); + } + } + value = builder.ins().band(value, is_present); + + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: AndFnCall) -> Self { + FnCallEnum::And(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no runtime inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_infer_types_requires_boolean_arguments() { + let expression = deserialize("(AND left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("left"), Some(&InferredTypeSet::BOOLEAN)); + assert_eq!(inferred_types.get("right"), Some(&InferredTypeSet::BOOLEAN)); + + let expression = deserialize("(AND)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::And, + expected: 1, + got: 0, + }) + )); + } + + #[test] + fn test_present_truth_table_and_variadic_inputs() { + assert_eq!(eval("(AND true)"), Some(true)); + assert_eq!(eval("(AND true true true)"), Some(true)); + assert_eq!(eval("(AND true false true)"), Some(false)); + assert_eq!(eval("(AND false false)"), Some(false)); + } + + #[test] + fn test_any_absent_argument_makes_result_absent() { + assert_eq!(eval("(AND true none)"), None); + assert_eq!(eval("(AND false none)"), None); + assert_eq!(eval("(AND none none)"), None); + } + + #[test] + fn test_runtime_null_propagation() { + let expression = deserialize("(AND left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::Bool), ("right", VarType::Bool)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(false), VariableValue::none()]) + .as_bool() + }, + None + ); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(true), VariableValue::some(false)]) + .as_bool() + }, + Some(false) + ); + } +} diff --git a/jitexpr/src/functions/ceil.rs b/jitexpr/src/functions/ceil.rs new file mode 100644 index 0000000000..39cdd7e4d3 --- /dev/null +++ b/jitexpr/src/functions/ceil.rs @@ -0,0 +1,213 @@ +//! `CEIL` returns the least integer greater than or equal to one numeric argument. +//! +//! The result type is always `i64`. Signed integers are returned unchanged. Unsigned integers are +//! returned unchanged when they fit in `i64`; larger values return null. Floating-point inputs are +//! rounded toward positive infinity and converted to `i64`. Null, NaN, infinity, and any result +//! outside the `i64` range return null. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{FloatCC, InstBuilder, IntCC, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct CeilFnCall { + arg: Box, +} + +impl FnCall for CeilFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::I64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Ceil, + expected: target_type, + got: InferredTypeSet::I64, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Ceil, + expected: 1, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::NUMERICAL, inferred_types)?; + Ok(InferredTypeSet::I64) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::I64)); + let arg = context.apply_types(&args[0], InferredTypeSet::NUMERICAL)?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type: VarType::I64, + ast: TypedExprAst::from_call(CeilFnCall { arg: Box::new(arg) }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("CEIL", std::iter::once(self.arg.as_ref()), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + if return_type != VarType::I64 { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Ceil, + return_type, + }); + } + let arg = context.compile_expr(&self.arg, builder)?; + let (value, value_is_valid) = match self.arg.return_type { + VarType::I64 => { + let valid = builder.ins().iconst(types::I8, 1); + (arg.value, valid) + } + VarType::U64 => { + let max = builder.ins().iconst(types::I64, i64::MAX); + let valid = builder + .ins() + .icmp(IntCC::UnsignedLessThanOrEqual, arg.value, max); + (arg.value, valid) + } + VarType::F64 => { + let rounded = builder.ins().ceil(arg.value); + let lower_bound = builder.ins().f64const(i64::MIN as f64); + let upper_bound = builder.ins().f64const(-(i64::MIN as f64)); + let above_lower = + builder + .ins() + .fcmp(FloatCC::GreaterThanOrEqual, rounded, lower_bound); + let below_upper = builder.ins().fcmp(FloatCC::LessThan, rounded, upper_bound); + let valid = builder.ins().band(above_lower, below_upper); + let value = builder.ins().fcvt_to_sint_sat(types::I64, rounded); + (value, valid) + } + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Ceil, + return_type: self.arg.return_type, + }); + } + }; + let is_present = builder.ins().band(arg.is_present, value_is_valid); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: CeilFnCall) -> Self { + FnCallEnum::Ceil(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable i64 value. + unsafe { compiled.call(&[]).as_i64() } + } + + #[test] + fn test_signature_and_output_type() { + let expression = deserialize("(CEIL value)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + assert_eq!(inferred.get("value"), Some(&InferredTypeSet::NUMERICAL)); + + for expression in ["(CEIL)", "(CEIL 1i64 2i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Ceil, + expected: 1, + .. + }) + )); + } + + let expression = deserialize("(CEIL 1.2f64)").unwrap(); + assert_eq!( + compile(&expression, &HashMap::new()).unwrap().result_type(), + VarType::I64 + ); + } + + #[test] + fn test_integer_identity_and_float_rounding() { + assert_eq!(eval("(CEIL -9223372036854775808i64)"), Some(i64::MIN)); + assert_eq!(eval("(CEIL 9223372036854775807u64)"), Some(i64::MAX)); + assert_eq!(eval("(CEIL 9223372036854775808u64)"), None); + assert_eq!(eval("(CEIL 1.2f64)"), Some(2)); + assert_eq!(eval("(CEIL -1.2f64)"), Some(-1)); + assert_eq!(eval("(CEIL -0f64)"), Some(0)); + } + + #[test] + fn test_null_and_exceptional_float_results() { + assert_eq!(eval("(CEIL none)"), None); + assert_eq!(eval("(CEIL nanf64)"), None); + assert_eq!(eval("(CEIL inff64)"), None); + assert_eq!(eval("(CEIL -inff64)"), None); + assert_eq!(eval("(CEIL 9223372036854775808f64)"), None); + assert_eq!(eval("(CEIL -9223372036854775808f64)"), Some(i64::MIN)); + } + + #[test] + fn test_runtime_null() { + let expression = deserialize("(CEIL value)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::F64)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable f64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(2.1f64)]).as_i64() }, + Some(3) + ); + // SAFETY: The compiled expression expects one nullable f64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_i64() }, + None + ); + } +} diff --git a/jitexpr/src/functions/comparison.rs b/jitexpr/src/functions/comparison.rs new file mode 100644 index 0000000000..81b4932fa6 --- /dev/null +++ b/jitexpr/src/functions/comparison.rs @@ -0,0 +1,370 @@ +//! Shared lowering and native helpers for ordered comparisons. + +use std::cmp::Ordering; +use std::collections::HashMap; + +use cranelift::codegen::ir::{ + FuncRef, Function as CraneliftFunction, Type, Value as CraneliftValue, 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, LoweredValue, LoweringContext, TypedExpr}; +use crate::types::VarType; + +const STRING_COMPARE_SYMBOL: &str = "jitexpr_string_compare"; +const F64_I64_COMPARE_SYMBOL: &str = "jitexpr_f64_i64_compare"; +const F64_U64_COMPARE_SYMBOL: &str = "jitexpr_f64_u64_compare"; +const UNORDERED: i8 = i8::MIN; + +const ORDERED_TYPES: InferredTypeSet = InferredTypeSet { + string: true, + i64: true, + u64: true, + f64: true, + boolean: false, +}; + +#[derive(Clone, Copy)] +pub(super) enum OrderedComparison { + GreaterThan, + LessThan, + GreaterThanOrEqual, + LessThanOrEqual, +} + +pub(super) fn infer_types<'a>( + function: Function, + 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, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function, + expected: 2, + got: args.len(), + }); + } + + for arg in args { + crate::ast::infer_types_aux(arg, ORDERED_TYPES, inferred_types)?; + } + Ok(InferredTypeSet::BOOLEAN) +} + +pub(super) fn apply_types( + args: &[UntypedExpr], + context: &mut CompileFnBuilder<'_, '_>, +) -> Result, CompileError> { + args.iter() + .map(|arg| match arg { + UntypedExpr::Literal(literal) => { + context.apply_types(arg, InferredTypeSet::singleton(literal.r#type())) + } + _ => context.apply_types(arg, ORDERED_TYPES), + }) + .collect::, _>>() + .map(Vec::into_boxed_slice) +} + +pub(super) fn lower( + args: &[TypedExpr], + comparison: OrderedComparison, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, +) -> Result { + let lhs_type = args[0].return_type; + let rhs_type = args[1].return_type; + let lhs = context.compile_expr(&args[0], builder)?; + let rhs = context.compile_expr(&args[1], builder)?; + let is_present = builder.ins().band(lhs.is_present, rhs.is_present); + let string_len = builder.ins().iconst(types::I64, 0); + + let compared = if lhs_type == VarType::None || rhs_type == VarType::None { + builder.ins().iconst(types::I8, 0) + } else if lhs_type == VarType::Str && rhs_type == VarType::Str { + let null = builder.ins().iconst(context.pointer_type(), 0); + let lhs_ptr = builder.ins().select(lhs.is_present, lhs.value, null); + let rhs_ptr = builder.ins().select(rhs.is_present, rhs.value, null); + let call = builder.ins().call( + context.native_functions().string_compare(), + &[lhs_ptr, lhs.string_len, rhs_ptr, rhs.string_len], + ); + compare_ordering(builder.inst_results(call)[0], comparison, builder) + } else if is_numerical(lhs_type) && is_numerical(rhs_type) { + lower_numeric(lhs, lhs_type, rhs, rhs_type, comparison, context, builder) + } else { + builder.ins().iconst(types::I8, 0) + }; + let value = builder.ins().band(compared, is_present); + Ok(LoweredValue { + value, + is_present, + string_len, + }) +} + +fn lower_numeric( + lhs: LoweredValue, + lhs_type: VarType, + rhs: LoweredValue, + rhs_type: VarType, + comparison: OrderedComparison, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, +) -> CraneliftValue { + match (lhs_type, rhs_type) { + (VarType::I64, VarType::I64) => { + builder + .ins() + .icmp(comparison.signed_int_cc(), lhs.value, rhs.value) + } + (VarType::U64, VarType::U64) => { + builder + .ins() + .icmp(comparison.unsigned_int_cc(), lhs.value, rhs.value) + } + (VarType::F64, VarType::F64) => { + builder + .ins() + .fcmp(comparison.float_cc(), lhs.value, rhs.value) + } + (VarType::I64, VarType::U64) => { + let ordering = compare_signed_unsigned(lhs.value, rhs.value, builder); + compare_ordering(ordering, comparison, builder) + } + (VarType::U64, VarType::I64) => { + let ordering = compare_signed_unsigned(rhs.value, lhs.value, builder); + let reversed = builder.ins().ineg(ordering); + compare_ordering(reversed, comparison, builder) + } + (VarType::F64, VarType::I64) => { + let call = builder.ins().call( + context.native_functions().f64_i64_compare(), + &[lhs.value, rhs.value], + ); + compare_ordering(builder.inst_results(call)[0], comparison, builder) + } + (VarType::I64, VarType::F64) => { + let call = builder.ins().call( + context.native_functions().f64_i64_compare(), + &[rhs.value, lhs.value], + ); + let ordering = builder.inst_results(call)[0]; + let reversed = builder.ins().ineg(ordering); + compare_ordering(reversed, comparison, builder) + } + (VarType::F64, VarType::U64) => { + let call = builder.ins().call( + context.native_functions().f64_u64_compare(), + &[lhs.value, rhs.value], + ); + compare_ordering(builder.inst_results(call)[0], comparison, builder) + } + (VarType::U64, VarType::F64) => { + let call = builder.ins().call( + context.native_functions().f64_u64_compare(), + &[rhs.value, lhs.value], + ); + let ordering = builder.inst_results(call)[0]; + let reversed = builder.ins().ineg(ordering); + compare_ordering(reversed, comparison, builder) + } + _ => unreachable!("ordered numeric comparison received non-numeric types"), + } +} + +fn compare_signed_unsigned( + signed: CraneliftValue, + unsigned: CraneliftValue, + builder: &mut FunctionBuilder<'_>, +) -> CraneliftValue { + let less = builder + .ins() + .icmp(IntCC::UnsignedLessThan, signed, unsigned); + let greater = builder + .ins() + .icmp(IntCC::UnsignedGreaterThan, signed, unsigned); + let minus_one = builder.ins().iconst(types::I8, -1); + let zero = builder.ins().iconst(types::I8, 0); + let one = builder.ins().iconst(types::I8, 1); + let greater_or_equal = builder.ins().select(greater, one, zero); + let nonnegative_ordering = builder.ins().select(less, minus_one, greater_or_equal); + let is_negative = builder.ins().icmp_imm_s(IntCC::SignedLessThan, signed, 0); + builder + .ins() + .select(is_negative, minus_one, nonnegative_ordering) +} + +fn compare_ordering( + ordering: CraneliftValue, + comparison: OrderedComparison, + builder: &mut FunctionBuilder<'_>, +) -> CraneliftValue { + let ordered = builder + .ins() + .icmp_imm_s(IntCC::NotEqual, ordering, i64::from(UNORDERED)); + let relation = builder + .ins() + .icmp_imm_s(comparison.ordering_cc(), ordering, 0); + builder.ins().band(ordered, relation) +} + +fn is_numerical(var_type: VarType) -> bool { + matches!(var_type, VarType::I64 | VarType::U64 | VarType::F64) +} + +impl OrderedComparison { + fn signed_int_cc(self) -> IntCC { + match self { + Self::GreaterThan => IntCC::SignedGreaterThan, + Self::LessThan => IntCC::SignedLessThan, + Self::GreaterThanOrEqual => IntCC::SignedGreaterThanOrEqual, + Self::LessThanOrEqual => IntCC::SignedLessThanOrEqual, + } + } + + fn unsigned_int_cc(self) -> IntCC { + match self { + Self::GreaterThan => IntCC::UnsignedGreaterThan, + Self::LessThan => IntCC::UnsignedLessThan, + Self::GreaterThanOrEqual => IntCC::UnsignedGreaterThanOrEqual, + Self::LessThanOrEqual => IntCC::UnsignedLessThanOrEqual, + } + } + + fn float_cc(self) -> FloatCC { + match self { + Self::GreaterThan => FloatCC::GreaterThan, + Self::LessThan => FloatCC::LessThan, + Self::GreaterThanOrEqual => FloatCC::GreaterThanOrEqual, + Self::LessThanOrEqual => FloatCC::LessThanOrEqual, + } + } + + fn ordering_cc(self) -> IntCC { + match self { + Self::GreaterThan => IntCC::SignedGreaterThan, + Self::LessThan => IntCC::SignedLessThan, + Self::GreaterThanOrEqual => IntCC::SignedGreaterThanOrEqual, + Self::LessThanOrEqual => IntCC::SignedLessThanOrEqual, + } + } +} + +pub(super) fn register_jit_symbols(jit_builder: &mut JITBuilder) { + jit_builder.symbol(STRING_COMPARE_SYMBOL, string_compare as *const u8); + jit_builder.symbol(F64_I64_COMPARE_SYMBOL, f64_i64_compare as *const u8); + jit_builder.symbol(F64_U64_COMPARE_SYMBOL, f64_u64_compare as *const u8); +} + +pub(super) struct NativeComparisonFunctions { + pub(super) string_compare: FuncRef, + pub(super) f64_i64_compare: FuncRef, + pub(super) f64_u64_compare: FuncRef, +} + +pub(super) fn declare_native_functions( + module: &mut JITModule, + function: &mut CraneliftFunction, + pointer_type: Type, +) -> Result { + let mut string_signature = module.make_signature(); + string_signature.params.extend([ + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + ]); + string_signature.returns.push(AbiParam::new(types::I8)); + let string_function = + module.declare_function(STRING_COMPARE_SYMBOL, Linkage::Import, &string_signature)?; + + let mut numeric_signature = module.make_signature(); + numeric_signature + .params + .extend([AbiParam::new(types::F64), AbiParam::new(types::I64)]); + numeric_signature.returns.push(AbiParam::new(types::I8)); + let f64_i64_function = + module.declare_function(F64_I64_COMPARE_SYMBOL, Linkage::Import, &numeric_signature)?; + let f64_u64_function = + module.declare_function(F64_U64_COMPARE_SYMBOL, Linkage::Import, &numeric_signature)?; + + Ok(NativeComparisonFunctions { + string_compare: module.declare_func_in_func(string_function, function), + f64_i64_compare: module.declare_func_in_func(f64_i64_function, function), + f64_u64_compare: module.declare_func_in_func(f64_u64_function, function), + }) +} + +unsafe extern "C" fn string_compare( + lhs_ptr: *const u8, + lhs_len: usize, + rhs_ptr: *const u8, + rhs_len: usize, +) -> i8 { + if lhs_ptr.is_null() || rhs_ptr.is_null() { + return 0; + } + // SAFETY: Generated code passes live UTF-8 pointers and their exact byte lengths for present + // string values. Null pointers were rejected above. + let lhs = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(lhs_ptr, lhs_len)) }; + // SAFETY: Same contract as `lhs`. + let rhs = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(rhs_ptr, rhs_len)) }; + ordering_code(lhs.cmp(rhs)) +} + +extern "C" fn f64_i64_compare(lhs: f64, rhs: i64) -> i8 { + if lhs.is_nan() { + return UNORDERED; + } + if lhs < i64::MIN as f64 { + return -1; + } + if lhs >= -(i64::MIN as f64) { + return 1; + } + let truncated = lhs as i64; + match truncated.cmp(&rhs) { + Ordering::Equal => ordering_code(lhs.partial_cmp(&(rhs as f64)).unwrap()), + ordering => ordering_code(ordering), + } +} + +extern "C" fn f64_u64_compare(lhs: f64, rhs: u64) -> i8 { + if lhs.is_nan() { + return UNORDERED; + } + if lhs < 0.0 { + return -1; + } + if lhs >= 2f64.powi(64) { + return 1; + } + let truncated = lhs as u64; + match truncated.cmp(&rhs) { + Ordering::Equal => ordering_code(lhs.partial_cmp(&(rhs as f64)).unwrap()), + ordering => ordering_code(ordering), + } +} + +fn ordering_code(ordering: Ordering) -> i8 { + match ordering { + Ordering::Less => -1, + Ordering::Equal => 0, + Ordering::Greater => 1, + } +} diff --git a/jitexpr/src/functions/concat.rs b/jitexpr/src/functions/concat.rs new file mode 100644 index 0000000000..a5e87fed07 --- /dev/null +++ b/jitexpr/src/functions/concat.rs @@ -0,0 +1,497 @@ +//! `CONCAT` joins two or more strings with configurable delimiter behavior. +//! +//! Its signature is `CONCAT(delimiter, ignore_empty, value1, value2, ...)`, with at least four +//! total arguments. `delimiter` and `ignore_empty` must be string literals. The second literal +//! enables skipping empty values only when it equals `"true"` case-insensitively; every other +//! spelling behaves as false. All remaining arguments are strings. +//! +//! Null propagation across value arguments is strict: any null value makes the result null. When +//! empty strings are not ignored, a delimiter is inserted only after output bytes have already been +//! written. Consequently `CONCAT(",", "false", "", "b")` is `"b"`, while the reversed values +//! produce `"b,"`. Empty output is present, not null. +//! +//! Constructed bytes live in the caller's fixed-capacity string arena; arena exhaustion returns +//! null. + +use std::collections::HashMap; +use std::sync::Arc; + +use cranelift::codegen::ir::{ + FuncRef, Function as CraneliftFunction, StackSlotData, StackSlotKind, Type, types, +}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, StringArena, TypedExpr, + TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_string_concat"; +const RAW_INPUT_SIZE: usize = size_of::(); + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct ConcatFnCall { + arguments: JoinArguments, +} + +#[derive(Clone, Debug, PartialEq)] +pub(super) struct JoinArguments { + delimiter: Arc, + ignore_empty: bool, + values: Box<[TypedExpr]>, +} + +impl FnCall for ConcatFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::AtLeast(4); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + validate_join_args(args) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + infer_join_types(Function::Concat, args, target_type, inferred_types) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let Some(arguments) = apply_join_types("CONCAT", args, context)? else { + return Ok(TypedExpr::none()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(ConcatFnCall { arguments }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + self.arguments.args_mut() + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + self.arguments.serialize("CONCAT", formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + self.arguments.emit_cranelift_ir(context, builder) + } +} + +pub(super) fn validate_join_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + super::validate_literal(args, 0, VarType::Str, |literal| { + matches!(literal, Literal::String(_)) + })?; + super::validate_literal(args, 1, VarType::Str, |literal| { + matches!(literal, Literal::String(_)) + }) +} + +pub(super) fn infer_join_types<'a>( + function: Function, + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, +) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() < 4 { + return Err(TypeError::InvalidNumberOfArguments { + function, + expected: 4, + got: args.len(), + }); + } + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::STRING, inferred_types)?; + } + Ok(InferredTypeSet::STRING) +} + +pub(super) fn apply_join_types( + _function_name: &str, + args: &[UntypedExpr], + context: &mut CompileFnBuilder<'_, '_>, +) -> Result, CompileError> { + let UntypedExpr::Literal(Literal::String(delimiter)) = &args[0] else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 1, + expected: VarType::Str, + } + .into()); + }; + let UntypedExpr::Literal(Literal::String(ignore_empty)) = &args[1] else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + .into()); + }; + + let values = args[2..] + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::STRING)) + .collect::, _>>()?; + if values + .iter() + .any(|value| value.return_type == VarType::None) + { + return Ok(None); + } + debug_assert!(values.iter().all(|value| value.return_type == VarType::Str)); + + Ok(Some(JoinArguments { + delimiter: Arc::from(delimiter.as_ref()), + ignore_empty: ignore_empty.eq_ignore_ascii_case("true"), + values: values.into_boxed_slice(), + })) +} + +impl JoinArguments { + pub(super) fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.values + } + + pub(super) fn serialize( + &self, + function_name: &str, + formatter: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + use std::fmt::Write as _; + + formatter.write_str(function_name)?; + formatter.write_char(' ')?; + crate::compile::format_string_literal(&self.delimiter, formatter)?; + formatter.write_char(' ')?; + crate::compile::format_string_literal( + if self.ignore_empty { "true" } else { "false" }, + formatter, + )?; + for value in &self.values { + write!(formatter, " {value}")?; + } + Ok(()) + } + + pub(super) fn emit_cranelift_ir( + &self, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let stack_size = self + .values + .len() + .checked_mul(RAW_INPUT_SIZE) + .and_then(|size| u32::try_from(size).ok()) + .expect("CONCAT argument descriptors exceed Cranelift stack-slot limits"); + let stack_slot = builder.create_sized_stack_slot(StackSlotData::new( + StackSlotKind::ExplicitSlot, + stack_size, + 3, + )); + + for (index, value) in self.values.iter().enumerate() { + let value = context.compile_expr(value, builder)?; + let offset = i32::try_from(index * RAW_INPUT_SIZE) + .expect("CONCAT argument descriptor offset exceeds i32"); + builder + .ins() + .stack_store(context.pointer_type(), value.value, stack_slot, offset); + builder.ins().stack_store( + context.pointer_type(), + value.string_len, + stack_slot, + offset + 8, + ); + let is_present = builder.ins().uextend(types::I64, value.is_present); + builder + .ins() + .stack_store(context.pointer_type(), is_present, stack_slot, offset + 16); + } + + let inputs_ptr = builder + .ins() + .stack_addr(context.pointer_type(), stack_slot, 0); + let input_count = builder.ins().iconst(types::I64, self.values.len() as i64); + let delimiter_ptr = builder.ins().iconst( + context.pointer_type(), + self.delimiter.as_ptr() as usize as i64, + ); + let delimiter_len = builder + .ins() + .iconst(types::I64, self.delimiter.len() as i64); + let ignore_empty = builder + .ins() + .iconst(types::I8, i64::from(self.ignore_empty)); + let string_arena_ptr = context.string_arena_ptr(builder); + let call = builder.ins().call( + context.native_functions().string_concat(), + &[ + inputs_ptr, + input_count, + delimiter_ptr, + delimiter_len, + ignore_empty, + string_arena_ptr, + ], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder + .ins() + .icmp_imm_u(cranelift::prelude::IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, string_concat 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(types::I64), + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + AbiParam::new(types::I8), + AbiParam::new(pointer_type), + ]); + signature.returns.push(AbiParam::new(pointer_type)); + signature.returns.push(AbiParam::new(types::I64)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct RawInput { + ptr: *const u8, + len: usize, + is_present: u64, +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } +} + +unsafe extern "C" fn string_concat( + inputs_ptr: *const RawInput, + input_count: usize, + delimiter_ptr: *const u8, + delimiter_len: usize, + ignore_empty: u8, + string_arena: *mut StringArena, +) -> RawStr { + if inputs_ptr.is_null() || delimiter_ptr.is_null() || string_arena.is_null() { + return RawStr::none(); + } + // SAFETY: Generated code constructs exactly `input_count` initialized descriptors in an + // aligned stack slot that remains live across this call. + let inputs = unsafe { std::slice::from_raw_parts(inputs_ptr, input_count) }; + // SAFETY: The delimiter is retained by the typed expression and this is its exact byte length. + let delimiter = unsafe { std::slice::from_raw_parts(delimiter_ptr, delimiter_len) }; + + let mut output_len = 0usize; + for input in inputs { + if input.is_present == 0 { + return RawStr::none(); + } + if ignore_empty != 0 && input.len == 0 { + continue; + } + if output_len > 0 && !delimiter.is_empty() { + let Some(next_len) = output_len.checked_add(delimiter.len()) else { + return RawStr::none(); + }; + output_len = next_len; + } + let Some(next_len) = output_len.checked_add(input.len) else { + return RawStr::none(); + }; + output_len = next_len; + } + + // SAFETY: The caller exclusively borrows and passes this arena for the duration of the call. + let Some(output_ptr) = (unsafe { &mut *string_arena }).allocate(output_len) else { + return RawStr::none(); + }; + let mut written = 0usize; + for input in inputs { + if ignore_empty != 0 && input.len == 0 { + continue; + } + if written > 0 && !delimiter.is_empty() { + // SAFETY: The exact output length was checked and allocated above. + unsafe { + std::ptr::copy_nonoverlapping( + delimiter.as_ptr(), + output_ptr.add(written), + delimiter.len(), + ); + } + written += delimiter.len(); + } + // SAFETY: Present string descriptors contain live UTF-8 pointers of exactly `len` bytes; + // source allocations precede and do not overlap the new output allocation. + unsafe { + std::ptr::copy_nonoverlapping(input.ptr, output_ptr.add(written), input.len); + } + written += input.len; + } + debug_assert_eq!(written, output_len); + RawStr { + ptr: output_ptr, + len: output_len, + } +} + +impl From for FnCallEnum { + fn from(call: ConcatFnCall) -> Self { + FnCallEnum::Concat(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::{STRING_ARENA_CAPACITY, compile}; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable strings. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_requires_four_or_more_string_arguments() { + let expression = deserialize(r#"(CONCAT "," "false" left right third)"#).unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("left"), Some(&InferredTypeSet::STRING)); + assert_eq!(inferred_types.get("right"), Some(&InferredTypeSet::STRING)); + assert_eq!(inferred_types.get("third"), Some(&InferredTypeSet::STRING)); + + let expression = deserialize(r#"(CONCAT "," "false" only_one_value)"#).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Concat, + expected: 4, + .. + }) + )); + } + + #[test] + fn test_delimiter_and_ignore_empty_behavior() { + assert_eq!( + eval(r#"(CONCAT "," "false" "a" "b" "c")"#).as_deref(), + Some("a,b,c") + ); + assert_eq!( + eval(r#"(CONCAT "," "TrUe" "a" "" "c")"#).as_deref(), + Some("a,c") + ); + assert_eq!( + eval(r#"(CONCAT "," "not-true" "a" "" "c")"#).as_deref(), + Some("a,,c") + ); + assert_eq!(eval(r#"(CONCAT "," "false" "" "b")"#).as_deref(), Some("b")); + assert_eq!( + eval(r#"(CONCAT "," "false" "b" "")"#).as_deref(), + Some("b,") + ); + assert_eq!(eval(r#"(CONCAT "," "true" "" "")"#).as_deref(), Some("")); + } + + #[test] + fn test_runtime_null_is_strict() { + let expression = deserialize(r#"(CONCAT ":" "true" left right)"#).unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some("a"), VariableValue::none()]) + .as_str() + }, + None + ); + } + + #[test] + fn test_nested_values_are_stable_and_arena_exhaustion_is_null() { + let expression = deserialize(r#"(CONCAT ":" "false" (UPPER left) (LOWER right))"#).unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some("ab"), VariableValue::some("CD")]) + .as_str() + }, + Some("AB:cd") + ); + + let oversized = "a".repeat(STRING_ARENA_CAPACITY); + assert_eq!( + unsafe { + compiled + .call(&[ + VariableValue::some(oversized.as_str()), + VariableValue::some("b"), + ]) + .as_str() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/divide.rs b/jitexpr/src/functions/divide.rs new file mode 100644 index 0000000000..74408a95c0 --- /dev/null +++ b/jitexpr/src/functions/divide.rs @@ -0,0 +1,212 @@ +//! `DIVIDE` performs floating-point division. +//! +//! It accepts exactly two numeric arguments and always coerces both to `f64`, even when both are +//! integers. +//! +//! The result is absent if either operand is absent or if the divisor is positive or negative +//! zero. Division by zero therefore yields NULL rather than infinity or NaN. Otherwise IEEE-754 +//! behavior applies, including propagation of NaN and infinities. + +use std::collections::HashMap; + +use cranelift::prelude::{FloatCC, FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct DivideFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for DivideFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::F64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Divide, + expected: target_type, + got: InferredTypeSet::F64, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Divide, + expected: 2, + got: args.len(), + }); + } + + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types)?; + } + Ok(InferredTypeSet::F64) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::F64)); + + let typed_args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::F64)) + .collect::, _>>()?; + if typed_args + .iter() + .any(|typed_arg| typed_arg.return_type == VarType::None) + { + return Ok(TypedExpr::none()); + } + + Ok(TypedExpr { + return_type: VarType::F64, + ast: TypedExprAst::from_call(DivideFnCall { + args: typed_args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("DIVIDE", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + if return_type != VarType::F64 { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Divide, + return_type, + }); + } + + let dividend = context.compile_expr(&self.args[0], builder)?; + let divisor = context.compile_expr(&self.args[1], builder)?; + let value = builder.ins().fdiv(dividend.value, divisor.value); + let zero = builder.ins().f64const(0.0); + let divisor_is_zero = builder.ins().fcmp(FloatCC::Equal, divisor.value, zero); + let divisor_is_nonzero = builder.ins().bxor_imm_u(divisor_is_zero, 1); + let both_present = builder.ins().band(dividend.is_present, divisor.is_present); + let is_present = builder.ins().band(both_present, divisor_is_nonzero); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: DivideFnCall) -> Self { + FnCallEnum::Divide(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_infer_types_requires_two_numeric_arguments_and_returns_float() { + let expression = deserialize("(DIVIDE left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!( + inferred_types.get("left"), + Some(&InferredTypeSet::NUMERICAL) + ); + assert_eq!( + inferred_types.get("right"), + Some(&InferredTypeSet::NUMERICAL) + ); + + for expression in ["(DIVIDE 1i64)", "(DIVIDE 1i64 2i64 3i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Divide, + expected: 2, + .. + }) + )); + } + } + + #[test] + fn test_integer_inputs_use_floating_point_division() { + let expression = deserialize("(DIVIDE 5i64 2i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + + assert_eq!(compiled.result_type(), VarType::F64); + // SAFETY: The expression has no inputs and returns f64. + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, Some(2.5)); + } + + #[test] + fn test_positive_and_negative_zero_divisors_return_none() { + for expression in ["(DIVIDE 1f64 0f64)", "(DIVIDE 1f64 -0f64)"] { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns nullable f64. + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, None); + } + } + + #[test] + fn test_nan_divisor_remains_present() { + let expression = deserialize("(DIVIDE 1f64 nanf64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + + // SAFETY: The expression has no inputs and returns f64. + assert!(unsafe { compiled.call(&[]).as_f64() }.unwrap().is_nan()); + } + + #[test] + fn test_runtime_null_propagation() { + let expression = deserialize("(DIVIDE left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::I64), ("right", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(9i64), VariableValue::some(2u64)]) + .as_f64() + }, + Some(4.5) + ); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::none(), VariableValue::some(2u64)]) + .as_f64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/eq.rs b/jitexpr/src/functions/eq.rs new file mode 100644 index 0000000000..2624a15104 --- /dev/null +++ b/jitexpr/src/functions/eq.rs @@ -0,0 +1,478 @@ +// 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 present types are not equal. If either +// operand is absent, the result is absent. +// +// 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 as CraneliftValue, 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, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::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 { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + 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 { + Self::ARG_COUNT.validate(args)?; + 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 serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("EQ", self.args.iter(), formatter) + } + + 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; + let lhs = context.compile_expr(&self.args[0], builder)?; + let rhs = context.compile_expr(&self.args[1], builder)?; + let is_present = builder.ins().band(lhs.is_present, rhs.is_present); + let string_len = builder.ins().iconst(types::I64, 0); + + if lhs_type == VarType::None || rhs_type == VarType::None { + let value = builder.ins().iconst(types::I8, 0); + return Ok(LoweredValue { + value, + is_present, + string_len, + }); + } + + if lhs_type == VarType::Str && rhs_type == VarType::Str { + let null = builder.ins().iconst(context.pointer_type(), 0); + let lhs_ptr = builder.ins().select(lhs.is_present, lhs.value, null); + let rhs_ptr = builder.ins().select(rhs.is_present, rhs.value, null); + let call = builder.ins().call( + context.native_functions().string_eq(), + &[lhs_ptr, lhs.string_len, rhs_ptr, rhs.string_len], + ); + return Ok(LoweredValue { + value: builder.inst_results(call)[0], + is_present, + string_len, + }); + } + + if !types_are_comparable(lhs_type, rhs_type) { + let value = builder.ins().iconst(types::I8, 0); + return Ok(LoweredValue { + value, + is_present, + string_len, + }); + } + + let values_equal = match (lhs_type, rhs_type) { + (VarType::Bool, VarType::Bool) + | (VarType::I64, VarType::I64) + | (VarType::U64, VarType::U64) => { + builder.ins().icmp(IntCC::Equal, lhs.value, rhs.value) + } + (VarType::F64, VarType::F64) => { + builder.ins().fcmp(FloatCC::Equal, lhs.value, rhs.value) + } + (VarType::I64, VarType::U64) => emit_signed_unsigned_eq(lhs.value, rhs.value, builder), + (VarType::U64, VarType::I64) => emit_signed_unsigned_eq(rhs.value, lhs.value, builder), + (VarType::F64, VarType::I64) => { + emit_float_integer_eq(lhs.value, rhs.value, VarType::I64, builder) + } + (VarType::I64, VarType::F64) => { + emit_float_integer_eq(rhs.value, lhs.value, VarType::I64, builder) + } + (VarType::F64, VarType::U64) => { + emit_float_integer_eq(lhs.value, rhs.value, VarType::U64, builder) + } + (VarType::U64, VarType::F64) => { + emit_float_integer_eq(rhs.value, lhs.value, VarType::U64, builder) + } + _ => unreachable!("the operand types were checked above"), + }; + let value = builder.ins().band(is_present, values_equal); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +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: CraneliftValue, + unsigned: CraneliftValue, + builder: &mut FunctionBuilder<'_>, +) -> CraneliftValue { + 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: CraneliftValue, + integer: CraneliftValue, + integer_type: VarType, + builder: &mut FunctionBuilder<'_>, +) -> CraneliftValue { + 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(types::I64), + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + ]); + 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_ptr: *const u8, + lhs_len: usize, + rhs_ptr: *const u8, + rhs_len: usize, +) -> u8 { + let lhs = if lhs_ptr.is_null() { + None + } else { + // SAFETY: Generated code passes a live UTF-8 string pointer and its + // exact byte length whenever the pointer is non-null. + Some(unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(lhs_ptr, lhs_len)) }) + }; + let rhs = if rhs_ptr.is_null() { + None + } else { + // SAFETY: Same contract as `lhs` above. + Some(unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(rhs_ptr, rhs_len)) }) + }; + match (lhs, rhs) { + (None, None) => 1, + (Some(lhs), Some(rhs)) => u8::from(lhs == rhs), + _ => 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 mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + let output = unsafe { compiled.call(&[]) }; + + unsafe { output.as_bool() }.unwrap() + } + + fn eval_nullable(expression: &str) -> Option { + let expression = ast::deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable boolean. + unsafe { compiled.call(&[]).as_bool() } + } + + #[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 mut compiled = compile(&expression, &variable_types).unwrap().context(); + + for (signed, unsigned, expected) in [(7i64, 7u64, true), (-1, u64::MAX, false)] { + let input = [VariableValue::some(signed), VariableValue::some(unsigned)]; + let output = unsafe { compiled.call(&input) }; + assert_eq!(unsafe { output.as_bool() }, Some(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 mut compiled = compile(&expression, &variable_types).unwrap().context(); + 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::some(float), VariableValue::some(integer)]; + let output = unsafe { compiled.call(&input) }; + assert_eq!(unsafe { output.as_bool() }, Some(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 mut compiled = compile(&expression, &variable_types).unwrap().context(); + 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::some(float), VariableValue::some(integer)]; + let output = unsafe { compiled.call(&input) }; + assert_eq!(unsafe { output.as_bool() }, Some(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 mut compiled = compile(&expression, &variable_types).unwrap().context(); + let left_value = String::from("same contents"); + let right_value = String::from("same contents"); + let input = [ + VariableValue::some(left_value.as_str()), + VariableValue::some(right_value.as_str()), + ]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_bool() }, Some(true)); + + let different_value = String::from("different contents"); + let input = [ + VariableValue::some(left_value.as_str()), + VariableValue::some(different_value.as_str()), + ]; + let output = unsafe { compiled.call(&input) }; + assert_eq!(unsafe { output.as_bool() }, Some(false)); + + let none_input = [VariableValue::none(), VariableValue::none()]; + let output = unsafe { compiled.call(&none_input) }; + assert_eq!(unsafe { output.as_bool() }, None); + } + + #[test] + fn test_compile_none_equality() { + assert_eq!(eval_nullable("(EQ none none)"), None); + assert_eq!(eval_nullable("(EQ none false)"), None); + } + + #[test] + fn test_compile_runtime_none_equality() { + let expression = ast::deserialize("(EQ left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::U64), ("right", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let output = unsafe { compiled.call(&[VariableValue::none(), VariableValue::none()]) }; + assert_eq!(unsafe { output.as_bool() }, None); + + let output = unsafe { compiled.call(&[VariableValue::none(), VariableValue::some(0u64)]) }; + assert_eq!(unsafe { output.as_bool() }, None); + } + + #[test] + fn test_compile_none_literal_equals_absent_variable() { + let expression = ast::deserialize("(EQ value none)").unwrap(); + let variable_types = HashMap::from([("value", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let output = unsafe { compiled.call(&[VariableValue::none()]) }; + + assert_eq!(unsafe { output.as_bool() }, None); + } + + #[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/floor.rs b/jitexpr/src/functions/floor.rs new file mode 100644 index 0000000000..3de1056461 --- /dev/null +++ b/jitexpr/src/functions/floor.rs @@ -0,0 +1,213 @@ +//! `FLOOR` returns the greatest integer less than or equal to one numeric argument. +//! +//! The result type is always `i64`. Signed integers are returned unchanged. Unsigned integers are +//! returned unchanged when they fit in `i64`; larger values return null. Floating-point inputs are +//! rounded toward negative infinity and converted to `i64`. Null, NaN, infinity, and any result +//! outside the `i64` range return null. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{FloatCC, InstBuilder, IntCC, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct FloorFnCall { + arg: Box, +} + +impl FnCall for FloorFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::I64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Floor, + expected: target_type, + got: InferredTypeSet::I64, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Floor, + expected: 1, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::NUMERICAL, inferred_types)?; + Ok(InferredTypeSet::I64) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::I64)); + let arg = context.apply_types(&args[0], InferredTypeSet::NUMERICAL)?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type: VarType::I64, + ast: TypedExprAst::from_call(FloorFnCall { arg: Box::new(arg) }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("FLOOR", std::iter::once(self.arg.as_ref()), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + if return_type != VarType::I64 { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Floor, + return_type, + }); + } + let arg = context.compile_expr(&self.arg, builder)?; + let (value, value_is_valid) = match self.arg.return_type { + VarType::I64 => { + let valid = builder.ins().iconst(types::I8, 1); + (arg.value, valid) + } + VarType::U64 => { + let max = builder.ins().iconst(types::I64, i64::MAX); + let valid = builder + .ins() + .icmp(IntCC::UnsignedLessThanOrEqual, arg.value, max); + (arg.value, valid) + } + VarType::F64 => { + let rounded = builder.ins().floor(arg.value); + let lower_bound = builder.ins().f64const(i64::MIN as f64); + let upper_bound = builder.ins().f64const(-(i64::MIN as f64)); + let above_lower = + builder + .ins() + .fcmp(FloatCC::GreaterThanOrEqual, rounded, lower_bound); + let below_upper = builder.ins().fcmp(FloatCC::LessThan, rounded, upper_bound); + let valid = builder.ins().band(above_lower, below_upper); + let value = builder.ins().fcvt_to_sint_sat(types::I64, rounded); + (value, valid) + } + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Floor, + return_type: self.arg.return_type, + }); + } + }; + let is_present = builder.ins().band(arg.is_present, value_is_valid); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: FloorFnCall) -> Self { + FnCallEnum::Floor(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable i64 value. + unsafe { compiled.call(&[]).as_i64() } + } + + #[test] + fn test_signature_and_output_type() { + let expression = deserialize("(FLOOR value)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + assert_eq!(inferred.get("value"), Some(&InferredTypeSet::NUMERICAL)); + + for expression in ["(FLOOR)", "(FLOOR 1i64 2i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Floor, + expected: 1, + .. + }) + )); + } + + let expression = deserialize("(FLOOR 1.2f64)").unwrap(); + assert_eq!( + compile(&expression, &HashMap::new()).unwrap().result_type(), + VarType::I64 + ); + } + + #[test] + fn test_integer_identity_and_float_rounding() { + assert_eq!(eval("(FLOOR -9223372036854775808i64)"), Some(i64::MIN)); + assert_eq!(eval("(FLOOR 9223372036854775807u64)"), Some(i64::MAX)); + assert_eq!(eval("(FLOOR 9223372036854775808u64)"), None); + assert_eq!(eval("(FLOOR 1.2f64)"), Some(1)); + assert_eq!(eval("(FLOOR -1.2f64)"), Some(-2)); + assert_eq!(eval("(FLOOR -0f64)"), Some(0)); + } + + #[test] + fn test_null_and_exceptional_float_results() { + assert_eq!(eval("(FLOOR none)"), None); + assert_eq!(eval("(FLOOR nanf64)"), None); + assert_eq!(eval("(FLOOR inff64)"), None); + assert_eq!(eval("(FLOOR -inff64)"), None); + assert_eq!(eval("(FLOOR 9223372036854775808f64)"), None); + assert_eq!(eval("(FLOOR -9223372036854775808f64)"), Some(i64::MIN)); + } + + #[test] + fn test_runtime_null() { + let expression = deserialize("(FLOOR value)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::F64)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable f64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(2.9f64)]).as_i64() }, + Some(2) + ); + // SAFETY: The compiled expression expects one nullable f64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_i64() }, + None + ); + } +} diff --git a/jitexpr/src/functions/gt.rs b/jitexpr/src/functions/gt.rs new file mode 100644 index 0000000000..d62ac25abf --- /dev/null +++ b/jitexpr/src/functions/gt.rs @@ -0,0 +1,162 @@ +//! `GT` tests whether one value is greater than another. +//! +//! It accepts exactly two operands. Ordered operands are strings or numbers; booleans are rejected. +//! Strings use lexicographic UTF-8 ordering. Numeric comparisons support `i64`, `u64`, and `f64` +//! combinations without converting large integers through a lossy `f64`. IEEE unordered +//! comparisons involving NaN return `false`. +//! +//! Null propagation is strict: if either operand is absent, the result is absent. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +use super::comparison::{self, OrderedComparison}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct GtFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for GtFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + comparison::infer_types(Function::Gt, args, target_type, inferred_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(GtFnCall { + args: comparison::apply_types(args, context)?, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("GT", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + comparison::lower(&self.args, OrderedComparison::GreaterThan, context, builder) + } +} + +impl From for FnCallEnum { + fn from(call: GtFnCall) -> Self { + FnCallEnum::Gt(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_requires_two_ordered_arguments() { + let expression = deserialize("(GT left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + let left = inferred_types.get("left").unwrap(); + assert!(left.string && left.i64 && left.u64 && left.f64 && !left.boolean); + + for expression in ["(GT 1i64)", "(GT 1i64 2i64 3i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Gt, + expected: 2, + .. + }) + )); + } + } + + #[test] + fn test_string_and_same_type_numeric_ordering() { + assert_eq!(eval(r#"(GT "beta" "alpha")"#), Some(true)); + assert_eq!(eval(r#"(GT "é" "z")"#), Some(true)); + assert_eq!(eval("(GT -1i64 0i64)"), Some(false)); + assert_eq!(eval("(GT 3.5f64 3f64)"), Some(true)); + assert_eq!(eval("(GT nanf64 0f64)"), Some(false)); + assert_eq!(eval("(GT nanf64 0i64)"), Some(false)); + assert_eq!(eval("(GT 0i64 nanf64)"), Some(false)); + } + + #[test] + fn test_mixed_numeric_boundaries_are_exact() { + assert_eq!(eval("(GT -1i64 18446744073709551615u64)"), Some(false)); + assert_eq!(eval("(GT 18446744073709551615u64 -1i64)"), Some(true)); + + let expression = deserialize("(GT float integer)").unwrap(); + let variable_types = HashMap::from([("float", VarType::F64), ("integer", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[ + VariableValue::some((1u64 << 53) as f64), + VariableValue::some((1u64 << 53) + 1), + ]) + .as_bool() + }, + Some(false) + ); + } + + #[test] + fn test_null_propagates() { + assert_eq!(eval("(GT none 1i64)"), None); + + let expression = deserialize("(GT left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::none(), VariableValue::some("a")]) + .as_bool() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/gt_eq.rs b/jitexpr/src/functions/gt_eq.rs new file mode 100644 index 0000000000..6cd46656f5 --- /dev/null +++ b/jitexpr/src/functions/gt_eq.rs @@ -0,0 +1,145 @@ +//! `GT_EQ` tests whether one value is greater than or equal to another. +//! +//! It accepts exactly two operands. Ordered operands are strings or numbers; booleans are rejected. +//! Strings use lexicographic UTF-8 ordering. Numeric comparisons support `i64`, `u64`, and `f64` +//! combinations without converting large integers through a lossy `f64`. IEEE unordered +//! comparisons involving NaN return `false`, including `NaN >= NaN`. +//! +//! Null propagation is strict: if either operand is absent, the result is absent. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +use super::comparison::{self, OrderedComparison}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct GtEqFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for GtEqFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + comparison::infer_types(Function::GtEq, args, target_type, inferred_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(GtEqFnCall { + args: comparison::apply_types(args, context)?, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("GT_EQ", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + comparison::lower( + &self.args, + OrderedComparison::GreaterThanOrEqual, + context, + builder, + ) + } +} + +impl From for FnCallEnum { + fn from(call: GtEqFnCall) -> Self { + FnCallEnum::GtEq(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_requires_two_ordered_arguments() { + let expression = deserialize("(GT_EQ left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + let left = inferred_types.get("left").unwrap(); + assert!(left.string && left.i64 && left.u64 && left.f64 && !left.boolean); + + let expression = deserialize("(GT_EQ 1i64 2i64 3i64)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::GtEq, + expected: 2, + .. + }) + )); + } + + #[test] + fn test_ordering_equality_nan_and_null() { + assert_eq!(eval(r#"(GT_EQ "same" "same")"#), Some(true)); + assert_eq!(eval(r#"(GT_EQ "alpha" "beta")"#), Some(false)); + assert_eq!(eval("(GT_EQ 0u64 -1i64)"), Some(true)); + assert_eq!( + eval("(GT_EQ 9007199254740993u64 9007199254740992f64)"), + Some(true) + ); + assert_eq!(eval("(GT_EQ nanf64 nanf64)"), Some(false)); + assert_eq!(eval("(GT_EQ nanf64 0i64)"), Some(false)); + assert_eq!(eval("(GT_EQ 0i64 nanf64)"), Some(false)); + assert_eq!(eval("(GT_EQ none 1i64)"), None); + } + + #[test] + fn test_runtime_null_propagates() { + let expression = deserialize("(GT_EQ left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some("a"), VariableValue::none()]) + .as_bool() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/if_fn.rs b/jitexpr/src/functions/if_fn.rs new file mode 100644 index 0000000000..9b3ee47a5f --- /dev/null +++ b/jitexpr/src/functions/if_fn.rs @@ -0,0 +1,218 @@ +//! `IF(condition, when_true, when_false)` selects one branch. +//! +//! The condition must be boolean and both branches are coerced to one common bool, string, or +//! numeric type. A null condition returns null. Otherwise only the selected branch's presence bit +//! controls the result; a null in the unselected branch is ignored. Both branch expressions are +//! lowered eagerly. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{InstBuilder, types}; + +use super::add::with_float_fallback; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct IfFnCall { + args: Box<[TypedExpr]>, +} + +fn common_types(left: InferredTypeSet, right: InferredTypeSet) -> InferredTypeSet { + let direct = left.intersect(right); + if direct.is_none() + && !left.intersect(InferredTypeSet::NUMERICAL).is_none() + && !right.intersect(InferredTypeSet::NUMERICAL).is_none() + { + InferredTypeSet::F64 + } else { + direct + } +} +fn select_type(types: InferredTypeSet) -> VarType { + if types.string { + VarType::Str + } else if types.boolean { + VarType::Bool + } else { + super::add::select_return_type(with_float_fallback( + types.intersect(InferredTypeSet::NUMERICAL), + )) + } +} + +impl FnCall for IfFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(3); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target: InferredTypeSet, + inferred: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if args.len() != 3 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::If, + expected: 3, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::BOOLEAN, inferred)?; + let left = crate::ast::infer_types_aux(&args[1], target, inferred)?; + let right = crate::ast::infer_types_aux(&args[2], target, inferred)?; + let result = common_types(left, right).intersect(target); + if result.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::If, + expected: target, + got: common_types(left, right), + }); + } + Ok(result) + } + + fn call_with_types( + args: &[UntypedExpr], + target: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let condition = context.apply_types(&args[0], InferredTypeSet::BOOLEAN)?; + let left = + crate::ast::infer_type_with_variable_types(&args[1], target, context.variable_types())?; + let right = + crate::ast::infer_type_with_variable_types(&args[2], target, context.variable_types())?; + let return_type = select_type(common_types(left, right).intersect(target)); + let branch_target = InferredTypeSet::singleton(return_type); + let when_true = context.apply_types(&args[1], branch_target)?; + let when_false = context.apply_types(&args[2], branch_target)?; + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(IfFnCall { + args: vec![condition, when_true, when_false].into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("IF", self.args.iter(), formatter) + } + fn emit_cranelift_ir( + &self, + _return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let condition = context.compile_expr(&self.args[0], builder)?; + let when_true = context.compile_expr(&self.args[1], builder)?; + let when_false = context.compile_expr(&self.args[2], builder)?; + let value = builder + .ins() + .select(condition.value, when_true.value, when_false.value); + let branch_present = + builder + .ins() + .select(condition.value, when_true.is_present, when_false.is_present); + let is_present = builder.ins().band(condition.is_present, branch_present); + let string_len = if self.args[1].return_type == VarType::Str { + builder + .ins() + .select(condition.value, when_true.string_len, when_false.string_len) + } else { + builder.ins().iconst(types::I64, 0) + }; + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} +impl From for FnCallEnum { + fn from(call: IfFnCall) -> Self { + FnCallEnum::If(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_signature_and_values() { + assert!(matches!( + infer_types(&deserialize("(IF true 1i64)").unwrap()), + Err(TypeError::InvalidNumberOfArguments { + function: Function::If, + expected: 3, + .. + }) + )); + for (expr, expected) in [("(IF true 7i64 9i64)", 7), ("(IF false 7i64 9i64)", 9)] { + let expr = deserialize(expr).unwrap(); + let mut compiled = compile(&expr, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, Some(expected)); + } + let expr = deserialize("(IF true \"yes\" \"no\")").unwrap(); + let mut compiled = compile(&expr, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_str() }, Some("yes")); + } + + #[test] + fn test_only_selected_null_and_null_condition_propagate() { + let expr = deserialize("(IF condition yes no)").unwrap(); + let types = HashMap::from([ + ("condition", VarType::Bool), + ("yes", VarType::I64), + ("no", VarType::I64), + ]); + let mut compiled = compile(&expr, &types).unwrap().context(); + assert_eq!( + unsafe { + compiled + .call(&[ + VariableValue::some(true), + VariableValue::some(1i64), + VariableValue::none(), + ]) + .as_i64() + }, + Some(1) + ); + assert_eq!( + unsafe { + compiled + .call(&[ + VariableValue::some(false), + VariableValue::some(1i64), + VariableValue::none(), + ]) + .as_i64() + }, + None + ); + assert_eq!( + unsafe { + compiled + .call(&[ + VariableValue::none(), + VariableValue::some(1i64), + VariableValue::some(2i64), + ]) + .as_i64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/int_mod.rs b/jitexpr/src/functions/int_mod.rs new file mode 100644 index 0000000000..ec9f025f4a --- /dev/null +++ b/jitexpr/src/functions/int_mod.rs @@ -0,0 +1,232 @@ +//! `INT_MOD` computes a floating-point modulo. +//! +//! Despite its calculated-field name, it accepts exactly two numeric arguments and always returns +//! `f64`. Integer inputs are converted to `f64` before the operation. The remainder is adjusted to +//! have the divisor's sign: for example `INT_MOD(-5, 3) = 1` and `INT_MOD(5, -3) = -1`. +//! +//! A positive or negative zero divisor returns null. Null operands propagate. Other IEEE values +//! remain present: NaN produces NaN, while infinities follow the combination of floating remainder +//! and the sign adjustment. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, FloatCC, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_float_mod"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct IntModFnCall { + args: Box<[TypedExpr]>, +} + +impl FnCall for IntModFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::F64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::IntMod, + expected: target_type, + got: InferredTypeSet::F64, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::IntMod, + expected: 2, + got: args.len(), + }); + } + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types)?; + } + Ok(InferredTypeSet::F64) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::F64)); + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::F64)) + .collect::, _>>()?; + if args.iter().any(|arg| arg.return_type == VarType::None) { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type: VarType::F64, + ast: TypedExprAst::from_call(IntModFnCall { + args: args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("INT_MOD", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::F64); + let value = context.compile_expr(&self.args[0], builder)?; + let modulus = context.compile_expr(&self.args[1], builder)?; + let call = builder.ins().call( + context.native_functions().float_mod(), + &[value.value, modulus.value], + ); + let result = builder.inst_results(call)[0]; + let zero = builder.ins().f64const(0.0); + let modulus_is_zero = builder.ins().fcmp(FloatCC::Equal, modulus.value, zero); + let modulus_is_nonzero = builder.ins().bxor_imm_u(modulus_is_zero, 1); + let both_present = builder.ins().band(value.is_present, modulus.is_present); + let is_present = builder.ins().band(both_present, modulus_is_nonzero); + Ok(LoweredValue { + value: result, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, float_mod as *const u8); +} + +pub(super) fn declare_native_function( + module: &mut JITModule, + function: &mut CraneliftFunction, +) -> Result { + let mut signature = module.make_signature(); + signature + .params + .extend(std::iter::repeat_n(AbiParam::new(types::F64), 2)); + signature.returns.push(AbiParam::new(types::F64)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +extern "C" fn float_mod(value: f64, modulus: f64) -> f64 { + if modulus == 0.0 { + return 0.0; + } + let remainder = value % modulus; + if (modulus > 0.0 && remainder < 0.0) || (modulus < 0.0 && remainder > 0.0) { + remainder + modulus + } else { + remainder + } +} + +impl From for FnCallEnum { + fn from(call: IntModFnCall) -> Self { + FnCallEnum::IntMod(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable f64 values. + unsafe { compiled.call(&[]).as_f64() } + } + + #[test] + fn test_requires_two_numeric_arguments_and_returns_float() { + let expression = deserialize("(INT_MOD left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!( + inferred_types.get("left"), + Some(&InferredTypeSet::NUMERICAL) + ); + assert_eq!( + inferred_types.get("right"), + Some(&InferredTypeSet::NUMERICAL) + ); + + let expression = deserialize("(INT_MOD 1i64)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::IntMod, + expected: 2, + .. + }) + )); + } + + #[test] + fn test_remainder_has_divisor_sign_and_output_is_float() { + assert_eq!(eval("(INT_MOD 5i64 3i64)"), Some(2.0)); + assert_eq!(eval("(INT_MOD -5i64 3i64)"), Some(1.0)); + assert_eq!(eval("(INT_MOD 5i64 -3i64)"), Some(-1.0)); + assert_eq!(eval("(INT_MOD -5i64 -3i64)"), Some(-2.0)); + assert_eq!(eval("(INT_MOD 5.5f64 2f64)"), Some(1.5)); + } + + #[test] + fn test_zero_null_nan_and_infinity_edges() { + assert_eq!(eval("(INT_MOD 1f64 0f64)"), None); + assert_eq!(eval("(INT_MOD 1f64 -0f64)"), None); + assert!(eval("(INT_MOD nanf64 2f64)").unwrap().is_nan()); + assert!(eval("(INT_MOD 2f64 nanf64)").unwrap().is_nan()); + assert_eq!(eval("(INT_MOD -2f64 inff64)"), Some(f64::INFINITY)); + assert_eq!(eval("(INT_MOD none 2f64)"), None); + } + + #[test] + fn test_runtime_null_propagates() { + let expression = deserialize("(INT_MOD value modulus)").unwrap(); + let variable_types = HashMap::from([("value", VarType::I64), ("modulus", VarType::F64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(-5i64), VariableValue::some(3.0f64)]) + .as_f64() + }, + Some(1.0) + ); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(-5i64), VariableValue::none()]) + .as_f64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/is_not_null.rs b/jitexpr/src/functions/is_not_null.rs new file mode 100644 index 0000000000..e58f83a2aa --- /dev/null +++ b/jitexpr/src/functions/is_not_null.rs @@ -0,0 +1,227 @@ +//! `IS_NOT_NULL` observes nullability without propagating it. +//! +//! It accepts exactly one expression of any supported type and always returns a +//! present boolean: `false` when its argument is absent and `true` when it is +//! present. The payload is irrelevant, so present values such as `false`, zero, +//! `NaN`, and the empty string all return `true`. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct IsNotNullFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for IsNotNullFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + 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::IsNotNull, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::IsNotNull, + expected: 1, + got: args.len(), + }); + } + + crate::ast::infer_types_aux(&args[0], InferredTypeSet::ALL, inferred_types)?; + Ok(InferredTypeSet::BOOLEAN) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + + let arg = context.apply_types(&args[0], InferredTypeSet::ALL)?; + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(IsNotNullFnCall { + args: vec![arg].into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("IS_NOT_NULL", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + let arg = context.compile_expr(&self.args[0], builder)?; + let is_present = builder.ins().iconst(types::I8, 1); + Ok(LoweredValue { + value: arg.is_present, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: IsNotNullFnCall) -> Self { + FnCallEnum::IsNotNull(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types, infer_types_with_target}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval_without_args(expression: &str) -> (bool, Option) { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert!(compiled.inputs.is_empty()); + // SAFETY: The compiled expression has no inputs. + let output = unsafe { compiled.call(&[]) }; + (unsafe { output.primitive.is_present }, unsafe { + output.as_bool() + }) + } + + fn eval_variable(var_type: VarType, input: VariableValue<'_>) -> (bool, Option) { + let expression = deserialize("(IS_NOT_NULL value)").unwrap(); + let variable_types = HashMap::from([("value", var_type)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + // SAFETY: `input` is constructed with the union member matching `var_type` + // at each call site, or is absent and therefore has no active payload. + let output = unsafe { compiled.call(&[input]) }; + (unsafe { output.primitive.is_present }, unsafe { + output.as_bool() + }) + } + + #[test] + fn test_infer_types_accepts_any_argument_and_returns_bool() { + let expression = deserialize("(IS_NOT_NULL value)").unwrap(); + + let inferred_types = infer_types(&expression).unwrap(); + + assert_eq!(inferred_types.get("value"), Some(&InferredTypeSet::ALL)); + assert!(matches!( + infer_types_with_target(&expression, InferredTypeSet::F64), + Err(TypeError::WrongFunctionReturnType { + function: Function::IsNotNull, + got: InferredTypeSet::BOOLEAN, + .. + }) + )); + } + + #[test] + fn test_rejects_no_arguments() { + let expression = deserialize("(IS_NOT_NULL)").unwrap(); + + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::IsNotNull, + expected: 1, + got: 0, + }) + )); + } + + #[test] + fn test_rejects_more_than_one_argument() { + let expression = deserialize("(IS_NOT_NULL value other)").unwrap(); + + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::IsNotNull, + expected: 1, + got: 2, + }) + )); + } + + #[test] + fn test_none_literal_and_missing_variable_are_false_but_present() { + let none = eval_without_args("(IS_NOT_NULL none)"); + assert!(none.0); + assert_eq!(none.1, Some(false)); + + let missing = eval_without_args("(IS_NOT_NULL missing)"); + assert!(missing.0); + assert_eq!(missing.1, Some(false)); + } + + #[test] + fn test_absent_runtime_values_of_every_type_are_false() { + for var_type in [ + VarType::Bool, + VarType::F64, + VarType::U64, + VarType::I64, + VarType::Str, + ] { + let output = eval_variable(var_type, VariableValue::none()); + assert!(output.0, "type: {var_type:?}"); + assert_eq!(output.1, Some(false), "type: {var_type:?}"); + } + } + + #[test] + fn test_present_edge_values_are_true() { + for (var_type, input) in [ + (VarType::Bool, VariableValue::some(false)), + (VarType::F64, VariableValue::some(f64::NAN)), + (VarType::U64, VariableValue::some(0u64)), + (VarType::I64, VariableValue::some(i64::MIN)), + ] { + let output = eval_variable(var_type, input); + assert_eq!(output.1, Some(true), "type: {var_type:?}"); + } + + let output = eval_variable(VarType::Str, VariableValue::some("")); + assert_eq!(output.1, Some(true)); + } + + #[test] + fn test_observes_nested_expression_nullability() { + assert_eq!( + eval_without_args(r#"(IS_NOT_NULL (REGEXP_EXTRACT "b" "(a*)b" 1u64))"#).1, + Some(true) + ); + assert_eq!( + eval_without_args(r#"(IS_NOT_NULL (REGEXP_EXTRACT "b" "(a+)" 1u64))"#).1, + Some(false) + ); + } +} diff --git a/jitexpr/src/functions/is_null.rs b/jitexpr/src/functions/is_null.rs new file mode 100644 index 0000000000..4f3ef66574 --- /dev/null +++ b/jitexpr/src/functions/is_null.rs @@ -0,0 +1,175 @@ +//! `IS_NULL` observes absence without propagating it. +//! +//! It accepts exactly one expression of any supported type and always returns a present boolean: +//! `true` when its argument is absent and `false` when it is present. Payload values such as +//! `false`, zero, `NaN`, and the empty string are present and therefore return `false`. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct IsNullFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for IsNullFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + 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::IsNull, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::IsNull, + expected: 1, + got: args.len(), + }); + } + + crate::ast::infer_types_aux(&args[0], InferredTypeSet::ALL, inferred_types)?; + Ok(InferredTypeSet::BOOLEAN) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + + let arg = context.apply_types(&args[0], InferredTypeSet::ALL)?; + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(IsNullFnCall { + args: vec![arg].into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("IS_NULL", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + let arg = context.compile_expr(&self.args[0], builder)?; + let value = builder.ins().bxor_imm_u(arg.is_present, 1); + Ok(LoweredValue { + value, + is_present: builder.ins().iconst(types::I8, 1), + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: IsNullFnCall) -> Self { + FnCallEnum::IsNull(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no runtime inputs and return booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_infer_types_accepts_any_single_argument() { + let expression = deserialize("(IS_NULL value)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + + assert_eq!(inferred_types.get("value"), Some(&InferredTypeSet::ALL)); + } + + #[test] + fn test_rejects_wrong_arity() { + for expression in ["(IS_NULL)", "(IS_NULL value other)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::IsNull, + expected: 1, + .. + }) + )); + } + } + + #[test] + fn test_absent_values_are_true() { + assert_eq!(eval("(IS_NULL none)"), Some(true)); + assert_eq!(eval("(IS_NULL missing)"), Some(true)); + assert_eq!( + eval(r#"(IS_NULL (REGEXP_EXTRACT "b" "(a+)" 1u64))"#), + Some(true) + ); + } + + #[test] + fn test_present_edge_values_are_false() { + for expression in [ + "(IS_NULL false)", + "(IS_NULL 0i64)", + "(IS_NULL nanf64)", + r#"(IS_NULL "")"#, + r#"(IS_NULL (REGEXP_EXTRACT "b" "(a*)b" 1u64))"#, + ] { + assert_eq!(eval(expression), Some(false), "expression: {expression}"); + } + } + + #[test] + fn test_runtime_absent_variable_is_true_and_present_variable_is_false() { + let expression = deserialize("(IS_NULL value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::I64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_bool() }, + Some(true) + ); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(0i64)]).as_bool() }, + Some(false) + ); + } +} diff --git a/jitexpr/src/functions/left.rs b/jitexpr/src/functions/left.rs new file mode 100644 index 0000000000..c3be387a68 --- /dev/null +++ b/jitexpr/src/functions/left.rs @@ -0,0 +1,201 @@ +//! `LEFT(input, length)` returns the first `length` UTF-8 bytes of a string. +//! +//! The length must be an integer constant and is stored as a `usize`. A length greater than the +//! input byte length returns the whole string, zero returns a present empty string, and null input +//! propagates. Negative lengths and lengths that split a UTF-8 code point return null. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{InstBuilder, IntCC}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct LeftFnCall { + input: Box, + length: usize, +} + +fn constant_length(expression: &UntypedExpr) -> Result, super::InvalidFunctionCall> { + let UntypedExpr::Literal(literal) = expression else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::I64, + }); + }; + if !literal.is_none() && !literal.types().contains(VarType::I64) { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::I64, + }); + } + Ok(match literal { + Literal::I64(value) => usize::try_from(*value).ok(), + Literal::U64(value) => usize::try_from(*value).ok(), + Literal::F64(value) => usize::try_from(*value as i64).ok(), + Literal::None => None, + Literal::Bool(_) | Literal::String(_) => None, + }) +} + +impl FnCall for LeftFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::I64, |literal| { + literal.is_none() || literal.types().contains(VarType::I64) + }) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Left, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Left, + expected: 2, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::I64, inferred_types)?; + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input = context.apply_types(&args[0], InferredTypeSet::STRING)?; + let Some(length) = constant_length(&args[1])? else { + return Ok(TypedExpr::none()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(LeftFnCall { + input: Box::new(input), + length, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "LEFT {} {}u64", self.input, self.length) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let input = context.compile_expr(&self.input, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, null); + let length = builder + .ins() + .iconst(context.pointer_type(), self.length as i64); + let call = builder.ins().call( + context.native_functions().substring(), + &[input_ptr, input.string_len, null, length], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let native_succeeded = builder.ins().icmp_imm_u(IntCC::NotEqual, value, 0); + let is_present = builder.ins().band(input.is_present, native_succeeded); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +impl From for FnCallEnum { + fn from(call: LeftFnCall) -> Self { + FnCallEnum::Left(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable string. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_signature_and_byte_lengths() { + for expression in ["(LEFT \"abc\")", "(LEFT \"abc\" 1i64 2i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Left, + expected: 2, + .. + }) + )); + } + + assert_eq!(eval("(LEFT \"abcdef\" 3i64)"), Some("abc".into())); + assert_eq!(eval("(LEFT \"éclair\" 2i64)"), Some("é".into())); + assert_eq!(eval("(LEFT \"éclair\" 1i64)"), None); + } + + #[test] + fn test_zero_clamping_and_invalid_length() { + assert_eq!(eval("(LEFT \"abc\" 0i64)"), Some(String::new())); + assert_eq!(eval("(LEFT \"abc\" 99i64)"), Some("abc".into())); + assert_eq!(eval("(LEFT \"abc\" -1i64)"), None); + } + + #[test] + fn test_runtime_null() { + let expression = deserialize("(LEFT value 2i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable string argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("abc")]).as_str() }, + Some("ab") + ); + // SAFETY: The compiled expression expects one nullable string argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + } +} diff --git a/jitexpr/src/functions/lower.rs b/jitexpr/src/functions/lower.rs new file mode 100644 index 0000000000..be3a5894ca --- /dev/null +++ b/jitexpr/src/functions/lower.rs @@ -0,0 +1,335 @@ +//! `LOWER` constructs the Unicode-lowercase form of a string. +//! +//! It accepts exactly one string and returns a newly allocated string without mutating its input. +//! It applies one-to-one Unicode simple case mappings. In particular, `İ` maps to plain `i` rather +//! than expanding to `i` followed by a combining dot. Null input returns null, while empty input +//! remains present. +//! +//! Constructed bytes live in the caller's fixed-capacity string arena; arena exhaustion returns +//! null. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, StringArena, TypedExpr, + TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_string_lowercase"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct LowerFnCall { + arg: Box, +} + +impl FnCall for LowerFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Lower, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Lower, + expected: 1, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let arg = context.apply_types(&args[0], InferredTypeSet::STRING)?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + debug_assert_eq!(arg.return_type, VarType::Str); + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(LowerFnCall { arg: Box::new(arg) }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("LOWER", std::iter::once(self.arg.as_ref()), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let arg = context.compile_expr(&self.arg, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(arg.is_present, arg.value, null); + let string_lowercase = context.native_functions().string_lowercase(); + let string_arena_ptr = context.string_arena_ptr(builder); + let call = builder.ins().call( + string_lowercase, + &[input_ptr, arg.string_len, string_arena_ptr], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder + .ins() + .icmp_imm_u(cranelift::prelude::IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, string_lowercase 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(types::I64), + AbiParam::new(pointer_type), + ]); + signature.returns.push(AbiParam::new(pointer_type)); + signature.returns.push(AbiParam::new(types::I64)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } +} + +unsafe extern "C" fn string_lowercase( + input_ptr: *const u8, + input_len: usize, + string_arena: *mut StringArena, +) -> RawStr { + if input_ptr.is_null() || string_arena.is_null() { + return RawStr::none(); + } + + let output_len = { + // SAFETY: Generated code passes a live UTF-8 string and its exact byte + // length for every present string value. + let input = unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) + }; + let mut output_len = 0usize; + for character in input.chars() { + let lowercase = simple_lowercase(character); + let Some(next_len) = output_len.checked_add(lowercase.len_utf8()) else { + return RawStr::none(); + }; + output_len = next_len; + } + output_len + }; + + // The input reference above is no longer live. This matters for nested + // LOWER calls whose input points into an earlier, disjoint arena allocation. + // SAFETY: The caller exclusively borrows and passes this arena for the call. + let Some(output_ptr) = (unsafe { &mut *string_arena }).allocate(output_len) else { + return RawStr::none(); + }; + + // SAFETY: Same input contract as above. The fixed arena never reallocates, + // and its new output range starts after any arena-backed input range. + let input = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) }; + let mut written = 0usize; + for character in input.chars() { + let lowercase = simple_lowercase(character); + let mut encoded = [0; 4]; + let encoded = lowercase.encode_utf8(&mut encoded).as_bytes(); + // SAFETY: output_len was computed from this exact transformation, + // and StringArena reserved that many bytes. + unsafe { + std::ptr::copy_nonoverlapping(encoded.as_ptr(), output_ptr.add(written), encoded.len()); + } + written += encoded.len(); + } + debug_assert_eq!(written, output_len); + RawStr { + ptr: output_ptr, + len: output_len, + } +} + +/// Use the first code point from a full lowercase mapping to keep the result one-to-one. +fn simple_lowercase(character: char) -> char { + character.to_lowercase().next().unwrap_or(character) +} + +impl From for FnCallEnum { + fn from(call: LowerFnCall) -> Self { + FnCallEnum::Lower(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::{STRING_ARENA_CAPACITY, compile}; + use crate::types::VariableValue; + + #[test] + fn test_infer_types_constrains_argument_to_string() { + let expression = deserialize("(LOWER value)").unwrap(); + + let inferred_types = infer_types(&expression).unwrap(); + + assert_eq!(expression.to_string(), "(LOWER value)"); + assert_eq!(inferred_types.get("value"), Some(&InferredTypeSet::STRING)); + } + + #[test] + fn test_infer_types_requires_one_argument() { + for expression in ["(LOWER)", "(LOWER one two)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Lower, + expected: 1, + .. + }) + )); + } + } + + #[test] + fn test_compile_lowercases_unicode_without_mutating_input() { + let expression = deserialize("(LOWER value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input_string = String::from("CAFÉ İSTANBUL"); + let input = [VariableValue::some(input_string.as_str())]; + + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, Some("café istanbul")); + assert_eq!(input_string, "CAFÉ İSTANBUL"); + } + + #[test] + fn test_compile_propagates_none_and_preserves_empty_string() { + let expression = deserialize("(LOWER value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + let none = unsafe { compiled.call(&[VariableValue::none()]) }; + assert_eq!(unsafe { none.as_str() }, None); + + let empty = unsafe { compiled.call(&[VariableValue::some("")]) }; + assert_eq!(unsafe { empty.as_str() }, Some("")); + } + + #[test] + fn test_nested_lower_uses_stable_arena_allocations() { + let expression = deserialize("(LOWER (LOWER value))").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("HeLLo")]; + + assert_eq!(unsafe { compiled.call(&input).as_str() }, Some("hello")); + } + + #[test] + fn test_multiple_lower_calls_keep_previous_allocations_valid() { + let expression = deserialize("(EQ (LOWER left) (LOWER right))").unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("FiRsT"), VariableValue::some("SeCoNd")]; + + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_bool() }, Some(false)); + assert_eq!(compiled.string_arena.used_bytes(), 11); + } + + #[test] + fn test_arena_exhaustion_returns_none_without_advancing_cursor() { + let expression = deserialize("(LOWER value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let too_large = "A".repeat(STRING_ARENA_CAPACITY + 1); + let input = [VariableValue::some(too_large.as_str())]; + + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, None); + assert_eq!(compiled.string_arena.used_bytes(), 0); + } + + #[test] + fn test_arena_cursor_is_cleared_before_each_call() { + let expression = deserialize("(LOWER value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let full = "A".repeat(STRING_ARENA_CAPACITY); + + { + let output = unsafe { compiled.call(&[VariableValue::some(full.as_str())]) }; + assert_eq!( + unsafe { output.as_str() }.unwrap().len(), + STRING_ARENA_CAPACITY + ); + } + assert_eq!(compiled.string_arena.used_bytes(), STRING_ARENA_CAPACITY); + + { + let output = unsafe { compiled.call(&[VariableValue::some("ABC")]) }; + assert_eq!(unsafe { output.as_str() }, Some("abc")); + } + assert_eq!(compiled.string_arena.used_bytes(), 3); + } +} diff --git a/jitexpr/src/functions/lt.rs b/jitexpr/src/functions/lt.rs new file mode 100644 index 0000000000..8539303a89 --- /dev/null +++ b/jitexpr/src/functions/lt.rs @@ -0,0 +1,146 @@ +//! `LT` tests whether one value is less than another. +//! +//! It accepts exactly two operands. Ordered operands are strings or numbers; booleans are rejected. +//! Strings use lexicographic UTF-8 ordering. Numeric comparisons support `i64`, `u64`, and `f64` +//! combinations without converting large integers through a lossy `f64`. IEEE unordered +//! comparisons involving NaN return `false`. +//! +//! Null propagation is strict: if either operand is absent, the result is absent. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +use super::comparison::{self, OrderedComparison}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct LtFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for LtFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + comparison::infer_types(Function::Lt, args, target_type, inferred_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(LtFnCall { + args: comparison::apply_types(args, context)?, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("LT", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + comparison::lower(&self.args, OrderedComparison::LessThan, context, builder) + } +} + +impl From for FnCallEnum { + fn from(call: LtFnCall) -> Self { + FnCallEnum::Lt(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_requires_two_ordered_arguments() { + let expression = deserialize("(LT left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!( + inferred_types.get("left"), + Some(&InferredTypeSet { + string: true, + i64: true, + u64: true, + f64: true, + boolean: false, + }) + ); + + let expression = deserialize("(LT 1i64)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Lt, + expected: 2, + .. + }) + )); + } + + #[test] + fn test_ordering_and_null_semantics() { + assert_eq!(eval(r#"(LT "alpha" "beta")"#), Some(true)); + assert_eq!(eval("(LT -1i64 0u64)"), Some(true)); + assert_eq!( + eval("(LT 9007199254740993u64 9007199254740992f64)"), + Some(false) + ); + assert_eq!(eval("(LT nanf64 0i64)"), Some(false)); + assert_eq!(eval("(LT 0i64 nanf64)"), Some(false)); + assert_eq!(eval("(LT none 1i64)"), None); + } + + #[test] + fn test_runtime_null_propagates() { + let expression = deserialize("(LT left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::I64), ("right", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::none(), VariableValue::some(0u64)]) + .as_bool() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/lt_eq.rs b/jitexpr/src/functions/lt_eq.rs new file mode 100644 index 0000000000..723f69e2ce --- /dev/null +++ b/jitexpr/src/functions/lt_eq.rs @@ -0,0 +1,145 @@ +//! `LT_EQ` tests whether one value is less than or equal to another. +//! +//! It accepts exactly two operands. Ordered operands are strings or numbers; booleans are rejected. +//! Strings use lexicographic UTF-8 ordering. Numeric comparisons support `i64`, `u64`, and `f64` +//! combinations without converting large integers through a lossy `f64`. IEEE unordered +//! comparisons involving NaN return `false`, including `NaN <= NaN`. +//! +//! Null propagation is strict: if either operand is absent, the result is absent. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +use super::comparison::{self, OrderedComparison}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct LtEqFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for LtEqFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + comparison::infer_types(Function::LtEq, args, target_type, inferred_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(LtEqFnCall { + args: comparison::apply_types(args, context)?, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("LT_EQ", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + comparison::lower( + &self.args, + OrderedComparison::LessThanOrEqual, + context, + builder, + ) + } +} + +impl From for FnCallEnum { + fn from(call: LtEqFnCall) -> Self { + FnCallEnum::LtEq(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_requires_two_ordered_arguments() { + let expression = deserialize("(LT_EQ left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + let left = inferred_types.get("left").unwrap(); + assert!(left.string && left.i64 && left.u64 && left.f64 && !left.boolean); + + let expression = deserialize("(LT_EQ)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::LtEq, + expected: 2, + .. + }) + )); + } + + #[test] + fn test_ordering_equality_nan_and_null() { + assert_eq!(eval(r#"(LT_EQ "same" "same")"#), Some(true)); + assert_eq!(eval(r#"(LT_EQ "beta" "alpha")"#), Some(false)); + assert_eq!(eval("(LT_EQ -1i64 0u64)"), Some(true)); + assert_eq!( + eval("(LT_EQ 9007199254740992f64 9007199254740993u64)"), + Some(true) + ); + assert_eq!(eval("(LT_EQ nanf64 nanf64)"), Some(false)); + assert_eq!(eval("(LT_EQ nanf64 0u64)"), Some(false)); + assert_eq!(eval("(LT_EQ 0u64 nanf64)"), Some(false)); + assert_eq!(eval("(LT_EQ none 1i64)"), None); + } + + #[test] + fn test_runtime_null_propagates() { + let expression = deserialize("(LT_EQ left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::F64), ("right", VarType::I64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(1.0f64), VariableValue::none()]) + .as_bool() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/max.rs b/jitexpr/src/functions/max.rs new file mode 100644 index 0000000000..71ffefcc81 --- /dev/null +++ b/jitexpr/src/functions/max.rs @@ -0,0 +1,210 @@ +//! `MAX` returns the greatest value among one or more numeric arguments. +//! +//! This is an expression function, not an aggregate. All arguments are coerced to one numeric type +//! and the result has that type. Any null argument makes the result null. The float implementation +//! starts at negative infinity and replaces it only on `>`, so NaNs are ignored; if every argument +//! is NaN, the result is negative infinity. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{FloatCC, InstBuilder, IntCC, types}; + +use super::add::{is_numerical, select_return_type, with_float_fallback}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct MaxFnCall { + args: Box<[TypedExpr]>, +} + +impl FnCall for MaxFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::AtLeast(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::NUMERICAL).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Max, + expected: target_type, + got: InferredTypeSet::NUMERICAL, + }); + } + if args.is_empty() { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Max, + expected: 1, + got: 0, + }); + } + let mut return_types = InferredTypeSet::NUMERICAL; + for arg in args { + return_types = return_types.intersect(crate::ast::infer_types_aux( + arg, + InferredTypeSet::NUMERICAL, + inferred_types, + )?); + } + let result = with_float_fallback(return_types).intersect(target_type); + if result.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Max, + expected: target_type, + got: return_types, + }); + } + Ok(result) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let mut return_types = InferredTypeSet::NUMERICAL.intersect(target_type_set); + for arg in args { + return_types = return_types.intersect(crate::ast::infer_type_with_variable_types( + arg, + InferredTypeSet::NUMERICAL, + context.variable_types(), + )?); + } + let return_type = select_return_type(with_float_fallback(return_types)); + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::singleton(return_type))) + .collect::, _>>()?; + if args.iter().any(|arg| !is_numerical(arg.return_type)) { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(MaxFnCall { + args: args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("MAX", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let mut value = match return_type { + VarType::I64 => builder.ins().iconst(types::I64, i64::MIN), + VarType::U64 => builder.ins().iconst(types::I64, 0), + VarType::F64 => builder.ins().f64const(f64::NEG_INFINITY), + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Max, + return_type, + }); + } + }; + let mut is_present = builder.ins().iconst(types::I8, 1); + for arg in &self.args { + let arg = context.compile_expr(arg, builder)?; + let is_greater = match return_type { + VarType::I64 => builder + .ins() + .icmp(IntCC::SignedGreaterThan, arg.value, value), + VarType::U64 => builder + .ins() + .icmp(IntCC::UnsignedGreaterThan, arg.value, value), + VarType::F64 => builder.ins().fcmp(FloatCC::GreaterThan, arg.value, value), + _ => unreachable!(), + }; + value = builder.ins().select(is_greater, arg.value, value); + is_present = builder.ins().band(is_present, arg.is_present); + } + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: MaxFnCall) -> Self { + FnCallEnum::Max(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_signature_and_variadic_values() { + let expression = deserialize("(MAX a b c)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + for name in ["a", "b", "c"] { + assert_eq!(inferred.get(name), Some(&InferredTypeSet::NUMERICAL)); + } + let empty = deserialize("(MAX)").unwrap(); + assert!(matches!( + infer_types(&empty), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Max, + expected: 1, + got: 0 + }) + )); + let expression = deserialize("(MAX 7i64 -3i64 12i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, Some(12)); + let expression = deserialize("(MAX 7u64 13u64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_u64() }, Some(13)); + } + + #[test] + fn test_float_nan_and_null_behavior() { + let expression = deserialize("(MAX nanf64 3f64 -2f64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, Some(3.0)); + let expression = deserialize("(MAX nanf64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!( + unsafe { compiled.call(&[]).as_f64() }, + Some(f64::NEG_INFINITY) + ); + let expression = deserialize("(MAX left right)").unwrap(); + let mut compiled = compile( + &expression, + &HashMap::from([("left", VarType::F64), ("right", VarType::F64)]), + ) + .unwrap() + .context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(1.0f64), VariableValue::none()]) + .as_f64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/min.rs b/jitexpr/src/functions/min.rs new file mode 100644 index 0000000000..a6bf8da952 --- /dev/null +++ b/jitexpr/src/functions/min.rs @@ -0,0 +1,207 @@ +//! `MIN` returns the least value among one or more numeric arguments. +//! +//! This is an expression function, not an aggregate. All arguments are coerced to one numeric type +//! and the result has that type. Any null argument makes the result null. The float implementation +//! starts at positive infinity and replaces it only on `<`, so NaNs are ignored; if every argument +//! is NaN, the result is positive infinity. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{FloatCC, InstBuilder, IntCC, types}; + +use super::add::{is_numerical, select_return_type, with_float_fallback}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct MinFnCall { + args: Box<[TypedExpr]>, +} + +impl FnCall for MinFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::AtLeast(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::NUMERICAL).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Min, + expected: target_type, + got: InferredTypeSet::NUMERICAL, + }); + } + if args.is_empty() { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Min, + expected: 1, + got: 0, + }); + } + let mut return_types = InferredTypeSet::NUMERICAL; + for arg in args { + return_types = return_types.intersect(crate::ast::infer_types_aux( + arg, + InferredTypeSet::NUMERICAL, + inferred_types, + )?); + } + let result = with_float_fallback(return_types).intersect(target_type); + if result.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Min, + expected: target_type, + got: return_types, + }); + } + Ok(result) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let mut return_types = InferredTypeSet::NUMERICAL.intersect(target_type_set); + for arg in args { + return_types = return_types.intersect(crate::ast::infer_type_with_variable_types( + arg, + InferredTypeSet::NUMERICAL, + context.variable_types(), + )?); + } + let return_type = select_return_type(with_float_fallback(return_types)); + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::singleton(return_type))) + .collect::, _>>()?; + if args.iter().any(|arg| !is_numerical(arg.return_type)) { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(MinFnCall { + args: args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("MIN", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let mut value = match return_type { + VarType::I64 => builder.ins().iconst(types::I64, i64::MAX), + VarType::U64 => builder.ins().iconst(types::I64, -1), + VarType::F64 => builder.ins().f64const(f64::INFINITY), + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Min, + return_type, + }); + } + }; + let mut is_present = builder.ins().iconst(types::I8, 1); + for arg in &self.args { + let arg = context.compile_expr(arg, builder)?; + let is_less = match return_type { + VarType::I64 => builder.ins().icmp(IntCC::SignedLessThan, arg.value, value), + VarType::U64 => builder + .ins() + .icmp(IntCC::UnsignedLessThan, arg.value, value), + VarType::F64 => builder.ins().fcmp(FloatCC::LessThan, arg.value, value), + _ => unreachable!(), + }; + value = builder.ins().select(is_less, arg.value, value); + is_present = builder.ins().band(is_present, arg.is_present); + } + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: MinFnCall) -> Self { + FnCallEnum::Min(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_signature_and_variadic_values() { + let expression = deserialize("(MIN a b c)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + for name in ["a", "b", "c"] { + assert_eq!(inferred.get(name), Some(&InferredTypeSet::NUMERICAL)); + } + let empty = deserialize("(MIN)").unwrap(); + assert!(matches!( + infer_types(&empty), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Min, + expected: 1, + got: 0 + }) + )); + + let expression = deserialize("(MIN 7i64 -3i64 2i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, Some(-3)); + let expression = deserialize("(MIN 7u64 3u64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_u64() }, Some(3)); + } + + #[test] + fn test_float_nan_and_null_behavior() { + let expression = deserialize("(MIN nanf64 3f64 -2f64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, Some(-2.0)); + let expression = deserialize("(MIN nanf64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, Some(f64::INFINITY)); + + let expression = deserialize("(MIN left right)").unwrap(); + let mut compiled = compile( + &expression, + &HashMap::from([("left", VarType::I64), ("right", VarType::I64)]), + ) + .unwrap() + .context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(1i64), VariableValue::none()]) + .as_i64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/mod.rs b/jitexpr/src/functions/mod.rs new file mode 100644 index 0000000000..8a56a47546 --- /dev/null +++ b/jitexpr/src/functions/mod.rs @@ -0,0 +1,875 @@ +mod abs; +mod add; +mod and; +mod ceil; +mod comparison; +mod concat; +mod divide; +mod eq; +mod floor; +mod gt; +mod gt_eq; +mod if_fn; +mod int_mod; +mod is_not_null; +mod is_null; +mod left; +mod lower; +mod lt; +mod lt_eq; +mod max; +mod min; +mod multiply; +mod native_function; +mod neq; +mod not; +mod or; +mod pow; +mod regexp_extract; +mod regexp_like; +mod right; +mod round; +mod split_after; +mod split_before; +mod sqrt; +mod substring; +mod substring_count; +mod subtract; +mod text_join; +mod trim; +mod upper; + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +pub(crate) use self::abs::AbsFnCall; +pub(crate) use self::add::AddFnCall; +pub(crate) use self::and::AndFnCall; +pub(crate) use self::ceil::CeilFnCall; +pub(crate) use self::concat::ConcatFnCall; +pub(crate) use self::divide::DivideFnCall; +pub(crate) use self::eq::EqFnCall; +pub(crate) use self::floor::FloorFnCall; +pub(crate) use self::gt::GtFnCall; +pub(crate) use self::gt_eq::GtEqFnCall; +pub(crate) use self::if_fn::IfFnCall; +pub(crate) use self::int_mod::IntModFnCall; +pub(crate) use self::is_not_null::IsNotNullFnCall; +pub(crate) use self::is_null::IsNullFnCall; +pub(crate) use self::left::LeftFnCall; +pub(crate) use self::lower::LowerFnCall; +pub(crate) use self::lt::LtFnCall; +pub(crate) use self::lt_eq::LtEqFnCall; +pub(crate) use self::max::MaxFnCall; +pub(crate) use self::min::MinFnCall; +pub(crate) use self::multiply::MultiplyFnCall; +pub(crate) use self::native_function::{ + NativeFunctions, declare_native_functions, register_jit_symbols, +}; +pub(crate) use self::neq::NeqFnCall; +pub(crate) use self::not::NotFnCall; +pub(crate) use self::or::OrFnCall; +pub(crate) use self::pow::PowFnCall; +pub(crate) use self::regexp_extract::RegexpExtractFnCall; +pub(crate) use self::regexp_like::RegexpLikeFnCall; +pub(crate) use self::right::RightFnCall; +pub(crate) use self::round::RoundFnCall; +pub(crate) use self::split_after::SplitAfterFnCall; +pub(crate) use self::split_before::SplitBeforeFnCall; +pub(crate) use self::sqrt::SqrtFnCall; +pub(crate) use self::substring::SubstringFnCall; +pub(crate) use self::substring_count::SubstringCountFnCall; +pub(crate) use self::subtract::SubtractFnCall; +pub(crate) use self::text_join::TextJoinFnCall; +pub(crate) use self::trim::TrimFnCall; +pub(crate) use self::upper::UpperFnCall; +use crate::ast::{InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr}; +use crate::types::VarType; + +/// A function supported by the first expression-language milestone. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum Function { + /// Computes the absolute value of a number. + Abs, + /// Conjoins one or more booleans with strict null propagation. + And, + /// Returns the least integer greater than or equal to a number. + Ceil, + /// Joins strings using a literal delimiter and empty-value policy. + Concat, + /// Adds zero or more numerical expressions. + Add, + /// Divides two numeric arguments using floating-point arithmetic. + Divide, + /// Compares two expressions for value equality. + Eq, + /// Returns the greatest integer less than or equal to a number. + Floor, + /// Tests whether the first ordered value is greater than the second. + Gt, + /// Tests whether the first ordered value is greater than or equal to the second. + GtEq, + /// Selects one of two values using a boolean condition. + If, + /// Computes floating-point modulo with the divisor's sign. + IntMod, + /// Returns the first requested number of bytes from a string. + Left, + /// Tests whether the first ordered value is less than the second. + Lt, + /// Tests whether the first ordered value is less than or equal to the second. + LtEq, + /// Tests whether an expression produced a present value. + IsNotNull, + /// Tests whether an expression produced an absent value. + IsNull, + /// Constructs the Unicode-lowercase form of a string. + Lower, + /// Returns the greatest of one or more numbers. + Max, + /// Returns the least of one or more numbers. + Min, + /// Multiplies two numeric arguments. + Multiply, + /// Tests inequality using `NOT(EQ(...))` null semantics. + Neq, + /// Negates a boolean, treating an absent input as false. + Not, + /// Disjoins one or more booleans, remaining present if any operand is present. + Or, + /// Raises a numeric base to a numeric exponent and returns a float. + Pow, + /// Returns the floating-point square root of a number, or null for a NaN result. + Sqrt, + /// Extracts a capture group from a string using a constant regular expression. + RegexpExtract, + /// Tests whether a constant regular expression matches a string. + RegexpLike, + /// Returns the last requested number of bytes from a string. + Right, + /// Rounds a number to a constant decimal precision. + Round, + /// Returns the suffix after a selected occurrence of a literal separator. + SplitAfter, + /// Returns the prefix before a selected occurrence of a literal separator. + SplitBefore, + /// Subtracts the second numeric argument from the first. + Subtract, + /// Returns a byte-indexed string slice. + Substring, + /// Counts non-overlapping occurrences of one string in another. + SubstringCount, + /// Joins strings with the same semantics as `CONCAT`. + TextJoin, + /// Removes a whole delimiter from selected ends of a string. + Trim, + /// Converts a string to Unicode uppercase. + Upper, +} + +impl Function { + pub(crate) fn call(self, args: Vec) -> Result { + match self { + Function::Abs => ::validate_args(&args)?, + Function::And => ::validate_args(&args)?, + Function::Ceil => ::validate_args(&args)?, + Function::Concat => ::validate_args(&args)?, + Function::Add => ::validate_args(&args)?, + Function::Divide => ::validate_args(&args)?, + Function::Eq => ::validate_args(&args)?, + Function::Floor => ::validate_args(&args)?, + Function::Gt => ::validate_args(&args)?, + Function::GtEq => ::validate_args(&args)?, + Function::If => ::validate_args(&args)?, + Function::IntMod => ::validate_args(&args)?, + Function::Left => ::validate_args(&args)?, + Function::Lt => ::validate_args(&args)?, + Function::LtEq => ::validate_args(&args)?, + Function::IsNotNull => ::validate_args(&args)?, + Function::IsNull => ::validate_args(&args)?, + Function::Lower => ::validate_args(&args)?, + Function::Max => ::validate_args(&args)?, + Function::Min => ::validate_args(&args)?, + Function::Multiply => ::validate_args(&args)?, + Function::Neq => ::validate_args(&args)?, + Function::Not => ::validate_args(&args)?, + Function::Or => ::validate_args(&args)?, + Function::Pow => ::validate_args(&args)?, + Function::Sqrt => ::validate_args(&args)?, + Function::RegexpExtract => ::validate_args(&args)?, + Function::RegexpLike => ::validate_args(&args)?, + Function::Right => ::validate_args(&args)?, + Function::Round => ::validate_args(&args)?, + Function::SplitAfter => ::validate_args(&args)?, + Function::SplitBefore => ::validate_args(&args)?, + Function::Subtract => ::validate_args(&args)?, + Function::Substring => ::validate_args(&args)?, + Function::SubstringCount => ::validate_args(&args)?, + Function::TextJoin => ::validate_args(&args)?, + Function::Trim => ::validate_args(&args)?, + Function::Upper => ::validate_args(&args)?, + } + + Ok(UntypedExpr::Call { + function: self, + args, + }) + } + + pub(crate) fn call_with_types( + self, + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + match self { + Function::Abs => ::call_with_types(args, target_type_set, context), + Function::And => ::call_with_types(args, target_type_set, context), + Function::Ceil => { + ::call_with_types(args, target_type_set, context) + } + Function::Concat => { + ::call_with_types(args, target_type_set, context) + } + Function::Add => ::call_with_types(args, target_type_set, context), + Function::Divide => { + ::call_with_types(args, target_type_set, context) + } + Function::Eq => ::call_with_types(args, target_type_set, context), + Function::Floor => { + ::call_with_types(args, target_type_set, context) + } + Function::Gt => ::call_with_types(args, target_type_set, context), + Function::GtEq => { + ::call_with_types(args, target_type_set, context) + } + Function::If => ::call_with_types(args, target_type_set, context), + Function::IntMod => { + ::call_with_types(args, target_type_set, context) + } + Function::Left => { + ::call_with_types(args, target_type_set, context) + } + Function::Lt => ::call_with_types(args, target_type_set, context), + Function::LtEq => { + ::call_with_types(args, target_type_set, context) + } + Function::IsNotNull => { + ::call_with_types(args, target_type_set, context) + } + Function::IsNull => { + ::call_with_types(args, target_type_set, context) + } + Function::Lower => { + ::call_with_types(args, target_type_set, context) + } + Function::Max => ::call_with_types(args, target_type_set, context), + Function::Min => ::call_with_types(args, target_type_set, context), + Function::Multiply => { + ::call_with_types(args, target_type_set, context) + } + Function::Neq => ::call_with_types(args, target_type_set, context), + Function::Not => ::call_with_types(args, target_type_set, context), + Function::Or => ::call_with_types(args, target_type_set, context), + Function::Pow => ::call_with_types(args, target_type_set, context), + Function::Sqrt => { + ::call_with_types(args, target_type_set, context) + } + Function::RegexpExtract => { + ::call_with_types(args, target_type_set, context) + } + Function::RegexpLike => { + ::call_with_types(args, target_type_set, context) + } + Function::Right => { + ::call_with_types(args, target_type_set, context) + } + Function::Round => { + ::call_with_types(args, target_type_set, context) + } + Function::SplitAfter => { + ::call_with_types(args, target_type_set, context) + } + Function::SplitBefore => { + ::call_with_types(args, target_type_set, context) + } + Function::Subtract => { + ::call_with_types(args, target_type_set, context) + } + Function::Substring => { + ::call_with_types(args, target_type_set, context) + } + Function::SubstringCount => { + ::call_with_types(args, target_type_set, context) + } + Function::TextJoin => { + ::call_with_types(args, target_type_set, context) + } + Function::Trim => { + ::call_with_types(args, target_type_set, context) + } + Function::Upper => { + ::call_with_types(args, target_type_set, context) + } + } + } + + pub(crate) fn infer_types<'a>( + self, + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + match self { + Function::Abs => ::infer_types(args, target_type, inferred_types), + Function::And => ::infer_types(args, target_type, inferred_types), + Function::Ceil => { + ::infer_types(args, target_type, inferred_types) + } + Function::Concat => { + ::infer_types(args, target_type, inferred_types) + } + Function::Add => ::infer_types(args, target_type, inferred_types), + Function::Divide => { + ::infer_types(args, target_type, inferred_types) + } + Function::Eq => ::infer_types(args, target_type, inferred_types), + Function::Floor => { + ::infer_types(args, target_type, inferred_types) + } + Function::Gt => ::infer_types(args, target_type, inferred_types), + Function::GtEq => { + ::infer_types(args, target_type, inferred_types) + } + Function::If => ::infer_types(args, target_type, inferred_types), + Function::IntMod => { + ::infer_types(args, target_type, inferred_types) + } + Function::Left => { + ::infer_types(args, target_type, inferred_types) + } + Function::Lt => ::infer_types(args, target_type, inferred_types), + Function::LtEq => { + ::infer_types(args, target_type, inferred_types) + } + Function::IsNotNull => { + ::infer_types(args, target_type, inferred_types) + } + Function::IsNull => { + ::infer_types(args, target_type, inferred_types) + } + Function::Lower => { + ::infer_types(args, target_type, inferred_types) + } + Function::Max => ::infer_types(args, target_type, inferred_types), + Function::Min => ::infer_types(args, target_type, inferred_types), + Function::Multiply => { + ::infer_types(args, target_type, inferred_types) + } + Function::Neq => ::infer_types(args, target_type, inferred_types), + Function::Not => ::infer_types(args, target_type, inferred_types), + Function::Or => ::infer_types(args, target_type, inferred_types), + Function::Pow => ::infer_types(args, target_type, inferred_types), + Function::Sqrt => { + ::infer_types(args, target_type, inferred_types) + } + Function::RegexpExtract => { + ::infer_types(args, target_type, inferred_types) + } + Function::RegexpLike => { + ::infer_types(args, target_type, inferred_types) + } + Function::Right => { + ::infer_types(args, target_type, inferred_types) + } + Function::Round => { + ::infer_types(args, target_type, inferred_types) + } + Function::SplitAfter => { + ::infer_types(args, target_type, inferred_types) + } + Function::SplitBefore => { + ::infer_types(args, target_type, inferred_types) + } + Function::Subtract => { + ::infer_types(args, target_type, inferred_types) + } + Function::Substring => { + ::infer_types(args, target_type, inferred_types) + } + Function::SubstringCount => { + ::infer_types(args, target_type, inferred_types) + } + Function::TextJoin => { + ::infer_types(args, target_type, inferred_types) + } + Function::Trim => { + ::infer_types(args, target_type, inferred_types) + } + Function::Upper => { + ::infer_types(args, target_type, inferred_types) + } + } + } + + pub fn call_untyped_expr(self, args: Vec) -> UntypedExpr { + UntypedExpr::Call { + function: self, + args, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum FnCallEnum { + Abs(AbsFnCall), + And(AndFnCall), + Ceil(CeilFnCall), + Concat(ConcatFnCall), + Add(AddFnCall), + Divide(DivideFnCall), + Eq(EqFnCall), + Floor(FloorFnCall), + Gt(GtFnCall), + GtEq(GtEqFnCall), + If(IfFnCall), + IntMod(IntModFnCall), + Left(LeftFnCall), + Lt(LtFnCall), + LtEq(LtEqFnCall), + IsNull(IsNullFnCall), + IsNotNull(IsNotNullFnCall), + Lower(LowerFnCall), + Max(MaxFnCall), + Min(MinFnCall), + Multiply(MultiplyFnCall), + Neq(NeqFnCall), + Not(NotFnCall), + Or(OrFnCall), + Pow(PowFnCall), + Sqrt(SqrtFnCall), + RegexpExtract(RegexpExtractFnCall), + RegexpLike(RegexpLikeFnCall), + Right(RightFnCall), + Round(RoundFnCall), + SplitAfter(SplitAfterFnCall), + SplitBefore(SplitBeforeFnCall), + Subtract(SubtractFnCall), + Substring(SubstringFnCall), + SubstringCount(SubstringCountFnCall), + TextJoin(TextJoinFnCall), + Trim(TrimFnCall), + Upper(UpperFnCall), +} + +impl FnCallEnum { + pub(crate) fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + FnCallEnum::Abs(call) => call.serialize(formatter), + FnCallEnum::And(call) => call.serialize(formatter), + FnCallEnum::Ceil(call) => call.serialize(formatter), + FnCallEnum::Concat(call) => call.serialize(formatter), + FnCallEnum::Add(call) => call.serialize(formatter), + FnCallEnum::Divide(call) => call.serialize(formatter), + FnCallEnum::Eq(call) => call.serialize(formatter), + FnCallEnum::Floor(call) => call.serialize(formatter), + FnCallEnum::Gt(call) => call.serialize(formatter), + FnCallEnum::GtEq(call) => call.serialize(formatter), + FnCallEnum::If(call) => call.serialize(formatter), + FnCallEnum::IntMod(call) => call.serialize(formatter), + FnCallEnum::Left(call) => call.serialize(formatter), + FnCallEnum::Lt(call) => call.serialize(formatter), + FnCallEnum::LtEq(call) => call.serialize(formatter), + FnCallEnum::IsNull(call) => call.serialize(formatter), + FnCallEnum::IsNotNull(call) => call.serialize(formatter), + FnCallEnum::Lower(call) => call.serialize(formatter), + FnCallEnum::Max(call) => call.serialize(formatter), + FnCallEnum::Min(call) => call.serialize(formatter), + FnCallEnum::Multiply(call) => call.serialize(formatter), + FnCallEnum::Neq(call) => call.serialize(formatter), + FnCallEnum::Not(call) => call.serialize(formatter), + FnCallEnum::Or(call) => call.serialize(formatter), + FnCallEnum::Pow(call) => call.serialize(formatter), + FnCallEnum::Sqrt(call) => call.serialize(formatter), + FnCallEnum::RegexpExtract(call) => call.serialize(formatter), + FnCallEnum::RegexpLike(call) => call.serialize(formatter), + FnCallEnum::Right(call) => call.serialize(formatter), + FnCallEnum::Round(call) => call.serialize(formatter), + FnCallEnum::SplitAfter(call) => call.serialize(formatter), + FnCallEnum::SplitBefore(call) => call.serialize(formatter), + FnCallEnum::Subtract(call) => call.serialize(formatter), + FnCallEnum::Substring(call) => call.serialize(formatter), + FnCallEnum::SubstringCount(call) => call.serialize(formatter), + FnCallEnum::TextJoin(call) => call.serialize(formatter), + FnCallEnum::Trim(call) => call.serialize(formatter), + FnCallEnum::Upper(call) => call.serialize(formatter), + } + } + + pub(crate) fn args_mut(&mut self) -> &mut [TypedExpr] { + match self { + FnCallEnum::Abs(call) => call.args_mut(), + FnCallEnum::And(call) => call.args_mut(), + FnCallEnum::Ceil(call) => call.args_mut(), + FnCallEnum::Concat(call) => call.args_mut(), + FnCallEnum::Add(call) => call.args_mut(), + FnCallEnum::Divide(call) => call.args_mut(), + FnCallEnum::Eq(call) => call.args_mut(), + FnCallEnum::Floor(call) => call.args_mut(), + FnCallEnum::Gt(call) => call.args_mut(), + FnCallEnum::GtEq(call) => call.args_mut(), + FnCallEnum::If(call) => call.args_mut(), + FnCallEnum::IntMod(call) => call.args_mut(), + FnCallEnum::Left(call) => call.args_mut(), + FnCallEnum::Lt(call) => call.args_mut(), + FnCallEnum::LtEq(call) => call.args_mut(), + FnCallEnum::IsNull(call) => call.args_mut(), + FnCallEnum::IsNotNull(call) => call.args_mut(), + FnCallEnum::Lower(call) => call.args_mut(), + FnCallEnum::Max(call) => call.args_mut(), + FnCallEnum::Min(call) => call.args_mut(), + FnCallEnum::Multiply(call) => call.args_mut(), + FnCallEnum::Neq(call) => call.args_mut(), + FnCallEnum::Not(call) => call.args_mut(), + FnCallEnum::Or(call) => call.args_mut(), + FnCallEnum::Pow(call) => call.args_mut(), + FnCallEnum::Sqrt(call) => call.args_mut(), + FnCallEnum::RegexpExtract(call) => call.args_mut(), + FnCallEnum::RegexpLike(call) => call.args_mut(), + FnCallEnum::Right(call) => call.args_mut(), + FnCallEnum::Round(call) => call.args_mut(), + FnCallEnum::SplitAfter(call) => call.args_mut(), + FnCallEnum::SplitBefore(call) => call.args_mut(), + FnCallEnum::Subtract(call) => call.args_mut(), + FnCallEnum::Substring(call) => call.args_mut(), + FnCallEnum::SubstringCount(call) => call.args_mut(), + FnCallEnum::TextJoin(call) => call.args_mut(), + FnCallEnum::Trim(call) => call.args_mut(), + FnCallEnum::Upper(call) => call.args_mut(), + } + } + + /// Produce CraneLift IR for the given function call. + pub(crate) fn lower( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + match self { + FnCallEnum::Abs(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::And(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Ceil(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Concat(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Add(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Divide(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Eq(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Floor(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Gt(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::GtEq(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::If(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::IntMod(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Left(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Lt(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::LtEq(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::IsNull(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::IsNotNull(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Lower(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Max(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Min(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Multiply(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Neq(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Not(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Or(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Pow(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Sqrt(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::RegexpExtract(call) => { + call.emit_cranelift_ir(return_type, context, builder) + } + FnCallEnum::RegexpLike(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Right(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Round(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::SplitAfter(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::SplitBefore(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Subtract(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Substring(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::SubstringCount(call) => { + call.emit_cranelift_ir(return_type, context, builder) + } + FnCallEnum::TextJoin(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Trim(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::Upper(call) => call.emit_cranelift_ir(return_type, context, builder), + } + } +} + +/// Error representing an invalid function call. +#[derive(Debug, Eq, PartialEq, thiserror::Error)] +pub enum InvalidFunctionCall { + #[error("invalid number of arguments: expected {expected}, got {provided}")] + InvalidNumberOfArguments { + expected: ArgumentCount, + provided: usize, + }, + #[error("argument {argument} must be a {expected:?} literal")] + ExpectedLiteral { argument: usize, expected: VarType }, + #[error("invalid value for argument {argument}: expected {expected}")] + InvalidLiteralValue { + argument: usize, + expected: &'static str, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ArgumentCount { + Any, + Exactly(usize), + AtLeast(usize), + Between { min: usize, max: usize }, +} + +impl std::fmt::Display for ArgumentCount { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match *self { + ArgumentCount::Any => formatter.write_str("any number of arguments"), + ArgumentCount::Exactly(1) => formatter.write_str("exactly 1 argument"), + ArgumentCount::Exactly(count) => { + write!(formatter, "exactly {count} arguments") + } + ArgumentCount::AtLeast(1) => formatter.write_str("at least 1 argument"), + ArgumentCount::AtLeast(count) => { + write!(formatter, "at least {count} arguments") + } + ArgumentCount::Between { min: 1, max: 1 } => formatter.write_str("exactly 1 argument"), + ArgumentCount::Between { min, max } if min == max => { + write!(formatter, "exactly {min} arguments") + } + ArgumentCount::Between { min, max } => { + write!(formatter, "between {min} and {max} arguments") + } + } + } +} + +impl ArgumentCount { + fn validate(self, args: &[UntypedExpr]) -> Result<(), InvalidFunctionCall> { + let provided = args.len(); + let is_valid = match self { + ArgumentCount::Any => true, + ArgumentCount::Exactly(expected) => provided == expected, + ArgumentCount::AtLeast(expected) => provided >= expected, + ArgumentCount::Between { min, max } => (min..=max).contains(&provided), + }; + if is_valid { + Ok(()) + } else { + Err(InvalidFunctionCall::InvalidNumberOfArguments { + expected: self, + provided, + }) + } + } +} + +pub(crate) fn validate_literal( + args: &[UntypedExpr], + index: usize, + expected: VarType, + is_valid: impl FnOnce(&Literal) -> bool, +) -> Result<(), InvalidFunctionCall> { + let Some(UntypedExpr::Literal(literal)) = args.get(index) else { + return Err(InvalidFunctionCall::ExpectedLiteral { + argument: index + 1, + expected, + }); + }; + if is_valid(literal) { + Ok(()) + } else { + Err(InvalidFunctionCall::ExpectedLiteral { + argument: index + 1, + expected, + }) + } +} + +/// Implements the type-inference, typed-AST, and lowering phases of a function call. +/// +/// The static methods operate on an [`UntypedExpr`] call before a concrete call node exists. +/// Once [`FnCall::call_with_types`] has produced that node, [`FnCall::args_mut`] and +/// [`FnCall::lower`] operate on its typed representation. +pub(crate) trait FnCall: std::fmt::Debug + Into { + const ARG_COUNT: ArgumentCount; + + /// Constrains the call and its arguments to the types accepted by its parent expression. + /// + /// Implementations validate their signature, recursively infer every argument, update + /// `inferred_types` with the accepted types for variables, and return the possible result + /// types that remain after intersecting with `target_type`. + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result + where + Self: Sized; + + fn validate_args(args: &[UntypedExpr]) -> Result<(), InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + Ok(()) + } + + /// Builds the typed call after concrete variable types have been supplied. + /// + /// `target_type_set` communicates the result types set accepted by the parent call. The + /// implementation selects a concrete result type, applies compatible target types to its + /// arguments through `context`, and stores any compilation resources on the typed call. + /// + /// The type of the returned is given to the caller in the TypedExpr object. + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result + where + Self: Sized; + + /// Returns the typed child expressions that participate in recursive AST passes. + /// + /// This is only used, to assign and deduplicate variable input slots. Compile-time + /// configuration stored directly on a call does not need to be returned. + /// + /// Today this is only used as a cheap visitor to allocate variable ids. + fn args_mut(&mut self) -> &mut [TypedExpr]; + + /// Serializes the function name and its normalized typed arguments. + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result; + + /// Emits Cranelift IR for an already typed call and returns its result SSA value. + /// + /// `return_type` is the concrete type selected during typed-AST construction. Implementations + /// lower child expressions through `context` and append their own instructions to `builder`. + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::compile::compile; + + fn call_error(function: Function, args: Vec) -> InvalidFunctionCall { + match UntypedExpr::call(function, args) { + Ok(_) => panic!("expected the function call to be rejected"), + Err(error) => error, + } + } + + #[test] + fn test_argument_count_display() { + assert_eq!(ArgumentCount::Any.to_string(), "any number of arguments"); + assert_eq!(ArgumentCount::Exactly(1).to_string(), "exactly 1 argument"); + assert_eq!(ArgumentCount::Exactly(2).to_string(), "exactly 2 arguments"); + assert_eq!(ArgumentCount::AtLeast(1).to_string(), "at least 1 argument"); + assert_eq!( + ArgumentCount::Between { min: 2, max: 3 }.to_string(), + "between 2 and 3 arguments" + ); + } + + #[test] + fn test_argument_count_validation() { + assert_eq!( + call_error(Function::Abs, Vec::new()), + InvalidFunctionCall::InvalidNumberOfArguments { + expected: ArgumentCount::Exactly(1), + provided: 0, + } + ); + assert_eq!( + call_error(Function::And, Vec::new()), + InvalidFunctionCall::InvalidNumberOfArguments { + expected: ArgumentCount::AtLeast(1), + provided: 0, + } + ); + assert_eq!( + call_error(Function::Round, Vec::new()), + InvalidFunctionCall::InvalidNumberOfArguments { + expected: ArgumentCount::Between { min: 1, max: 2 }, + provided: 0, + } + ); + assert!(UntypedExpr::call(Function::Add, Vec::new()).is_ok()); + } + + #[test] + fn test_literal_argument_validation() { + assert_eq!( + call_error( + Function::RegexpLike, + vec![ + UntypedExpr::variable("input"), + UntypedExpr::variable("pattern") + ], + ), + InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + ); + assert_eq!( + call_error( + Function::RegexpExtract, + vec![ + UntypedExpr::variable("input"), + UntypedExpr::literal("pattern"), + UntypedExpr::literal(1i64), + ], + ), + InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::U64, + } + ); + } + + #[test] + fn test_typed_construction_validates_unchecked_ast() { + let expression = Function::Abs.call_untyped_expr(Vec::new()); + let error = match compile(&expression, &HashMap::new()) { + Ok(_) => panic!("expected compilation to reject the unchecked AST"), + Err(error) => error, + }; + assert!(matches!( + error, + CompileError::InvalidArguments(InvalidFunctionCall::InvalidNumberOfArguments { + expected: ArgumentCount::Exactly(1), + provided: 0, + }) + )); + + let expression = Function::RegexpLike.call_untyped_expr(vec![ + UntypedExpr::variable("input"), + UntypedExpr::variable("pattern"), + ]); + let variable_types = HashMap::from([("input", VarType::Str), ("pattern", VarType::Str)]); + let error = match compile(&expression, &variable_types) { + Ok(_) => panic!("expected compilation to reject the non-literal pattern"), + Err(error) => error, + }; + assert!(matches!( + error, + CompileError::InvalidArguments(InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + }) + )); + } +} diff --git a/jitexpr/src/functions/multiply.rs b/jitexpr/src/functions/multiply.rs new file mode 100644 index 0000000000..24c791241b --- /dev/null +++ b/jitexpr/src/functions/multiply.rs @@ -0,0 +1,233 @@ +//! `MULTIPLY` multiplies two numeric expressions. +//! +//! It accepts exactly two numeric arguments. Both operands are coerced to one common numeric type +//! using the same rules as `ADD`: prefer a common `i64`, then `u64`, and fall back to `f64` when no +//! integer type represents both operands. Integer multiplication wraps at 64 bits; floating-point +//! multiplication follows IEEE-754, including NaN and infinity behavior. +//! +//! Null propagation is strict: if either operand is absent, the result is absent. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use super::add::{is_numerical, select_return_type, with_float_fallback}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct MultiplyFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for MultiplyFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::NUMERICAL).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Multiply, + expected: target_type, + got: InferredTypeSet::NUMERICAL, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Multiply, + expected: 2, + got: args.len(), + }); + } + + let mut return_types = InferredTypeSet::NUMERICAL; + for arg in args { + let arg_types = + crate::ast::infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types)?; + return_types = return_types.intersect(arg_types); + } + return_types = with_float_fallback(return_types); + let constrained_return_types = return_types.intersect(target_type); + if constrained_return_types.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Multiply, + expected: target_type, + got: return_types, + }); + } + Ok(constrained_return_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let mut return_types = InferredTypeSet::NUMERICAL.intersect(target_type_set); + for arg in args { + let arg_types = crate::ast::infer_type_with_variable_types( + arg, + InferredTypeSet::NUMERICAL, + context.variable_types(), + )?; + return_types = return_types.intersect(arg_types); + } + let return_type = select_return_type(with_float_fallback(return_types)); + let typed_args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::singleton(return_type))) + .collect::, _>>()?; + if typed_args + .iter() + .any(|typed_arg| !is_numerical(typed_arg.return_type)) + { + return Ok(TypedExpr::none()); + } + + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(MultiplyFnCall { + args: typed_args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("MULTIPLY", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let left = context.compile_expr(&self.args[0], builder)?; + let right = context.compile_expr(&self.args[1], builder)?; + let value = match return_type { + VarType::U64 | VarType::I64 => builder.ins().imul(left.value, right.value), + VarType::F64 => builder.ins().fmul(left.value, right.value), + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Multiply, + return_type, + }); + } + }; + let is_present = builder.ins().band(left.is_present, right.is_present); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: MultiplyFnCall) -> Self { + FnCallEnum::Multiply(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_infer_types_requires_two_numeric_arguments() { + let expression = deserialize("(MULTIPLY left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!( + inferred_types.get("left"), + Some(&InferredTypeSet::NUMERICAL) + ); + assert_eq!( + inferred_types.get("right"), + Some(&InferredTypeSet::NUMERICAL) + ); + + for expression in ["(MULTIPLY 1i64)", "(MULTIPLY 1i64 2i64 3i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Multiply, + expected: 2, + .. + }) + )); + } + } + + #[test] + fn test_signed_unsigned_and_float_multiplication() { + let expression = deserialize("(MULTIPLY -7i64 3i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns i64. + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, Some(-21)); + + let expression = deserialize("(MULTIPLY 9223372036854775808u64 2u64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::U64); + // SAFETY: The expression has no inputs and returns u64; multiplication wraps. + assert_eq!(unsafe { compiled.call(&[]).as_u64() }, Some(0)); + + let expression = deserialize("(MULTIPLY 1.5f64 2f64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns f64. + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, Some(3.0)); + } + + #[test] + fn test_nan_is_preserved() { + let expression = deserialize("(MULTIPLY nanf64 2f64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + + // SAFETY: The expression has no inputs and returns f64. + assert!(unsafe { compiled.call(&[]).as_f64() }.unwrap().is_nan()); + } + + #[test] + fn test_runtime_null_propagation_and_integer_coercion() { + let expression = deserialize("(MULTIPLY value 3i64)").unwrap(); + let variable_types = HashMap::from([("value", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::U64); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(4u64)]).as_u64() }, + Some(12) + ); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_u64() }, + None + ); + } + + #[test] + fn test_compile_time_none_propagates() { + let expression = deserialize("(MULTIPLY none 2i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + + assert_eq!(compiled.result_type(), VarType::None); + // SAFETY: The expression has no inputs and returns an absent value. + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, None); + } +} diff --git a/jitexpr/src/functions/native_function.rs b/jitexpr/src/functions/native_function.rs new file mode 100644 index 0000000000..30b5ef6f32 --- /dev/null +++ b/jitexpr/src/functions/native_function.rs @@ -0,0 +1,157 @@ +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type}; +use cranelift_jit::{JITBuilder, JITModule}; + +use super::{ + comparison, concat, eq, int_mod, lower, pow, regexp_extract, regexp_like, round, split_after, + split_before, substring, substring_count, trim, upper, +}; +use crate::compile::CompileError; + +/// References to native functions imported into the current Cranelift function. +pub(crate) struct NativeFunctions { + string_eq: FuncRef, + string_lowercase: FuncRef, + string_uppercase: FuncRef, + string_trim: FuncRef, + substring_count: FuncRef, + substring: FuncRef, + split_after: FuncRef, + split_before: FuncRef, + string_concat: FuncRef, + float_mod: FuncRef, + float_pow: FuncRef, + round_float: FuncRef, + round_int_to_i64: FuncRef, + round_float_to_i64: FuncRef, + regexp_extract: FuncRef, + regexp_like: FuncRef, + string_compare: FuncRef, + f64_i64_compare: FuncRef, + f64_u64_compare: FuncRef, +} + +impl NativeFunctions { + pub(crate) fn string_eq(&self) -> FuncRef { + self.string_eq + } + + pub(crate) fn string_lowercase(&self) -> FuncRef { + self.string_lowercase + } + + pub(crate) fn string_uppercase(&self) -> FuncRef { + self.string_uppercase + } + pub(crate) fn string_trim(&self) -> FuncRef { + self.string_trim + } + pub(crate) fn substring_count(&self) -> FuncRef { + self.substring_count + } + + pub(crate) fn substring(&self) -> FuncRef { + self.substring + } + + pub(crate) fn split_after(&self) -> FuncRef { + self.split_after + } + + pub(crate) fn split_before(&self) -> FuncRef { + self.split_before + } + + pub(crate) fn string_concat(&self) -> FuncRef { + self.string_concat + } + + pub(crate) fn float_mod(&self) -> FuncRef { + self.float_mod + } + + pub(crate) fn float_pow(&self) -> FuncRef { + self.float_pow + } + + pub(crate) fn round_float(&self) -> FuncRef { + self.round_float + } + + pub(crate) fn round_int_to_i64(&self) -> FuncRef { + self.round_int_to_i64 + } + + pub(crate) fn round_float_to_i64(&self) -> FuncRef { + self.round_float_to_i64 + } + + pub(crate) fn regexp_extract(&self) -> FuncRef { + self.regexp_extract + } + + pub(crate) fn regexp_like(&self) -> FuncRef { + self.regexp_like + } + + pub(crate) fn string_compare(&self) -> FuncRef { + self.string_compare + } + + pub(crate) fn f64_i64_compare(&self) -> FuncRef { + self.f64_i64_compare + } + + pub(crate) fn f64_u64_compare(&self) -> FuncRef { + self.f64_u64_compare + } +} + +/// 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); + comparison::register_jit_symbols(jit_builder); + lower::register_jit_symbol(jit_builder); + upper::register_jit_symbol(jit_builder); + trim::register_jit_symbol(jit_builder); + substring_count::register_jit_symbol(jit_builder); + substring::register_jit_symbol(jit_builder); + split_after::register_jit_symbol(jit_builder); + split_before::register_jit_symbol(jit_builder); + concat::register_jit_symbol(jit_builder); + int_mod::register_jit_symbol(jit_builder); + pow::register_jit_symbol(jit_builder); + round::register_jit_symbols(jit_builder); + regexp_extract::register_jit_symbol(jit_builder); + regexp_like::register_jit_symbol(jit_builder); +} + +/// Declares every native function imported by the expression being compiled. +pub(crate) fn declare_native_functions( + module: &mut JITModule, + function: &mut CraneliftFunction, + pointer_type: Type, +) -> Result { + let comparison = comparison::declare_native_functions(module, function, pointer_type)?; + let round = round::declare_native_functions(module, function)?; + Ok(NativeFunctions { + string_eq: eq::declare_native_function(module, function, pointer_type)?, + string_lowercase: lower::declare_native_function(module, function, pointer_type)?, + string_uppercase: upper::declare_native_function(module, function, pointer_type)?, + string_trim: trim::declare_native_function(module, function, pointer_type)?, + substring_count: substring_count::declare_native_function(module, function, pointer_type)?, + substring: substring::declare_native_function(module, function, pointer_type)?, + split_after: split_after::declare_native_function(module, function, pointer_type)?, + split_before: split_before::declare_native_function(module, function, pointer_type)?, + string_concat: concat::declare_native_function(module, function, pointer_type)?, + float_mod: int_mod::declare_native_function(module, function)?, + float_pow: pow::declare_native_function(module, function)?, + round_float: round.round_float, + round_int_to_i64: round.round_int_to_i64, + round_float_to_i64: round.round_float_to_i64, + regexp_extract: regexp_extract::declare_native_function(module, function, pointer_type)?, + regexp_like: regexp_like::declare_native_function(module, function, pointer_type)?, + string_compare: comparison.string_compare, + f64_i64_compare: comparison.f64_i64_compare, + f64_u64_compare: comparison.f64_u64_compare, + }) +} diff --git a/jitexpr/src/functions/neq.rs b/jitexpr/src/functions/neq.rs new file mode 100644 index 0000000000..dc9e44ec0c --- /dev/null +++ b/jitexpr/src/functions/neq.rs @@ -0,0 +1,172 @@ +//! `NEQ` tests two values for inequality by composing `NOT(EQ(...))` semantics. +//! +//! It accepts exactly two operands of any supported type. Values of the same type compare normally, +//! numeric types compare across `i64`, `u64`, and `f64`, and unrelated present types are unequal. +//! NaN is unequal to every value, including itself. +//! +//! This is deliberately not a strict-null primitive. `EQ` first produces null when either operand +//! is null; `NOT` then converts that null result to present `true`. Consequently `NEQ(null, +//! value)`, `NEQ(value, null)`, and `NEQ(null, null)` are all present and true. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{InstBuilder, IntCC, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{EqFnCall, FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct NeqFnCall { + eq: EqFnCall, +} + +impl FnCall for NeqFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + 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::Neq, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Neq, + expected: 2, + got: args.len(), + }); + } + 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 { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + 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(NeqFnCall { + eq: EqFnCall { + args: typed_args.into_boxed_slice(), + }, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + self.eq.args_mut() + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("NEQ", self.eq.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + let equal = self.eq.emit_cranelift_ir(VarType::Bool, context, builder)?; + let value = builder.ins().icmp_imm_u(IntCC::Equal, equal.value, 0); + Ok(LoweredValue { + value, + is_present: builder.ins().iconst(types::I8, 1), + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: NeqFnCall) -> Self { + FnCallEnum::Neq(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_requires_two_arguments_of_any_type() { + let expression = deserialize("(NEQ left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("left"), Some(&InferredTypeSet::ALL)); + assert_eq!(inferred_types.get("right"), Some(&InferredTypeSet::ALL)); + + let expression = deserialize("(NEQ 1i64)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Neq, + expected: 2, + .. + }) + )); + } + + #[test] + fn test_present_values_and_numeric_cross_types() { + assert_eq!(eval("(NEQ 1i64 1u64)"), Some(false)); + assert_eq!(eval("(NEQ 1i64 2f64)"), Some(true)); + assert_eq!(eval(r#"(NEQ "same" "same")"#), Some(false)); + assert_eq!(eval(r#"(NEQ "1" 1i64)"#), Some(true)); + assert_eq!(eval("(NEQ nanf64 nanf64)"), Some(true)); + } + + #[test] + fn test_null_is_converted_to_present_true() { + assert_eq!(eval("(NEQ none 1i64)"), Some(true)); + assert_eq!(eval("(NEQ 1i64 none)"), Some(true)); + assert_eq!(eval("(NEQ none none)"), Some(true)); + + let expression = deserialize("(NEQ left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::none(), VariableValue::some("value")]) + .as_bool() + }, + Some(true) + ); + } +} diff --git a/jitexpr/src/functions/not.rs b/jitexpr/src/functions/not.rs new file mode 100644 index 0000000000..9cd44ff4b7 --- /dev/null +++ b/jitexpr/src/functions/not.rs @@ -0,0 +1,170 @@ +//! `NOT` negates a nullable boolean using Datadog's non-SQL null semantics. +//! +//! It accepts exactly one boolean expression and always returns a present boolean. Present values +//! are inverted normally; an absent input returns `true`. In particular, `NOT(NULL) = TRUE`, not +//! NULL as it would under SQL three-valued logic. +//! +//! This preserves the behavior where a negated predicate also matches documents for which the field +//! is absent. The implementation must therefore select `true` from the child's presence bit instead +//! of reading or negating the unspecified payload of an absent runtime value. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct NotFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for NotFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + 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::Not, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Not, + expected: 1, + got: args.len(), + }); + } + + crate::ast::infer_types_aux(&args[0], InferredTypeSet::BOOLEAN, inferred_types)?; + Ok(InferredTypeSet::BOOLEAN) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + + let arg = context.apply_types(&args[0], InferredTypeSet::BOOLEAN)?; + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(NotFnCall { + args: vec![arg].into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("NOT", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + let arg = context.compile_expr(&self.args[0], builder)?; + let true_value = builder.ins().iconst(types::I8, 1); + let value = if self.args[0].return_type == VarType::None { + true_value + } else { + let negated = builder.ins().bxor_imm_u(arg.value, 1); + builder.ins().select(arg.is_present, negated, true_value) + }; + Ok(LoweredValue { + value, + is_present: true_value, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: NotFnCall) -> Self { + FnCallEnum::Not(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types, infer_types_with_target}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no runtime inputs and return booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_infer_types_requires_one_boolean_argument() { + let expression = deserialize("(NOT value)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("value"), Some(&InferredTypeSet::BOOLEAN)); + + assert!(matches!( + infer_types_with_target(&expression, InferredTypeSet::STRING), + Err(TypeError::WrongFunctionReturnType { + function: Function::Not, + .. + }) + )); + + for expression in ["(NOT)", "(NOT true false)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Not, + expected: 1, + .. + }) + )); + } + } + + #[test] + fn test_negates_present_booleans() { + assert_eq!(eval("(NOT true)"), Some(false)); + assert_eq!(eval("(NOT false)"), Some(true)); + assert_eq!(eval("(NOT (EQ 1i64 1i64))"), Some(false)); + } + + #[test] + fn test_absent_input_returns_present_true() { + assert_eq!(eval("(NOT none)"), Some(true)); + + let expression = deserialize("(NOT value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::Bool)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_bool() }, + Some(true) + ); + } +} diff --git a/jitexpr/src/functions/or.rs b/jitexpr/src/functions/or.rs new file mode 100644 index 0000000000..293edeba44 --- /dev/null +++ b/jitexpr/src/functions/or.rs @@ -0,0 +1,189 @@ +//! `OR` combines one or more nullable boolean expressions. +//! +//! All arguments must be boolean. Present values are combined with ordinary disjunction. The null +//! rule is intentionally unlike SQL: the result is absent only when every operand is absent. +//! Therefore `TRUE OR NULL = TRUE`, `FALSE OR NULL = FALSE`, and `NULL OR NULL = NULL`. +//! +//! One or more operands are required. The implementation combines their presence bits and ignores +//! the unspecified payload of absent operands. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct OrFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for OrFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::AtLeast(1); + + 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::Or, + expected: target_type, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.is_empty() { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Or, + expected: 1, + got: 0, + }); + } + + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::BOOLEAN, inferred_types)?; + } + Ok(InferredTypeSet::BOOLEAN) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::Bool)); + + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::BOOLEAN)) + .collect::, _>>()?; + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(OrFnCall { + args: args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("OR", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Bool); + let mut value = builder.ins().iconst(types::I8, 0); + let mut is_present = value; + + for arg in &self.args { + let lowered = context.compile_expr(arg, builder)?; + is_present = builder.ins().bor(is_present, lowered.is_present); + if arg.return_type != VarType::None { + let present_value = builder.ins().band(lowered.value, lowered.is_present); + value = builder.ins().bor(value, present_value); + } + } + + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: OrFnCall) -> Self { + FnCallEnum::Or(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no runtime inputs and return nullable booleans. + unsafe { compiled.call(&[]).as_bool() } + } + + #[test] + fn test_infer_types_requires_boolean_arguments() { + let expression = deserialize("(OR left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("left"), Some(&InferredTypeSet::BOOLEAN)); + assert_eq!(inferred_types.get("right"), Some(&InferredTypeSet::BOOLEAN)); + + let expression = deserialize("(OR)").unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Or, + expected: 1, + got: 0, + }) + )); + } + + #[test] + fn test_present_truth_table_and_variadic_inputs() { + assert_eq!(eval("(OR false)"), Some(false)); + assert_eq!(eval("(OR false false false)"), Some(false)); + assert_eq!(eval("(OR false true false)"), Some(true)); + assert_eq!(eval("(OR true true)"), Some(true)); + } + + #[test] + fn test_result_is_absent_only_when_every_argument_is_absent() { + assert_eq!(eval("(OR true none)"), Some(true)); + assert_eq!(eval("(OR false none)"), Some(false)); + assert_eq!(eval("(OR none none)"), None); + assert_eq!(eval("(OR none false none)"), Some(false)); + } + + #[test] + fn test_runtime_null_handling() { + let expression = deserialize("(OR left right)").unwrap(); + let variable_types = HashMap::from([("left", VarType::Bool), ("right", VarType::Bool)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(false), VariableValue::none()]) + .as_bool() + }, + Some(false) + ); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::none(), VariableValue::none()]) + .as_bool() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/pow.rs b/jitexpr/src/functions/pow.rs new file mode 100644 index 0000000000..6c6d079cec --- /dev/null +++ b/jitexpr/src/functions/pow.rs @@ -0,0 +1,219 @@ +//! `POW` raises a numeric base to a numeric exponent. +//! +//! It accepts exactly two numeric arguments. Both are converted to `f64`, and the result is always +//! `f64`. A negative base with a non-integral exponent represents a complex result and is returned +//! as null. Null operands propagate. +//! +//! Other IEEE-754 results remain present: overflow may produce infinity, and a NaN base may +//! produce NaN. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, FloatCC, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_float_pow"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct PowFnCall { + args: Box<[TypedExpr]>, +} + +impl FnCall for PowFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::F64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Pow, + expected: target_type, + got: InferredTypeSet::F64, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Pow, + expected: 2, + got: args.len(), + }); + } + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types)?; + } + Ok(InferredTypeSet::F64) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::F64)); + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::F64)) + .collect::, _>>()?; + if args.iter().any(|arg| arg.return_type == VarType::None) { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type: VarType::F64, + ast: TypedExprAst::from_call(PowFnCall { + args: args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("POW", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::F64); + let base = context.compile_expr(&self.args[0], builder)?; + let exponent = context.compile_expr(&self.args[1], builder)?; + let call = builder.ins().call( + context.native_functions().float_pow(), + &[base.value, exponent.value], + ); + let result = builder.inst_results(call)[0]; + + let zero = builder.ins().f64const(0.0); + let base_is_negative = builder.ins().fcmp(FloatCC::LessThan, base.value, zero); + let result_is_nan = builder.ins().fcmp(FloatCC::Unordered, result, result); + let complex_result = builder.ins().band(base_is_negative, result_is_nan); + let is_real = builder.ins().bxor_imm_u(complex_result, 1); + let both_present = builder.ins().band(base.is_present, exponent.is_present); + let is_present = builder.ins().band(both_present, is_real); + + Ok(LoweredValue { + value: result, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, float_pow as *const u8); +} + +pub(super) fn declare_native_function( + module: &mut JITModule, + function: &mut CraneliftFunction, +) -> Result { + let mut signature = module.make_signature(); + signature + .params + .extend(std::iter::repeat_n(AbiParam::new(types::F64), 2)); + signature.returns.push(AbiParam::new(types::F64)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +extern "C" fn float_pow(base: f64, exponent: f64) -> f64 { + base.powf(exponent) +} + +impl From for FnCallEnum { + fn from(call: PowFnCall) -> Self { + FnCallEnum::Pow(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable f64 values. + unsafe { compiled.call(&[]).as_f64() } + } + + #[test] + fn test_signature_and_output_type() { + let expression = deserialize("(POW base exponent)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + assert_eq!(inferred.get("base"), Some(&InferredTypeSet::NUMERICAL)); + assert_eq!(inferred.get("exponent"), Some(&InferredTypeSet::NUMERICAL)); + + for expression in ["(POW 2i64)", "(POW 2i64 3i64 4i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Pow, + expected: 2, + .. + }) + )); + } + let expression = deserialize("(POW 2i64 3i64)").unwrap(); + assert_eq!( + compile(&expression, &HashMap::new()).unwrap().result_type(), + VarType::F64 + ); + } + + #[test] + fn test_numeric_and_ieee_edges() { + assert_eq!(eval("(POW 2i64 3i64)"), Some(8.0)); + assert_eq!(eval("(POW -2f64 3f64)"), Some(-8.0)); + assert_eq!(eval("(POW -2f64 0.5f64)"), None); + assert_eq!(eval("(POW none 2f64)"), None); + assert!(eval("(POW nanf64 2f64)").unwrap().is_nan()); + assert!(eval("(POW 1e308f64 2f64)").unwrap().is_infinite()); + } + + #[test] + fn test_runtime_null_propagates() { + let expression = deserialize("(POW base exponent)").unwrap(); + let variable_types = HashMap::from([("base", VarType::I64), ("exponent", VarType::F64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(3i64), VariableValue::some(2.0f64)]) + .as_f64() + }, + Some(9.0) + ); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(3i64), VariableValue::none()]) + .as_f64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/regexp_extract.rs b/jitexpr/src/functions/regexp_extract.rs new file mode 100644 index 0000000000..fe091e4dac --- /dev/null +++ b/jitexpr/src/functions/regexp_extract.rs @@ -0,0 +1,433 @@ +// RegexpExtract extracts a regular-expression match from a string. +// +// It takes two or three arguments: +// - string: the input string +// - const string: a regular-expression pattern literal. This one CANNOT be the result of another +// expression +// - optional const u64: capture index literal. It defaults to 0 during conversion to the typed +// expression. Capture index 0 returns the full match, while indexes 1 and above return the +// corresponding explicit capture group. +// +// It returns None when the input is None, the pattern does not match, or the requested capture +// group is absent or did not participate in the match. + +use std::collections::HashMap; +use std::sync::Arc; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; +use regex::Regex; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum, InvalidFunctionCall}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_regexp_extract"; + +#[derive(Clone, Debug)] +pub(crate) struct RegexpExtractFnCall { + regex: Arc, + haystack: Box, + capture_index: u64, +} + +impl PartialEq for RegexpExtractFnCall { + fn eq(&self, other: &Self) -> bool { + self.regex.as_str() == other.regex.as_str() + && self.haystack == other.haystack + && self.capture_index == other.capture_index + } +} + +impl FnCall for RegexpExtractFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Between { min: 2, max: 3 }; + + fn validate_args(args: &[UntypedExpr]) -> Result<(), InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::Str, |literal| { + matches!(literal, Literal::String(_)) + })?; + if args.len() == 3 { + super::validate_literal(args, 2, VarType::U64, |literal| { + matches!(literal, Literal::U64(_)) + })?; + } + Ok(()) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::RegexpExtract, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if !(2..=3).contains(&args.len()) { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::RegexpExtract, + expected: 3, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::STRING, inferred_types)?; + if let Some(capture_index) = args.get(2) { + crate::ast::infer_types_aux(capture_index, InferredTypeSet::NUMERICAL, inferred_types)?; + } + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let haystack = context.apply_types(&args[0], target_type_set)?; + if haystack.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + assert_eq!(haystack.return_type, VarType::Str); + + let UntypedExpr::Literal(Literal::String(pattern)) = &args[1] else { + return Err(InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + .into()); + }; + let regex = Arc::new( + Regex::new(pattern).map_err(|source| CompileError::InvalidRegex { + pattern: pattern.to_string(), + source, + })?, + ); + + let capture_index = match args.get(2) { + None => 0, + Some(UntypedExpr::Literal(Literal::U64(capture_index))) => *capture_index, + Some(_) => { + return Err(InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::U64, + } + .into()); + } + }; + + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(RegexpExtractFnCall { + regex, + haystack: Box::new(haystack), + capture_index, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.haystack) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "REGEXP_EXTRACT {} ", self.haystack)?; + crate::compile::format_string_literal(self.regex.as_str(), formatter)?; + write!(formatter, " {}u64", self.capture_index) + } + + /// Produce CraneLift IR for the given function call. + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + + let haystack = context.compile_expr(&self.haystack, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let haystack_ptr = builder + .ins() + .select(haystack.is_present, haystack.value, null); + let regex_ptr = builder + .ins() + .iconst(context.pointer_type(), Arc::as_ptr(&self.regex) as i64); + let capture_index = builder.ins().iconst(types::I64, self.capture_index as i64); + let call = builder.ins().call( + context.native_functions().regexp_extract(), + &[regex_ptr, haystack_ptr, haystack.string_len, capture_index], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder + .ins() + .icmp_imm_u(cranelift::prelude::IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, regexp_extract 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(std::iter::repeat_n(AbiParam::new(pointer_type), 2)); + signature.params.push(AbiParam::new(types::I64)); + signature.params.push(AbiParam::new(types::I64)); + signature.returns.push(AbiParam::new(pointer_type)); + signature.returns.push(AbiParam::new(types::I64)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +/// Raw two-word string result returned to generated code. +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } + + fn some(value: &str) -> Self { + Self { + ptr: value.as_ptr(), + len: value.len(), + } + } +} + +/// Runtime implementation called by generated code for `RegexpExtract`. +/// +/// The JIT forwards a nullable UTF-8 pointer and byte length. The returned +/// pointer and length borrow directly from the haystack. +unsafe extern "C" fn regexp_extract( + regex: *const Regex, + haystack_ptr: *const u8, + haystack_len: usize, + capture_index: u64, +) -> RawStr { + if haystack_ptr.is_null() { + return RawStr::none(); + } + let Ok(capture_index) = usize::try_from(capture_index) else { + return RawStr::none(); + }; + // SAFETY: Generated code embeds a pointer to the Arc-owned Regex stored in + // the typed expression retained by CompiledFn. + let regex = unsafe { &*regex }; + // SAFETY: The contract of CompiledFn::call requires a live UTF-8 string + // pointer and its exact byte length for every present string input. + let haystack = unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(haystack_ptr, haystack_len)) + }; + let Some(regex_match) = regex + .captures(haystack) + .and_then(|captures| captures.get(capture_index)) + else { + return RawStr::none(); + }; + + RawStr::some(regex_match.as_str()) +} + +impl From for FnCallEnum { + fn from(call: RegexpExtractFnCall) -> Self { + FnCallEnum::RegexpExtract(call) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::ast::{self, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_infer_types_constrains_haystack_to_string() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)")"#).unwrap(); + + let inferred_types = infer_types(&expression).unwrap(); + + assert_eq!( + inferred_types.get("message"), + Some(&InferredTypeSet::STRING) + ); + } + + #[test] + fn test_infer_types_accepts_optional_capture_index() { + for expression in [ + r#"(REGEXP_EXTRACT message "([a-z]+)")"#, + r#"(REGEXP_EXTRACT message "([a-z]+)" 1u64)"#, + ] { + let expression = ast::deserialize(expression).unwrap(); + assert!(infer_types(&expression).is_ok()); + } + + for expression in [ + r#"(REGEXP_EXTRACT message)"#, + r#"(REGEXP_EXTRACT message "([a-z]+)" 0u64 1u64)"#, + ] { + let expression = ast::deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::RegexpExtract, + expected: 3, + .. + }) + )); + } + } + + #[test] + fn test_compile_returns_borrowed_capture() { + let expression = + ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)-(\\d+)" 1u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let haystack = "prefix user-123 suffix"; + let input = [VariableValue::some(haystack)]; + let output = unsafe { compiled.call(&input) }; + + let extracted = unsafe { output.as_str() }.unwrap(); + assert_eq!(extracted, "user"); + assert_eq!(extracted.as_ptr(), haystack[7..].as_ptr()); + } + + #[test] + fn test_compile_selects_capture_by_index() { + let expression = + ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)-(\\d+)" 2u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("user-123")]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, Some("123")); + } + + #[test] + fn test_compile_returns_none_without_capture() { + let expression = + ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)-(\\d+)" 0u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("no digits here")]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, None); + } + + #[test] + fn test_compile_propagates_none_haystack() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)" 0u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::none()]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, None); + } + + #[test] + fn test_compile_propagates_compile_time_none_haystack() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT missing "([a-z]+)" 0u64)"#).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::None); + let output = unsafe { compiled.call(&[]) }; + + assert_eq!(unsafe { output.as_str() }, None); + } + + #[test] + fn test_compile_distinguishes_empty_capture_from_none() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT "b" "(a*)b" 1u64)"#).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + let output = unsafe { compiled.call(&[]) }; + + assert_eq!(unsafe { output.as_str() }, Some("")); + } + + #[test] + fn test_compile_omitted_group_defaults_to_full_match() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT message "[a-z]+-\\d+")"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("prefix user-123 suffix")]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, Some("user-123")); + } + + #[test] + fn test_compile_group_zero_returns_full_match_without_capture_groups() { + let expression = + ast::deserialize(r#"(REGEXP_EXTRACT message "[a-z]+-\\d+" 0u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("prefix user-123 suffix")]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, Some("user-123")); + } + + #[test] + fn test_compile_nested_calls_use_their_own_regexes() { + let expression = ast::deserialize( + r#"(REGEXP_EXTRACT + (REGEXP_EXTRACT message "([a-z]+-\\d+)" 1u64) + "([a-z]+)" + 1u64)"#, + ) + .unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input = [VariableValue::some("id=user-123!")]; + let output = unsafe { compiled.call(&input) }; + + assert_eq!(unsafe { output.as_str() }, Some("user")); + } + + #[test] + fn test_compile_rejects_invalid_pattern() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT "anything" "(" 0u64)"#).unwrap(); + let error = compile(&expression, &HashMap::new()).err().unwrap(); + assert!(matches!( + error, + CompileError::InvalidRegex { pattern, .. } if pattern == "(" + )); + } +} diff --git a/jitexpr/src/functions/regexp_like.rs b/jitexpr/src/functions/regexp_like.rs new file mode 100644 index 0000000000..32503acfc6 --- /dev/null +++ b/jitexpr/src/functions/regexp_like.rs @@ -0,0 +1,217 @@ +//! `REGEXP_LIKE(input, pattern)` tests whether a regular expression matches anywhere in a string. +//! +//! The pattern must be constant. Before compilation, the placeholder ICU converter removes +//! Java-style named-group names (`(?` becomes `(`) and changes atomic groups (`(?>` becomes +//! `(?:`). Invalid patterns fail compilation. Null input and definitively non-string input both +//! return present `false`; this deliberate mismatch behavior is different from ordinary strict +//! functions. + +use std::collections::HashMap; +use std::sync::Arc; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; +use regex::Regex; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_regexp_like"; +#[derive(Clone, Debug)] +pub(crate) struct RegexpLikeFnCall { + regex: Arc, + input: Box, +} +impl PartialEq for RegexpLikeFnCall { + fn eq(&self, other: &Self) -> bool { + self.regex.as_str() == other.regex.as_str() && self.input == other.input + } +} + +fn convert_pattern(pattern: &str) -> String { + let named = Regex::new(r"\(\?<[^>]+>").expect("static regex"); + named.replace_all(pattern, "(").replace("(?>", "(?:") +} + +impl FnCall for RegexpLikeFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::Str, |literal| { + matches!(literal, Literal::String(_)) + }) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target: InferredTypeSet, + inferred: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target.intersect(InferredTypeSet::BOOLEAN).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::RegexpLike, + expected: target, + got: InferredTypeSet::BOOLEAN, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::RegexpLike, + expected: 2, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::ALL, inferred)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::STRING, inferred)?; + Ok(InferredTypeSet::BOOLEAN) + } + fn call_with_types( + args: &[UntypedExpr], + _target: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input_target = match &args[0] { + UntypedExpr::Literal(literal) => InferredTypeSet::singleton(literal.r#type()), + _ => InferredTypeSet::ALL, + }; + let input = context.apply_types(&args[0], input_target)?; + let UntypedExpr::Literal(Literal::String(pattern)) = &args[1] else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + .into()); + }; + let converted = convert_pattern(pattern); + let regex = + Arc::new( + Regex::new(&converted).map_err(|source| CompileError::InvalidRegex { + pattern: pattern.to_string(), + source, + })?, + ); + Ok(TypedExpr { + return_type: VarType::Bool, + ast: TypedExprAst::from_call(RegexpLikeFnCall { + regex, + input: Box::new(input), + }), + }) + } + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "REGEXP_LIKE {} ", self.input)?; + crate::compile::format_string_literal(self.regex.as_str(), formatter) + } + fn emit_cranelift_ir( + &self, + _return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + if self.input.return_type != VarType::Str { + return Ok(LoweredValue { + value: builder.ins().iconst(types::I8, 0), + is_present: builder.ins().iconst(types::I8, 1), + string_len: builder.ins().iconst(types::I64, 0), + }); + } + let input = context.compile_expr(&self.input, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, null); + let regex_ptr = builder + .ins() + .iconst(context.pointer_type(), Arc::as_ptr(&self.regex) as i64); + let call = builder.ins().call( + context.native_functions().regexp_like(), + &[regex_ptr, input_ptr, input.string_len], + ); + Ok(LoweredValue { + value: builder.inst_results(call)[0], + is_present: builder.ins().iconst(types::I8, 1), + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} +pub(super) fn register_jit_symbol(builder: &mut JITBuilder) { + builder.symbol(SYMBOL, regexp_like as *const u8); +} +pub(super) fn declare_native_function( + module: &mut JITModule, + function: &mut CraneliftFunction, + pointer: Type, +) -> Result { + let mut signature = module.make_signature(); + signature.params.extend([ + AbiParam::new(pointer), + AbiParam::new(pointer), + AbiParam::new(types::I64), + ]); + signature.returns.push(AbiParam::new(types::I8)); + let id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(id, function)) +} +unsafe extern "C" fn regexp_like(regex: *const Regex, input: *const u8, len: usize) -> u8 { + if input.is_null() { + return 0; + } + let regex = unsafe { &*regex }; + let input = unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input, len)) }; + u8::from(regex.is_match(input)) +} +impl From for FnCallEnum { + fn from(call: RegexpLikeFnCall) -> Self { + FnCallEnum::RegexpLike(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + fn eval(expr: &str) -> Option { + let expr = deserialize(expr).unwrap(); + let mut compiled = compile(&expr, &HashMap::new()).unwrap().context(); + unsafe { compiled.call(&[]).as_bool() } + } + #[test] + fn test_signature_matching_and_conversion() { + assert!(matches!( + infer_types(&deserialize("(REGEXP_LIKE \"a\")").unwrap()), + Err(TypeError::InvalidNumberOfArguments { + function: Function::RegexpLike, + expected: 2, + .. + }) + )); + assert_eq!(eval("(REGEXP_LIKE \"prefix-123\" \"[0-9]+\")"), Some(true)); + assert_eq!(eval("(REGEXP_LIKE \"abc\" \"^z\")"), Some(false)); + assert_eq!(eval("(REGEXP_LIKE \"abc\" \"(?abc)\")"), Some(true)); + } + #[test] + fn test_null_and_non_string_are_false() { + assert_eq!(eval("(REGEXP_LIKE none \"x\")"), Some(false)); + assert_eq!(eval("(REGEXP_LIKE 123i64 \"123\")"), Some(false)); + let expr = deserialize("(REGEXP_LIKE value \"x\")").unwrap(); + let mut compiled = compile(&expr, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_bool() }, + Some(false) + ); + } +} diff --git a/jitexpr/src/functions/right.rs b/jitexpr/src/functions/right.rs new file mode 100644 index 0000000000..ab25e460f8 --- /dev/null +++ b/jitexpr/src/functions/right.rs @@ -0,0 +1,215 @@ +//! `RIGHT(input, length)` returns the last `length` UTF-8 bytes of a string. +//! +//! The length must be an integer constant and is stored as a `usize`. A length greater than the +//! input byte length returns the whole string, zero returns a present empty string, and null input +//! propagates. Negative lengths and lengths that place the start inside a UTF-8 code point return +//! null. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{InstBuilder, IntCC}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct RightFnCall { + input: Box, + length: usize, +} + +fn constant_length(expression: &UntypedExpr) -> Result, super::InvalidFunctionCall> { + let UntypedExpr::Literal(literal) = expression else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::I64, + }); + }; + if !literal.is_none() && !literal.types().contains(VarType::I64) { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::I64, + }); + } + Ok(match literal { + Literal::I64(value) => usize::try_from(*value).ok(), + Literal::U64(value) => usize::try_from(*value).ok(), + Literal::F64(value) => usize::try_from(*value as i64).ok(), + Literal::None => None, + Literal::Bool(_) | Literal::String(_) => None, + }) +} + +impl FnCall for RightFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::I64, |literal| { + literal.is_none() || literal.types().contains(VarType::I64) + }) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Right, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Right, + expected: 2, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::I64, inferred_types)?; + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input = context.apply_types(&args[0], InferredTypeSet::STRING)?; + let Some(length) = constant_length(&args[1])? else { + return Ok(TypedExpr::none()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(RightFnCall { + input: Box::new(input), + length, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "RIGHT {} {}u64", self.input, self.length) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let input = context.compile_expr(&self.input, builder)?; + let zero = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, zero); + let length = builder + .ins() + .iconst(context.pointer_type(), self.length as i64); + let length_is_shorter = + builder + .ins() + .icmp(IntCC::UnsignedLessThan, length, input.string_len); + let suffix_start = builder.ins().isub(input.string_len, length); + let start = builder.ins().select(length_is_shorter, suffix_start, zero); + let call = builder.ins().call( + context.native_functions().substring(), + &[input_ptr, input.string_len, start, length], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let native_succeeded = builder.ins().icmp_imm_u(IntCC::NotEqual, value, 0); + let is_present = builder.ins().band(input.is_present, native_succeeded); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +impl From for FnCallEnum { + fn from(call: RightFnCall) -> Self { + FnCallEnum::Right(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable string. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_signature_and_byte_lengths() { + for expression in ["(RIGHT \"abc\")", "(RIGHT \"abc\" 1i64 2i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Right, + expected: 2, + .. + }) + )); + } + + assert_eq!(eval("(RIGHT \"abcdef\" 3i64)"), Some("def".into())); + assert_eq!(eval("(RIGHT \"éclair\" 5i64)"), Some("clair".into())); + assert_eq!(eval("(RIGHT \"éclair\" 6i64)"), None); + } + + #[test] + fn test_zero_clamping_and_invalid_length() { + assert_eq!(eval("(RIGHT \"abc\" 0i64)"), Some(String::new())); + assert_eq!(eval("(RIGHT \"abc\" 99i64)"), Some("abc".into())); + assert_eq!(eval("(RIGHT \"abc\" -1i64)"), None); + } + + #[test] + fn test_runtime_null() { + let expression = deserialize("(RIGHT value 2i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("abc")]).as_str() }, + Some("bc") + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + } + + #[test] + fn test_utf8_boundary_break_returns_null() { + let expression = deserialize(r#"(RIGHT "下北沢" 1i64)"#).unwrap(); + let mut compiled = compile(&expression, &HashMap::default()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_str() }, None); + let expression = deserialize(r#"(RIGHT "下北沢" 3i64)"#).unwrap(); + let mut compiled = compile(&expression, &HashMap::default()).unwrap().context(); + assert_eq!(unsafe { compiled.call(&[]).as_str() }, Some("沢")); + } +} diff --git a/jitexpr/src/functions/round.rs b/jitexpr/src/functions/round.rs new file mode 100644 index 0000000000..ecf1b29288 --- /dev/null +++ b/jitexpr/src/functions/round.rs @@ -0,0 +1,503 @@ +//! `ROUND(value, precision)` rounds one numeric argument to a decimal precision. +//! +//! Precision is an optional integer constant and defaults to zero. Halfway values are rounded away +//! from zero. A positive precision keeps digits to the right of the decimal point and returns +//! `f64`; a zero or negative precision rounds to units, tens, hundreds, and so on and returns +//! `i64`. Integer inputs use exact integer arithmetic whenever the result type is `i64`. +//! +//! Null input or null precision returns null. An `i64` result is also null when the rounded value +//! is NaN, infinite, or outside the `i64` range. An `f64` result preserves NaN and infinity. Very +//! large positive precisions leave finite values unchanged, while very large negative precisions +//! round finite values to zero. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder, IntCC}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const ROUND_FLOAT_SYMBOL: &str = "jitexpr_round_float"; +const ROUND_INT_TO_I64_SYMBOL: &str = "jitexpr_round_int_to_i64"; +const ROUND_FLOAT_TO_I64_SYMBOL: &str = "jitexpr_round_float_to_i64"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct RoundFnCall { + arg: Box, + precision: i64, +} + +fn constant_precision( + expression: Option<&UntypedExpr>, +) -> Result, super::InvalidFunctionCall> { + let Some(expression) = expression else { + return Ok(Some(0)); + }; + let UntypedExpr::Literal(literal) = expression else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::I64, + }); + }; + if !literal.is_none() && !literal.types().contains(VarType::I64) { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::I64, + }); + } + Ok(match literal { + Literal::I64(value) => Some(*value), + Literal::U64(value) => i64::try_from(*value).ok(), + Literal::F64(value) + if value.is_finite() + && value.fract() == 0.0 + && *value >= i64::MIN as f64 + && *value < -(i64::MIN as f64) => + { + Some(*value as i64) + } + Literal::None => None, + Literal::F64(_) | Literal::Bool(_) | Literal::String(_) => None, + }) +} + +fn return_type_for_precision(precision: i64) -> VarType { + if precision > 0 { + VarType::F64 + } else { + VarType::I64 + } +} + +impl FnCall for RoundFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Between { min: 1, max: 2 }; + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + if args.len() == 2 { + super::validate_literal(args, 1, VarType::I64, |literal| { + literal.is_none() || literal.types().contains(VarType::I64) + })?; + } + Ok(()) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if !(1..=2).contains(&args.len()) { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Round, + expected: 2, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::NUMERICAL, inferred_types)?; + if let Some(precision) = args.get(1) { + crate::ast::infer_types_aux(precision, InferredTypeSet::I64, inferred_types)?; + } + let precision = constant_precision(args.get(1))?.unwrap_or(0); + let return_types = InferredTypeSet::singleton(return_type_for_precision(precision)); + if target_type.intersect(return_types).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Round, + expected: target_type, + got: return_types, + }); + } + Ok(return_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let Some(precision) = constant_precision(args.get(1))? else { + return Ok(TypedExpr::none()); + }; + let return_type = return_type_for_precision(precision); + debug_assert!(target_type_set.contains(return_type)); + let arg = context.apply_types(&args[0], InferredTypeSet::NUMERICAL)?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(RoundFnCall { + arg: Box::new(arg), + precision, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "ROUND {} {}i64", self.arg, self.precision) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let arg = context.compile_expr(&self.arg, builder)?; + let precision = builder.ins().iconst(types::I64, self.precision); + let string_len = builder.ins().iconst(types::I64, 0); + + if return_type == VarType::F64 { + let float_value = match self.arg.return_type { + VarType::I64 => builder.ins().fcvt_from_sint(types::F64, arg.value), + VarType::U64 => builder.ins().fcvt_from_uint(types::F64, arg.value), + VarType::F64 => arg.value, + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Round, + return_type: self.arg.return_type, + }); + } + }; + let call = builder.ins().call( + context.native_functions().round_float(), + &[float_value, precision], + ); + return Ok(LoweredValue { + value: builder.inst_results(call)[0], + is_present: arg.is_present, + string_len, + }); + } + + if return_type != VarType::I64 { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Round, + return_type, + }); + } + let call = match self.arg.return_type { + VarType::I64 | VarType::U64 => { + let is_signed = builder + .ins() + .iconst(types::I64, i64::from(self.arg.return_type == VarType::I64)); + builder.ins().call( + context.native_functions().round_int_to_i64(), + &[arg.value, is_signed, precision], + ) + } + VarType::F64 => builder.ins().call( + context.native_functions().round_float_to_i64(), + &[arg.value, precision], + ), + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Round, + return_type: self.arg.return_type, + }); + } + }; + let value = builder.inst_results(call)[0]; + let native_is_present = builder.inst_results(call)[1]; + let native_is_present = builder + .ins() + .icmp_imm_u(IntCC::NotEqual, native_is_present, 0); + let is_present = builder.ins().band(arg.is_present, native_is_present); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbols(jit_builder: &mut JITBuilder) { + jit_builder.symbol(ROUND_FLOAT_SYMBOL, round_float as *const u8); + jit_builder.symbol(ROUND_INT_TO_I64_SYMBOL, round_int_to_i64 as *const u8); + jit_builder.symbol(ROUND_FLOAT_TO_I64_SYMBOL, round_float_to_i64 as *const u8); +} + +pub(super) struct NativeRoundFunctions { + pub(super) round_float: FuncRef, + pub(super) round_int_to_i64: FuncRef, + pub(super) round_float_to_i64: FuncRef, +} + +pub(super) fn declare_native_functions( + module: &mut JITModule, + function: &mut CraneliftFunction, +) -> Result { + let mut round_float_signature = module.make_signature(); + round_float_signature + .params + .extend([AbiParam::new(types::F64), AbiParam::new(types::I64)]); + round_float_signature + .returns + .push(AbiParam::new(types::F64)); + let round_float_id = + module.declare_function(ROUND_FLOAT_SYMBOL, Linkage::Import, &round_float_signature)?; + + let mut round_int_signature = module.make_signature(); + round_int_signature.params.extend([ + AbiParam::new(types::I64), + AbiParam::new(types::I64), + AbiParam::new(types::I64), + ]); + round_int_signature + .returns + .extend([AbiParam::new(types::I64), AbiParam::new(types::I64)]); + let round_int_id = module.declare_function( + ROUND_INT_TO_I64_SYMBOL, + Linkage::Import, + &round_int_signature, + )?; + + let mut round_float_to_i64_signature = module.make_signature(); + round_float_to_i64_signature + .params + .extend([AbiParam::new(types::F64), AbiParam::new(types::I64)]); + round_float_to_i64_signature + .returns + .extend([AbiParam::new(types::I64), AbiParam::new(types::I64)]); + let round_float_to_i64_id = module.declare_function( + ROUND_FLOAT_TO_I64_SYMBOL, + Linkage::Import, + &round_float_to_i64_signature, + )?; + + Ok(NativeRoundFunctions { + round_float: module.declare_func_in_func(round_float_id, function), + round_int_to_i64: module.declare_func_in_func(round_int_id, function), + round_float_to_i64: module.declare_func_in_func(round_float_to_i64_id, function), + }) +} + +#[repr(C)] +struct RawI64 { + value: i64, + is_present: usize, +} + +impl RawI64 { + fn none() -> Self { + Self { + value: 0, + is_present: 0, + } + } + + fn some(value: i64) -> Self { + Self { + value, + is_present: 1, + } + } +} + +extern "C" fn round_float(value: f64, precision: i64) -> f64 { + round_f64_with_precision(value, precision) +} + +extern "C" fn round_int_to_i64(value: u64, is_signed: usize, precision: i64) -> RawI64 { + let value = if is_signed != 0 { + value as i64 + } else { + let Ok(value) = i64::try_from(value) else { + return RawI64::none(); + }; + value + }; + match round_i64_with_precision(value, precision) { + Some(value) => RawI64::some(value), + None => RawI64::none(), + } +} + +extern "C" fn round_float_to_i64(value: f64, precision: i64) -> RawI64 { + let rounded = round_f64_with_precision(value, precision); + if !rounded.is_finite() || rounded < i64::MIN as f64 || rounded >= -(i64::MIN as f64) { + return RawI64::none(); + } + RawI64::some(rounded as i64) +} + +fn round_i64_with_precision(value: i64, precision: i64) -> Option { + if precision >= 0 { + return Some(value); + } + let decimal_places = precision.unsigned_abs(); + if decimal_places > 19 { + return Some(0); + } + let factor = 10i128.pow(decimal_places as u32); + let value = i128::from(value); + let quotient = value / factor; + let remainder = value % factor; + let adjustment = if remainder.abs() * 2 >= factor { + value.signum() + } else { + 0 + }; + i64::try_from((quotient + adjustment) * factor).ok() +} + +fn round_f64_with_precision(value: f64, precision: i64) -> f64 { + if !value.is_finite() { + return value; + } + if precision == 0 { + return value.round(); + } + if precision > 0 { + let Ok(decimal_places) = u32::try_from(precision) else { + return value; + }; + if decimal_places > 308 { + return value; + } + let factor = 10f64.powi(decimal_places as i32); + let scaled = value * factor; + if !scaled.is_finite() { + return value; + } + return scaled.round() / factor; + } + + let decimal_places = precision.unsigned_abs(); + if decimal_places > 308 { + return 0.0f64.copysign(value); + } + let factor = 10f64.powi(decimal_places as i32); + (value / factor).round() * factor +} + +impl From for FnCallEnum { + fn from(call: RoundFnCall) -> Self { + FnCallEnum::Round(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval_i64(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable i64 value. + unsafe { compiled.call(&[]).as_i64() } + } + + fn eval_f64(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable f64 value. + unsafe { compiled.call(&[]).as_f64() } + } + + #[test] + fn test_signature_precision_and_output_type() { + let expression = deserialize("(ROUND value)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + assert_eq!(inferred.get("value"), Some(&InferredTypeSet::NUMERICAL)); + + for expression in ["(ROUND)", "(ROUND 1i64 2i64 3i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Round, + expected: 2, + .. + }) + )); + } + + for (expression, expected) in [ + ("(ROUND 1.2f64)", VarType::I64), + ("(ROUND 1.2f64 0i64)", VarType::I64), + ("(ROUND 1.2f64 -1i64)", VarType::I64), + ("(ROUND 1.2f64 1i64)", VarType::F64), + ] { + let expression = deserialize(expression).unwrap(); + assert_eq!( + compile(&expression, &HashMap::new()).unwrap().result_type(), + expected + ); + } + } + + #[test] + fn test_units_and_negative_precision_round_away_from_zero() { + assert_eq!(eval_i64("(ROUND 1.4f64)"), Some(1)); + assert_eq!(eval_i64("(ROUND 1.5f64)"), Some(2)); + assert_eq!(eval_i64("(ROUND -1.4f64)"), Some(-1)); + assert_eq!(eval_i64("(ROUND -1.5f64)"), Some(-2)); + assert_eq!(eval_i64("(ROUND 149i64 -2i64)"), Some(100)); + assert_eq!(eval_i64("(ROUND 150i64 -2i64)"), Some(200)); + assert_eq!(eval_i64("(ROUND -150i64 -2i64)"), Some(-200)); + assert_eq!( + eval_i64("(ROUND -9223372036854775808i64 0i64)"), + Some(i64::MIN) + ); + } + + #[test] + fn test_positive_precision_returns_float() { + assert_eq!(eval_f64("(ROUND 1.234f64 2i64)"), Some(1.23)); + assert_eq!(eval_f64("(ROUND 1.235f64 2i64)"), Some(1.24)); + assert_eq!(eval_f64("(ROUND -1.235f64 2i64)"), Some(-1.24)); + assert_eq!(eval_f64("(ROUND 123i64 2i64)"), Some(123.0)); + assert_eq!(eval_f64("(ROUND 1e308f64 2i64)"), Some(1e308)); + assert_eq!(eval_f64("(ROUND 1.25f64 400i64)"), Some(1.25)); + } + + #[test] + fn test_null_nonfinite_and_range_edges() { + assert_eq!(eval_i64("(ROUND none)"), None); + assert_eq!(eval_i64("(ROUND 1.2f64 none)"), None); + assert_eq!(eval_i64("(ROUND nanf64)"), None); + assert_eq!(eval_i64("(ROUND inff64)"), None); + assert_eq!(eval_i64("(ROUND 9223372036854775808f64)"), None); + assert_eq!(eval_i64("(ROUND 18446744073709551615u64)"), None); + assert_eq!(eval_i64("(ROUND 123i64 -400i64)"), Some(0)); + + assert!(eval_f64("(ROUND nanf64 2i64)").unwrap().is_nan()); + assert_eq!(eval_f64("(ROUND inff64 2i64)"), Some(f64::INFINITY)); + } + + #[test] + fn test_runtime_null_and_integer_rounding() { + let expression = deserialize("(ROUND value -1i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::I64)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable i64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(155i64)]).as_i64() }, + Some(160) + ); + // SAFETY: The compiled expression expects one nullable i64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_i64() }, + None + ); + } +} diff --git a/jitexpr/src/functions/split_after.rs b/jitexpr/src/functions/split_after.rs new file mode 100644 index 0000000000..74adffc1a5 --- /dev/null +++ b/jitexpr/src/functions/split_after.rs @@ -0,0 +1,340 @@ +//! `SPLIT_AFTER(input, separator, occurrence)` returns the part of a string after a separator. +//! +//! The separator must be a string constant. The optional occurrence must be a nonnegative integer +//! constant, uses zero-based indexing, and defaults to zero. Matches are literal, case-sensitive, +//! and non-overlapping. A missing occurrence returns a present empty string, while an empty +//! separator returns the input unchanged. Null input, a null constant, or a negative occurrence +//! returns null. + +use std::collections::HashMap; +use std::sync::Arc; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder, IntCC}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_split_after"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct SplitAfterFnCall { + input: Box, + separator: Arc, + occurrence: usize, +} + +fn constant_occurrence( + expression: Option<&UntypedExpr>, +) -> Result, super::InvalidFunctionCall> { + let Some(expression) = expression else { + return Ok(Some(0)); + }; + let UntypedExpr::Literal(literal) = expression else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::I64, + }); + }; + if !literal.is_none() && !literal.types().contains(VarType::I64) { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::I64, + }); + } + Ok(match literal { + Literal::I64(value) => usize::try_from(*value).ok(), + Literal::U64(value) => usize::try_from(*value).ok(), + Literal::F64(value) => usize::try_from(*value as i64).ok(), + Literal::None => None, + Literal::Bool(_) | Literal::String(_) => None, + }) +} + +impl FnCall for SplitAfterFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Between { min: 2, max: 3 }; + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::Str, |literal| { + matches!(literal, Literal::String(_) | Literal::None) + })?; + if args.len() == 3 { + super::validate_literal(args, 2, VarType::I64, |literal| { + literal.is_none() || literal.types().contains(VarType::I64) + })?; + } + Ok(()) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::SplitAfter, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if !(2..=3).contains(&args.len()) { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::SplitAfter, + expected: 3, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::STRING, inferred_types)?; + if let Some(occurrence) = args.get(2) { + crate::ast::infer_types_aux(occurrence, InferredTypeSet::I64, inferred_types)?; + } + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input = context.apply_types(&args[0], InferredTypeSet::STRING)?; + let separator = match &args[1] { + UntypedExpr::Literal(Literal::String(separator)) => Arc::clone(separator), + UntypedExpr::Literal(Literal::None) => return Ok(TypedExpr::none()), + _ => { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + .into()); + } + }; + let Some(occurrence) = constant_occurrence(args.get(2))? else { + return Ok(TypedExpr::none()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(SplitAfterFnCall { + input: Box::new(input), + separator, + occurrence, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "SPLIT_AFTER {} ", self.input)?; + crate::compile::format_string_literal(&self.separator, formatter)?; + write!(formatter, " {}u64", self.occurrence) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let input = context.compile_expr(&self.input, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, null); + let separator_ptr = builder + .ins() + .iconst(context.pointer_type(), self.separator.as_ptr() as i64); + let separator_len = builder + .ins() + .iconst(types::I64, self.separator.len() as i64); + let occurrence = builder.ins().iconst(types::I64, self.occurrence as i64); + let call = builder.ins().call( + context.native_functions().split_after(), + &[ + input_ptr, + input.string_len, + separator_ptr, + separator_len, + occurrence, + ], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder.ins().icmp_imm_u(IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(builder: &mut JITBuilder) { + builder.symbol(SYMBOL, split_after 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(types::I64), + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + AbiParam::new(types::I64), + ]); + signature + .returns + .extend([AbiParam::new(pointer_type), AbiParam::new(types::I64)]); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } + + fn some(value: &str) -> Self { + Self { + ptr: value.as_ptr(), + len: value.len(), + } + } +} + +unsafe extern "C" fn split_after( + input_ptr: *const u8, + input_len: usize, + separator_ptr: *const u8, + separator_len: usize, + occurrence: usize, +) -> RawStr { + if input_ptr.is_null() { + return RawStr::none(); + } + // SAFETY: CompiledFn supplies a live UTF-8 input and the typed call owns the separator. + let input = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) }; + // SAFETY: The separator pointer and length come from the call's live Arc. + let separator = unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(separator_ptr, separator_len)) + }; + if separator.is_empty() { + return RawStr::some(input); + } + let Some((separator_start, _)) = input.match_indices(separator).nth(occurrence) else { + return RawStr::some(&input[..0]); + }; + RawStr::some(&input[separator_start + separator.len()..]) +} + +impl From for FnCallEnum { + fn from(call: SplitAfterFnCall) -> Self { + FnCallEnum::SplitAfter(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable string. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_signature_and_occurrences() { + for expression in [ + "(SPLIT_AFTER \"a.b\")", + "(SPLIT_AFTER \"a.b\" \".\" 0i64 1i64)", + ] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::SplitAfter, + expected: 3, + .. + }) + )); + } + + assert_eq!(eval("(SPLIT_AFTER \"a.b.c\" \".\")"), Some("b.c".into())); + assert_eq!( + eval("(SPLIT_AFTER \"a.b.c\" \".\" 0i64)"), + Some("b.c".into()) + ); + assert_eq!(eval("(SPLIT_AFTER \"a.b.c\" \".\" 1i64)"), Some("c".into())); + } + + #[test] + fn test_literal_non_overlapping_and_unicode_matches() { + assert_eq!( + eval("(SPLIT_AFTER \"......\" \"...\" 1i64)"), + Some(String::new()) + ); + assert_eq!(eval("(SPLIT_AFTER \"a...\" \"..\" 0i64)"), Some(".".into())); + assert_eq!(eval("(SPLIT_AFTER \"α→β→γ\" \"→\" 1i64)"), Some("γ".into())); + } + + #[test] + fn test_empty_missing_and_invalid_occurrences() { + assert_eq!( + eval("(SPLIT_AFTER \"abc\" \".\" 0i64)"), + Some(String::new()) + ); + assert_eq!(eval("(SPLIT_AFTER \"abc\" \"\" 3i64)"), Some("abc".into())); + assert_eq!(eval("(SPLIT_AFTER \"a.b\" \".\" -1i64)"), None); + assert_eq!(eval("(SPLIT_AFTER \"a.b\" \".\" none)"), None); + assert_eq!(eval("(SPLIT_AFTER \"a.b\" none)"), None); + } + + #[test] + fn test_runtime_null() { + let expression = deserialize("(SPLIT_AFTER value \".\")").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable string argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("a.b")]).as_str() }, + Some("b") + ); + // SAFETY: The compiled expression expects one nullable string argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + } +} diff --git a/jitexpr/src/functions/split_before.rs b/jitexpr/src/functions/split_before.rs new file mode 100644 index 0000000000..42947b82d9 --- /dev/null +++ b/jitexpr/src/functions/split_before.rs @@ -0,0 +1,351 @@ +//! `SPLIT_BEFORE(input, separator, occurrence)` returns the part of a string before a separator. +//! +//! The separator must be a string constant. The optional occurrence must be a nonnegative integer +//! constant, uses zero-based indexing, and defaults to zero. Matches are literal, case-sensitive, +//! and non-overlapping. A missing occurrence or an empty separator returns a present empty string. +//! Null input, a null constant, or a negative occurrence returns null. + +use std::collections::HashMap; +use std::sync::Arc; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder, IntCC}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_split_before"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct SplitBeforeFnCall { + input: Box, + separator: Arc, + occurrence: usize, +} + +fn constant_occurrence( + expression: Option<&UntypedExpr>, +) -> Result, super::InvalidFunctionCall> { + let Some(expression) = expression else { + return Ok(Some(0)); + }; + let UntypedExpr::Literal(literal) = expression else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::I64, + }); + }; + if !literal.is_none() && !literal.types().contains(VarType::I64) { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::I64, + }); + } + Ok(match literal { + Literal::I64(value) => usize::try_from(*value).ok(), + Literal::U64(value) => usize::try_from(*value).ok(), + Literal::F64(value) => usize::try_from(*value as i64).ok(), + Literal::None => None, + Literal::Bool(_) | Literal::String(_) => None, + }) +} + +impl FnCall for SplitBeforeFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Between { min: 2, max: 3 }; + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::Str, |literal| { + matches!(literal, Literal::String(_) | Literal::None) + })?; + if args.len() == 3 { + super::validate_literal(args, 2, VarType::I64, |literal| { + literal.is_none() || literal.types().contains(VarType::I64) + })?; + } + Ok(()) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::SplitBefore, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if !(2..=3).contains(&args.len()) { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::SplitBefore, + expected: 3, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::STRING, inferred_types)?; + if let Some(occurrence) = args.get(2) { + crate::ast::infer_types_aux(occurrence, InferredTypeSet::I64, inferred_types)?; + } + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input = context.apply_types(&args[0], InferredTypeSet::STRING)?; + let separator = match &args[1] { + UntypedExpr::Literal(Literal::String(separator)) => Arc::clone(separator), + UntypedExpr::Literal(Literal::None) => return Ok(TypedExpr::none()), + _ => { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + .into()); + } + }; + let Some(occurrence) = constant_occurrence(args.get(2))? else { + return Ok(TypedExpr::none()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(SplitBeforeFnCall { + input: Box::new(input), + separator, + occurrence, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(formatter, "SPLIT_BEFORE {} ", self.input)?; + crate::compile::format_string_literal(&self.separator, formatter)?; + write!(formatter, " {}u64", self.occurrence) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let input = context.compile_expr(&self.input, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, null); + let separator_ptr = builder + .ins() + .iconst(context.pointer_type(), self.separator.as_ptr() as i64); + let separator_len = builder + .ins() + .iconst(types::I64, self.separator.len() as i64); + let occurrence = builder.ins().iconst(types::I64, self.occurrence as i64); + let call = builder.ins().call( + context.native_functions().split_before(), + &[ + input_ptr, + input.string_len, + separator_ptr, + separator_len, + occurrence, + ], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder.ins().icmp_imm_u(IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(builder: &mut JITBuilder) { + builder.symbol(SYMBOL, split_before 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(types::I64), + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + AbiParam::new(types::I64), + ]); + signature + .returns + .extend([AbiParam::new(pointer_type), AbiParam::new(types::I64)]); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } + + fn some(value: &str) -> Self { + Self { + ptr: value.as_ptr(), + len: value.len(), + } + } +} + +unsafe extern "C" fn split_before( + input_ptr: *const u8, + input_len: usize, + separator_ptr: *const u8, + separator_len: usize, + occurrence: usize, +) -> RawStr { + if input_ptr.is_null() { + return RawStr::none(); + } + // SAFETY: CompiledFn supplies a live UTF-8 input and the typed call owns the separator. + let input = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) }; + // SAFETY: The separator pointer and length come from the call's live Arc. + let separator = unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(separator_ptr, separator_len)) + }; + if separator.is_empty() { + return RawStr::some(&input[..0]); + } + let Some((separator_start, _)) = input.match_indices(separator).nth(occurrence) else { + return RawStr::some(&input[..0]); + }; + RawStr::some(&input[..separator_start]) +} + +impl From for FnCallEnum { + fn from(call: SplitBeforeFnCall) -> Self { + FnCallEnum::SplitBefore(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable string. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_signature_and_occurrences() { + for expression in [ + "(SPLIT_BEFORE \"a.b\")", + "(SPLIT_BEFORE \"a.b\" \".\" 0i64 1i64)", + ] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::SplitBefore, + expected: 3, + .. + }) + )); + } + + assert_eq!(eval("(SPLIT_BEFORE \"a.b.c\" \".\")"), Some("a".into())); + assert_eq!( + eval("(SPLIT_BEFORE \"a.b.c\" \".\" 0i64)"), + Some("a".into()) + ); + assert_eq!( + eval("(SPLIT_BEFORE \"a.b.c\" \".\" 1i64)"), + Some("a.b".into()) + ); + } + + #[test] + fn test_literal_non_overlapping_and_unicode_matches() { + assert_eq!( + eval("(SPLIT_BEFORE \"......\" \"...\" 1i64)"), + Some("...".into()) + ); + assert_eq!( + eval("(SPLIT_BEFORE \"a...\" \"..\" 0i64)"), + Some("a".into()) + ); + assert_eq!( + eval("(SPLIT_BEFORE \"α→β→γ\" \"→\" 1i64)"), + Some("α→β".into()) + ); + } + + #[test] + fn test_empty_missing_and_invalid_occurrences() { + assert_eq!( + eval("(SPLIT_BEFORE \"abc\" \".\" 0i64)"), + Some(String::new()) + ); + assert_eq!( + eval("(SPLIT_BEFORE \"abc\" \"\" 3i64)"), + Some(String::new()) + ); + assert_eq!(eval("(SPLIT_BEFORE \"a.b\" \".\" -1i64)"), None); + assert_eq!(eval("(SPLIT_BEFORE \"a.b\" \".\" none)"), None); + assert_eq!(eval("(SPLIT_BEFORE \"a.b\" none)"), None); + } + + #[test] + fn test_runtime_null() { + let expression = deserialize("(SPLIT_BEFORE value \".\")").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable string argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("a.b")]).as_str() }, + Some("a") + ); + // SAFETY: The compiled expression expects one nullable string argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + } +} diff --git a/jitexpr/src/functions/sqrt.rs b/jitexpr/src/functions/sqrt.rs new file mode 100644 index 0000000000..a10d3dda98 --- /dev/null +++ b/jitexpr/src/functions/sqrt.rs @@ -0,0 +1,185 @@ +//! `SQRT` computes the square root of one numeric argument. +//! +//! The argument is coerced to `f64`, and the result is always `f64`. Null input returns null. A +//! NaN result also returns null, covering NaN input and negative values other than negative zero. +//! Positive infinity and signed zero retain their IEEE-754 behavior. + +use std::collections::HashMap; + +use cranelift::prelude::{FloatCC, FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct SqrtFnCall { + arg: Box, +} + +impl FnCall for SqrtFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::F64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Sqrt, + expected: target_type, + got: InferredTypeSet::F64, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Sqrt, + expected: 1, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::NUMERICAL, inferred_types)?; + Ok(InferredTypeSet::F64) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + debug_assert!(target_type_set.contains(VarType::F64)); + let arg = context.apply_types(&args[0], InferredTypeSet::F64)?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type: VarType::F64, + ast: TypedExprAst::from_call(SqrtFnCall { arg: Box::new(arg) }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("SQRT", std::iter::once(self.arg.as_ref()), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + if return_type != VarType::F64 { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Sqrt, + return_type, + }); + } + let arg = context.compile_expr(&self.arg, builder)?; + let value = builder.ins().sqrt(arg.value); + let is_nan = builder.ins().fcmp(FloatCC::Unordered, value, value); + let is_not_nan = builder.ins().bxor_imm_u(is_nan, 1); + let is_present = builder.ins().band(arg.is_present, is_not_nan); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: SqrtFnCall) -> Self { + FnCallEnum::Sqrt(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable f64 value. + unsafe { compiled.call(&[]).as_f64() } + } + + #[test] + fn test_signature_and_output_type() { + let expression = deserialize("(SQRT value)").unwrap(); + let inferred = infer_types(&expression).unwrap(); + assert_eq!(inferred.get("value"), Some(&InferredTypeSet::NUMERICAL)); + + for expression in ["(SQRT)", "(SQRT 1i64 2i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Sqrt, + expected: 1, + .. + }) + )); + } + + let expression = deserialize("(SQRT 9i64)").unwrap(); + assert_eq!( + compile(&expression, &HashMap::new()).unwrap().result_type(), + VarType::F64 + ); + } + + #[test] + fn test_numeric_inputs_are_converted_to_float() { + assert_eq!(eval("(SQRT 9i64)"), Some(3.0)); + assert_eq!(eval("(SQRT 2.25f64)"), Some(1.5)); + assert_eq!(eval("(SQRT 2u64)"), Some(2.0f64.sqrt())); + } + + #[test] + fn test_null_nan_and_ieee_edges() { + assert_eq!(eval("(SQRT none)"), None); + assert_eq!(eval("(SQRT -1i64)"), None); + assert_eq!(eval("(SQRT nanf64)"), None); + assert_eq!(eval("(SQRT inff64)"), Some(f64::INFINITY)); + + let negative_zero = eval("(SQRT -0f64)").unwrap(); + assert_eq!(negative_zero.to_bits(), (-0.0f64).to_bits()); + } + + #[test] + fn test_runtime_null_and_negative_input() { + let expression = deserialize("(SQRT value)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::I64)])) + .unwrap() + .context(); + + // SAFETY: The compiled expression expects one nullable i64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(16i64)]).as_f64() }, + Some(4.0) + ); + // SAFETY: The compiled expression expects one nullable i64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(-16i64)]).as_f64() }, + None + ); + // SAFETY: The compiled expression expects one nullable i64 argument. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_f64() }, + None + ); + } +} diff --git a/jitexpr/src/functions/substring.rs b/jitexpr/src/functions/substring.rs new file mode 100644 index 0000000000..203d010ce4 --- /dev/null +++ b/jitexpr/src/functions/substring.rs @@ -0,0 +1,308 @@ +//! `SUBSTRING(input, start, length)` returns a byte-indexed string slice. +//! +//! The start and length arguments must be integer constants. They count UTF-8 bytes, not Unicode +//! scalar values or grapheme clusters. For example, `SUBSTRING("éclair", 2, 5)` returns `"clair"` +//! because `é` occupies bytes 0 and 1. The end is clamped to the input length, an out-of-range +//! start or zero length returns an empty string, and null input propagates. +//! +//! Negative bounds and ranges that split a UTF-8 code point return null. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder, IntCC}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_substring"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct SubstringFnCall { + input: Box, + start: usize, + length: usize, +} + +fn constant_usize( + expression: &UntypedExpr, + argument: usize, +) -> Result, super::InvalidFunctionCall> { + let UntypedExpr::Literal(literal) = expression else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument, + expected: VarType::I64, + }); + }; + if !literal.is_none() && !literal.types().contains(VarType::I64) { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument, + expected: VarType::I64, + }); + } + Ok(match literal { + Literal::I64(value) => usize::try_from(*value).ok(), + Literal::U64(value) => i64::try_from(*value) + .ok() + .and_then(|value| usize::try_from(value).ok()), + Literal::F64(value) if literal.types().contains(VarType::I64) => { + usize::try_from(*value as i64).ok() + } + Literal::None => None, + Literal::Bool(_) | Literal::F64(_) | Literal::String(_) => None, + }) +} + +impl FnCall for SubstringFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(3); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + for index in 1..=2 { + super::validate_literal(args, index, VarType::I64, |literal| { + literal.is_none() || literal.types().contains(VarType::I64) + })?; + } + Ok(()) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Substring, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() != 3 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Substring, + expected: 3, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + crate::ast::infer_types_aux(&args[1], InferredTypeSet::I64, inferred_types)?; + crate::ast::infer_types_aux(&args[2], InferredTypeSet::I64, inferred_types)?; + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input = context.apply_types(&args[0], InferredTypeSet::STRING)?; + let start = constant_usize(&args[1], 2)?; + let length = constant_usize(&args[2], 3)?; + let (Some(start), Some(length)) = (start, length) else { + return Ok(TypedExpr::none()); + }; + if input.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(SubstringFnCall { + input: Box::new(input), + start, + length, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + formatter, + "SUBSTRING {} {}u64 {}u64", + self.input, self.start, self.length + ) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let input = context.compile_expr(&self.input, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, null); + let start = builder + .ins() + .iconst(context.pointer_type(), self.start as i64); + let length = builder + .ins() + .iconst(context.pointer_type(), self.length as i64); + let call = builder.ins().call( + context.native_functions().substring(), + &[input_ptr, input.string_len, start, length], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let native_succeeded = builder.ins().icmp_imm_u(IntCC::NotEqual, value, 0); + let is_present = builder.ins().band(input.is_present, native_succeeded); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, substring 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(types::I64), + AbiParam::new(types::I64), + AbiParam::new(types::I64), + ]); + signature + .returns + .extend([AbiParam::new(pointer_type), AbiParam::new(types::I64)]); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } + + fn some(value: &str) -> Self { + Self { + ptr: value.as_ptr(), + len: value.len(), + } + } +} + +unsafe extern "C" fn substring( + input_ptr: *const u8, + input_len: usize, + start: usize, + length: usize, +) -> RawStr { + if input_ptr.is_null() { + return RawStr::none(); + } + // SAFETY: CompiledFn's call contract guarantees a live UTF-8 string pointer and exact length. + let input = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) }; + let end = start.saturating_add(length).min(input.len()); + if start >= end { + return RawStr::some(&input[..0]); + } + if !input.is_char_boundary(start) || !input.is_char_boundary(end) { + return RawStr::none(); + } + RawStr::some(&input[start..end]) +} + +impl From for FnCallEnum { + fn from(call: SubstringFnCall) -> Self { + FnCallEnum::Substring(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: The expression has no inputs and returns a nullable string. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_signature_and_byte_offsets() { + for expression in [ + "(SUBSTRING \"abc\" 1i64)", + "(SUBSTRING \"abc\" 1i64 2i64 3i64)", + ] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Substring, + expected: 3, + .. + }) + )); + } + assert_eq!(eval("(SUBSTRING \"abcdef\" 1i64 3i64)"), Some("bcd".into())); + assert_eq!( + eval("(SUBSTRING \"éclair\" 2i64 5i64)"), + Some("clair".into()) + ); + } + + #[test] + fn test_empty_and_clamping() { + assert_eq!(eval("(SUBSTRING \"abc\" 1i64 99i64)"), Some("bc".into())); + assert_eq!(eval("(SUBSTRING \"abc\" 3i64 1i64)"), Some(String::new())); + assert_eq!(eval("(SUBSTRING \"abc\" 9i64 1i64)"), Some(String::new())); + assert_eq!(eval("(SUBSTRING \"abcdef\" 2i64 2i64)"), Some("cd".into())); + assert_eq!(eval("(SUBSTRING \"abc\" 1i64 0i64)"), Some(String::new())); + } + + #[test] + fn test_invalid_bounds_and_runtime_null() { + assert_eq!(eval("(SUBSTRING \"é\" 1i64 2i64)"), None); + assert_eq!(eval("(SUBSTRING \"abc\" -1i64 1i64)"), None); + + let expression = deserialize("(SUBSTRING value 0i64 2i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("abc")]).as_str() }, + Some("ab") + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + } +} diff --git a/jitexpr/src/functions/substring_count.rs b/jitexpr/src/functions/substring_count.rs new file mode 100644 index 0000000000..d747776878 --- /dev/null +++ b/jitexpr/src/functions/substring_count.rs @@ -0,0 +1,200 @@ +//! `SUBSTRING_COUNT(haystack, needle)` counts non-overlapping substring occurrences. +//! +//! Both arguments are strings and null propagates. Matching is case-sensitive and byte-based; +//! `SUBSTRING_COUNT("aaaa", "aa")` is 2. An empty needle returns zero rather than counting string +//! boundaries. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_substring_count"; +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct SubstringCountFnCall { + args: Box<[TypedExpr]>, +} + +impl FnCall for SubstringCountFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target: InferredTypeSet, + inferred: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target.intersect(InferredTypeSet::I64).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::SubstringCount, + expected: target, + got: InferredTypeSet::I64, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::SubstringCount, + expected: 2, + got: args.len(), + }); + } + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::STRING, inferred)?; + } + Ok(InferredTypeSet::I64) + } + fn call_with_types( + args: &[UntypedExpr], + _target: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::STRING)) + .collect::, _>>()?; + if args.iter().any(|arg| arg.return_type == VarType::None) { + return Ok(TypedExpr::none()); + } + Ok(TypedExpr { + return_type: VarType::I64, + ast: TypedExprAst::from_call(SubstringCountFnCall { + args: args.into_boxed_slice(), + }), + }) + } + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("SUBSTRING_COUNT", self.args.iter(), formatter) + } + fn emit_cranelift_ir( + &self, + _return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let haystack = context.compile_expr(&self.args[0], builder)?; + let needle = context.compile_expr(&self.args[1], builder)?; + let call = builder.ins().call( + context.native_functions().substring_count(), + &[ + haystack.value, + haystack.string_len, + needle.value, + needle.string_len, + ], + ); + let value = builder.inst_results(call)[0]; + let is_present = builder.ins().band(haystack.is_present, needle.is_present); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} +pub(super) fn register_jit_symbol(builder: &mut JITBuilder) { + builder.symbol(SYMBOL, substring_count as *const u8); +} +pub(super) fn declare_native_function( + module: &mut JITModule, + function: &mut CraneliftFunction, + pointer: Type, +) -> Result { + let mut signature = module.make_signature(); + signature.params.extend([ + AbiParam::new(pointer), + AbiParam::new(types::I64), + AbiParam::new(pointer), + AbiParam::new(types::I64), + ]); + signature.returns.push(AbiParam::new(types::I64)); + let id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(id, function)) +} +unsafe extern "C" fn substring_count( + haystack: *const u8, + haystack_len: usize, + needle: *const u8, + needle_len: usize, +) -> i64 { + if needle_len == 0 { + return 0; + } + let haystack = unsafe { std::slice::from_raw_parts(haystack, haystack_len) }; + let needle = unsafe { std::slice::from_raw_parts(needle, needle_len) }; + let mut count = 0i64; + let mut rest = haystack; + while needle.len() <= rest.len() { + let Some(pos) = rest + .windows(needle.len()) + .position(|window| window == needle) + else { + break; + }; + count += 1; + rest = &rest[pos + needle.len()..]; + } + count +} +impl From for FnCallEnum { + fn from(call: SubstringCountFnCall) -> Self { + FnCallEnum::SubstringCount(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + fn eval(expr: &str) -> Option { + let expr = deserialize(expr).unwrap(); + let mut compiled = compile(&expr, &HashMap::new()).unwrap().context(); + unsafe { compiled.call(&[]).as_i64() } + } + #[test] + fn test_signature_and_counts() { + assert!(matches!( + infer_types(&deserialize("(SUBSTRING_COUNT \"a\")").unwrap()), + Err(TypeError::InvalidNumberOfArguments { + function: Function::SubstringCount, + expected: 2, + .. + }) + )); + assert_eq!(eval("(SUBSTRING_COUNT \"aaaa\" \"aa\")"), Some(2)); + assert_eq!(eval("(SUBSTRING_COUNT \"Abab\" \"ab\")"), Some(1)); + assert_eq!(eval("(SUBSTRING_COUNT \"abc\" \"\")"), Some(0)); + } + #[test] + fn test_runtime_null() { + let expr = deserialize("(SUBSTRING_COUNT text needle)").unwrap(); + let mut compiled = compile( + &expr, + &HashMap::from([("text", VarType::Str), ("needle", VarType::Str)]), + ) + .unwrap() + .context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some("abc"), VariableValue::none()]) + .as_i64() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/subtract.rs b/jitexpr/src/functions/subtract.rs new file mode 100644 index 0000000000..36d284a7ae --- /dev/null +++ b/jitexpr/src/functions/subtract.rs @@ -0,0 +1,253 @@ +//! `SUBTRACT` subtracts one numeric expression from another. +//! +//! It accepts exactly two numeric arguments, `left` and `right`, and computes `left - right`. +//! Both operands are coerced to one common numeric type using the same rules as `ADD`: prefer a +//! common `i64`, then `u64`, and fall back to `f64` when the operands cannot be represented by one +//! integer type. Integer arithmetic uses two's-complement wrapping at 64 bits; floating-point +//! arithmetic preserves IEEE values such as NaN. +//! +//! Null propagation is strict: if either operand is absent, the result is absent. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use super::add::{is_numerical, select_return_type, with_float_fallback}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct SubtractFnCall { + pub(crate) args: Box<[TypedExpr]>, +} + +impl FnCall for SubtractFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(2); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::NUMERICAL).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Subtract, + expected: target_type, + got: InferredTypeSet::NUMERICAL, + }); + } + if args.len() != 2 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Subtract, + expected: 2, + got: args.len(), + }); + } + + let mut return_types = InferredTypeSet::NUMERICAL; + for arg in args { + let arg_types = + crate::ast::infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types)?; + return_types = return_types.intersect(arg_types); + } + return_types = with_float_fallback(return_types); + let constrained_return_types = return_types.intersect(target_type); + if constrained_return_types.is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Subtract, + expected: target_type, + got: return_types, + }); + } + Ok(constrained_return_types) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let mut return_types = InferredTypeSet::NUMERICAL.intersect(target_type_set); + for arg in args { + let arg_types = crate::ast::infer_type_with_variable_types( + arg, + InferredTypeSet::NUMERICAL, + context.variable_types(), + )?; + return_types = return_types.intersect(arg_types); + } + let return_type = select_return_type(with_float_fallback(return_types)); + let typed_args = args + .iter() + .map(|arg| context.apply_types(arg, InferredTypeSet::singleton(return_type))) + .collect::, _>>()?; + if typed_args + .iter() + .any(|typed_arg| !is_numerical(typed_arg.return_type)) + { + return Ok(TypedExpr::none()); + } + + Ok(TypedExpr { + return_type, + ast: TypedExprAst::from_call(SubtractFnCall { + args: typed_args.into_boxed_slice(), + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + &mut self.args + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("SUBTRACT", self.args.iter(), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let left = context.compile_expr(&self.args[0], builder)?; + let right = context.compile_expr(&self.args[1], builder)?; + let value = match return_type { + VarType::U64 | VarType::I64 => builder.ins().isub(left.value, right.value), + VarType::F64 => builder.ins().fsub(left.value, right.value), + _ => { + return Err(CompileError::UnsupportedFunctionType { + function: Function::Subtract, + return_type, + }); + } + }; + let is_present = builder.ins().band(left.is_present, right.is_present); + Ok(LoweredValue { + value, + is_present, + string_len: builder.ins().iconst(types::I64, 0), + }) + } +} + +impl From for FnCallEnum { + fn from(call: SubtractFnCall) -> Self { + FnCallEnum::Subtract(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + #[test] + fn test_infer_types_requires_two_numeric_arguments() { + let expression = deserialize("(SUBTRACT left right)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!( + inferred_types.get("left"), + Some(&InferredTypeSet::NUMERICAL) + ); + assert_eq!( + inferred_types.get("right"), + Some(&InferredTypeSet::NUMERICAL) + ); + + for expression in ["(SUBTRACT 1i64)", "(SUBTRACT 1i64 2i64 3i64)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Subtract, + expected: 2, + .. + }) + )); + } + } + + #[test] + fn test_signed_unsigned_and_float_subtraction() { + let expression = deserialize("(SUBTRACT 7i64 10i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::I64); + // SAFETY: The expression has no inputs and returns i64. + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, Some(-3)); + + let expression = deserialize("(SUBTRACT 9223372036854775810u64 1u64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::U64); + // SAFETY: The expression has no inputs and returns u64. + assert_eq!( + unsafe { compiled.call(&[]).as_u64() }, + Some(9223372036854775809) + ); + + let expression = deserialize("(SUBTRACT 7.5f64 2f64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::F64); + // SAFETY: The expression has no inputs and returns f64. + assert_eq!(unsafe { compiled.call(&[]).as_f64() }, Some(5.5)); + } + + #[test] + fn test_unsigned_subtraction_wraps() { + let expression = deserialize("(SUBTRACT value 1u64)").unwrap(); + let variable_types = HashMap::from([("value", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(0u64)]).as_u64() }, + Some(u64::MAX) + ); + } + + #[test] + fn test_runtime_null_propagation_and_mixed_literal_coercion() { + let expression = deserialize("(SUBTRACT value 1i64)").unwrap(); + let variable_types = HashMap::from([("value", VarType::U64)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!(compiled.result_type(), VarType::U64); + + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::some(5u64)]).as_u64() }, + Some(4) + ); + // SAFETY: The input and output types match the compiled signature. + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_u64() }, + None + ); + } + + #[test] + fn test_compile_time_none_propagates() { + let expression = deserialize("(SUBTRACT none 1i64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + + assert_eq!(compiled.result_type(), VarType::None); + // SAFETY: The expression has no inputs and returns an absent value. + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, None); + } + + #[test] + fn test_compile_time_i64() { + let expression = deserialize("(SUBTRACT 1u64 10u64)").unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + + assert_eq!(compiled.result_type(), VarType::I64); + // SAFETY: The expression has no inputs and returns an absent value. + assert_eq!(unsafe { compiled.call(&[]).as_i64() }, Some(-9i64)); + } +} diff --git a/jitexpr/src/functions/text_join.rs b/jitexpr/src/functions/text_join.rs new file mode 100644 index 0000000000..fbf99f4735 --- /dev/null +++ b/jitexpr/src/functions/text_join.rs @@ -0,0 +1,146 @@ +//! `TEXT_JOIN` is another name for the `CONCAT` string-joining operation. +//! +//! It accepts `TEXT_JOIN(delimiter, ignore_empty, value1, value2, ...)`, with at least four total +//! arguments. The first two arguments must be string literals. Only a case-insensitive `"true"` +//! enables empty-value skipping; other flag strings mean false. Any null value makes the +//! result null. Delimiters are inserted only after nonempty output bytes have been written, so a +//! leading empty value does not produce a leading delimiter even when empty values are retained. +//! +//! Constructed bytes use the fixed call arena and arena exhaustion returns null. + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +use super::concat::{self, JoinArguments}; +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct TextJoinFnCall { + arguments: JoinArguments, +} + +impl FnCall for TextJoinFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::AtLeast(4); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + concat::validate_join_args(args) + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + concat::infer_join_types(Function::TextJoin, args, target_type, inferred_types) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let Some(arguments) = concat::apply_join_types("TEXT_JOIN", args, context)? else { + return Ok(TypedExpr::none()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(TextJoinFnCall { arguments }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + self.arguments.args_mut() + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + self.arguments.serialize("TEXT_JOIN", formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + self.arguments.emit_cranelift_ir(context, builder) + } +} + +impl From for FnCallEnum { + fn from(call: TextJoinFnCall) -> Self { + FnCallEnum::TextJoin(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expression: &str) -> Option { + let expression = deserialize(expression).unwrap(); + let mut compiled = compile(&expression, &HashMap::new()).unwrap().context(); + // SAFETY: These expressions have no inputs and return nullable strings. + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_requires_four_or_more_string_arguments() { + let expression = deserialize(r#"(TEXT_JOIN "," "false" left right)"#).unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("left"), Some(&InferredTypeSet::STRING)); + assert_eq!(inferred_types.get("right"), Some(&InferredTypeSet::STRING)); + + let expression = deserialize(r#"(TEXT_JOIN "," "false" one)"#).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::TextJoin, + expected: 4, + .. + }) + )); + } + + #[test] + fn test_matches_concat_behavior() { + assert_eq!( + eval(r#"(TEXT_JOIN " / " "TRUE" "one" "" "two")"#).as_deref(), + Some("one / two") + ); + assert_eq!( + eval(r#"(TEXT_JOIN "," "false" "" "two")"#).as_deref(), + Some("two") + ); + assert_eq!( + eval(r#"(TEXT_JOIN "," "false" "two" "")"#).as_deref(), + Some("two,") + ); + } + + #[test] + fn test_runtime_null_is_strict() { + let expression = deserialize(r#"(TEXT_JOIN "" "false" left right)"#).unwrap(); + let variable_types = HashMap::from([("left", VarType::Str), ("right", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::none(), VariableValue::some("right")]) + .as_str() + }, + None + ); + } +} diff --git a/jitexpr/src/functions/trim.rs b/jitexpr/src/functions/trim.rs new file mode 100644 index 0000000000..9558b3ef61 --- /dev/null +++ b/jitexpr/src/functions/trim.rs @@ -0,0 +1,323 @@ +//! `TRIM` removes a whole delimiter once from selected ends of a string. +//! +//! It takes `(input, delimiter, mode)`, where the latter two arguments are string literals emitted +//! by the calculated-field producer. Mode is case-insensitive `leading`, `trailing`, or `both`. +//! Unlike SQL character-set trimming, `"xy"` is removed only as one complete prefix/suffix and at +//! most once per selected side. Null input propagates; an empty delimiter leaves the input intact. + +use std::collections::HashMap; +use std::sync::Arc; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder, IntCC}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, Literal, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_string_trim"; + +#[derive(Clone, Copy, Debug, PartialEq)] +enum TrimMode { + Leading, + Trailing, + Both, +} + +impl TrimMode { + fn as_str(self) -> &'static str { + match self { + TrimMode::Leading => "leading", + TrimMode::Trailing => "trailing", + TrimMode::Both => "both", + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct TrimFnCall { + input: Box, + delimiter: Arc, + mode: TrimMode, +} + +impl FnCall for TrimFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(3); + + fn validate_args(args: &[UntypedExpr]) -> Result<(), super::InvalidFunctionCall> { + Self::ARG_COUNT.validate(args)?; + super::validate_literal(args, 1, VarType::Str, |literal| { + matches!(literal, Literal::String(_)) + })?; + let UntypedExpr::Literal(Literal::String(mode)) = &args[2] else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::Str, + }); + }; + if ["leading", "trailing", "both"] + .iter() + .any(|valid_mode| mode.eq_ignore_ascii_case(valid_mode)) + { + Ok(()) + } else { + Err(super::InvalidFunctionCall::InvalidLiteralValue { + argument: 3, + expected: "leading, trailing, or both", + }) + } + } + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Trim, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() != 3 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Trim, + expected: 3, + got: args.len(), + }); + } + for arg in args { + crate::ast::infer_types_aux(arg, InferredTypeSet::STRING, inferred_types)?; + } + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let input = context.apply_types(&args[0], InferredTypeSet::STRING)?; + if input.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + let UntypedExpr::Literal(Literal::String(delimiter)) = &args[1] else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 2, + expected: VarType::Str, + } + .into()); + }; + let UntypedExpr::Literal(Literal::String(mode)) = &args[2] else { + return Err(super::InvalidFunctionCall::ExpectedLiteral { + argument: 3, + expected: VarType::Str, + } + .into()); + }; + let mode = if mode.eq_ignore_ascii_case("leading") { + TrimMode::Leading + } else if mode.eq_ignore_ascii_case("trailing") { + TrimMode::Trailing + } else if mode.eq_ignore_ascii_case("both") { + TrimMode::Both + } else { + return Err(super::InvalidFunctionCall::InvalidLiteralValue { + argument: 3, + expected: "leading, trailing, or both", + } + .into()); + }; + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(TrimFnCall { + input: Box::new(input), + delimiter: Arc::clone(delimiter), + mode, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.input) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + use std::fmt::Write as _; + + write!(formatter, "TRIM {} ", self.input)?; + crate::compile::format_string_literal(&self.delimiter, formatter)?; + formatter.write_char(' ')?; + crate::compile::format_string_literal(self.mode.as_str(), formatter) + } + + fn emit_cranelift_ir( + &self, + _return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + let input = context.compile_expr(&self.input, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(input.is_present, input.value, null); + let delimiter_ptr = builder + .ins() + .iconst(context.pointer_type(), self.delimiter.as_ptr() as i64); + let delimiter_len = builder + .ins() + .iconst(types::I64, self.delimiter.len() as i64); + let mode = builder.ins().iconst(types::I64, self.mode as i64); + let call = builder.ins().call( + context.native_functions().string_trim(), + &[ + input_ptr, + input.string_len, + delimiter_ptr, + delimiter_len, + mode, + ], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder.ins().icmp_imm_u(IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(builder: &mut JITBuilder) { + builder.symbol(SYMBOL, string_trim 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(types::I64), + AbiParam::new(pointer_type), + AbiParam::new(types::I64), + AbiParam::new(types::I64), + ]); + signature + .returns + .extend([AbiParam::new(pointer_type), AbiParam::new(types::I64)]); + let id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(id, function)) +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} +unsafe extern "C" fn string_trim( + input_ptr: *const u8, + input_len: usize, + delimiter_ptr: *const u8, + delimiter_len: usize, + mode: u64, +) -> RawStr { + if input_ptr.is_null() { + return RawStr { + ptr: std::ptr::null(), + len: 0, + }; + } + let input = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) }; + let delimiter = unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(delimiter_ptr, delimiter_len)) + }; + let leading = mode == TrimMode::Leading as u64 || mode == TrimMode::Both as u64; + let trailing = mode == TrimMode::Trailing as u64 || mode == TrimMode::Both as u64; + let value = if leading { + input.strip_prefix(delimiter).unwrap_or(input) + } else { + input + }; + let value = if trailing { + value.strip_suffix(delimiter).unwrap_or(value) + } else { + value + }; + RawStr { + ptr: value.as_ptr(), + len: value.len(), + } +} + +impl From for FnCallEnum { + fn from(call: TrimFnCall) -> Self { + FnCallEnum::Trim(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::compile; + use crate::types::VariableValue; + + fn eval(expr: &str) -> Option { + let expr = deserialize(expr).unwrap(); + let mut compiled = compile(&expr, &HashMap::new()).unwrap().context(); + unsafe { compiled.call(&[]).as_str().map(str::to_owned) } + } + + #[test] + fn test_signature_and_modes() { + assert!(infer_types(&deserialize("(TRIM value \"xy\" \"both\")").unwrap()).is_ok()); + assert!(matches!( + infer_types(&deserialize("(TRIM value)").unwrap()), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Trim, + expected: 3, + .. + }) + )); + assert_eq!( + eval("(TRIM \"xyhelloxy\" \"xy\" \"both\")"), + Some("hello".into()) + ); + assert_eq!( + eval("(TRIM \"xyxy\" \"xy\" \"leading\")"), + Some("xy".into()) + ); + assert_eq!( + eval("(TRIM \"xyxy\" \"xy\" \"TRAILING\")"), + Some("xy".into()) + ); + assert_eq!(eval("(TRIM \"hello\" \"\" \"both\")"), Some("hello".into())); + } + + #[test] + fn test_runtime_null_and_unicode() { + let expr = deserialize("(TRIM value \"é\" \"both\")").unwrap(); + let mut compiled = compile(&expr, &HashMap::from([("value", VarType::Str)])) + .unwrap() + .context(); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("éhelloé")]).as_str() }, + Some("hello") + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + } +} diff --git a/jitexpr/src/functions/upper.rs b/jitexpr/src/functions/upper.rs new file mode 100644 index 0000000000..0270273df9 --- /dev/null +++ b/jitexpr/src/functions/upper.rs @@ -0,0 +1,288 @@ +//! `UPPER` constructs the Unicode-uppercase form of a string. +//! +//! It accepts exactly one string and returns a newly allocated string without mutating its input. +//! It applies one-to-one Unicode simple case mappings; mappings that would expand one character +//! into several (for example `ß` to `SS`) are not applied. Null input returns null, while an empty +//! input returns a present empty string. +//! +//! Constructed bytes live in the caller's fixed-capacity string arena; arena exhaustion returns +//! null. + +use std::collections::HashMap; + +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, types}; +use cranelift::frontend::FunctionBuilder; +use cranelift::prelude::{AbiParam, InstBuilder}; +use cranelift_jit::{JITBuilder, JITModule}; +use cranelift_module::{Linkage, Module}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{ + CompileError, CompileFnBuilder, LoweredValue, LoweringContext, StringArena, TypedExpr, + TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::VarType; + +const SYMBOL: &str = "jitexpr_string_uppercase"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct UpperFnCall { + arg: Box, +} + +impl FnCall for UpperFnCall { + const ARG_COUNT: super::ArgumentCount = super::ArgumentCount::Exactly(1); + + fn infer_types<'a>( + args: &'a [UntypedExpr], + target_type: InferredTypeSet, + inferred_types: &mut HashMap<&'a str, InferredTypeSet>, + ) -> Result { + if target_type.intersect(InferredTypeSet::STRING).is_none() { + return Err(TypeError::WrongFunctionReturnType { + function: Function::Upper, + expected: target_type, + got: InferredTypeSet::STRING, + }); + } + if args.len() != 1 { + return Err(TypeError::InvalidNumberOfArguments { + function: Function::Upper, + expected: 1, + got: args.len(), + }); + } + crate::ast::infer_types_aux(&args[0], InferredTypeSet::STRING, inferred_types)?; + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + _target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + Self::ARG_COUNT.validate(args)?; + let arg = context.apply_types(&args[0], InferredTypeSet::STRING)?; + if arg.return_type == VarType::None { + return Ok(TypedExpr::none()); + } + debug_assert_eq!(arg.return_type, VarType::Str); + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(UpperFnCall { arg: Box::new(arg) }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.arg) + } + + fn serialize(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + crate::compile::format_function_call("UPPER", std::iter::once(self.arg.as_ref()), formatter) + } + + fn emit_cranelift_ir( + &self, + return_type: VarType, + context: &mut LoweringContext<'_>, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + debug_assert_eq!(return_type, VarType::Str); + let arg = context.compile_expr(&self.arg, builder)?; + let null = builder.ins().iconst(context.pointer_type(), 0); + let input_ptr = builder.ins().select(arg.is_present, arg.value, null); + let string_arena_ptr = context.string_arena_ptr(builder); + let call = builder.ins().call( + context.native_functions().string_uppercase(), + &[input_ptr, arg.string_len, string_arena_ptr], + ); + let value = builder.inst_results(call)[0]; + let string_len = builder.inst_results(call)[1]; + let is_present = builder + .ins() + .icmp_imm_u(cranelift::prelude::IntCC::NotEqual, value, 0); + Ok(LoweredValue { + value, + is_present, + string_len, + }) + } +} + +pub(super) fn register_jit_symbol(jit_builder: &mut JITBuilder) { + jit_builder.symbol(SYMBOL, string_uppercase 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(types::I64), + AbiParam::new(pointer_type), + ]); + signature.returns.push(AbiParam::new(pointer_type)); + signature.returns.push(AbiParam::new(types::I64)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +#[repr(C)] +struct RawStr { + ptr: *const u8, + len: usize, +} + +impl RawStr { + fn none() -> Self { + Self { + ptr: std::ptr::null(), + len: 0, + } + } +} + +unsafe extern "C" fn string_uppercase( + input_ptr: *const u8, + input_len: usize, + string_arena: *mut StringArena, +) -> RawStr { + if input_ptr.is_null() || string_arena.is_null() { + return RawStr::none(); + } + + let output_len = { + // SAFETY: Generated code passes a live UTF-8 string and its exact byte length for every + // present string value. + let input = unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) + }; + let mut output_len = 0usize; + for character in input.chars() { + let uppercase = simple_uppercase(character); + let Some(next_len) = output_len.checked_add(uppercase.len_utf8()) else { + return RawStr::none(); + }; + output_len = next_len; + } + output_len + }; + + // The input borrow has ended. Nested calls may pass an earlier, disjoint arena allocation. + // SAFETY: The caller exclusively borrows and passes this arena for the duration of the call. + let Some(output_ptr) = (unsafe { &mut *string_arena }).allocate(output_len) else { + return RawStr::none(); + }; + + // SAFETY: The arena never reallocates and the newly allocated output range does not overlap an + // earlier arena-backed input range. + let input = + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(input_ptr, input_len)) }; + let mut written = 0usize; + for character in input.chars() { + let uppercase = simple_uppercase(character); + let mut encoded = [0; 4]; + let encoded = uppercase.encode_utf8(&mut encoded).as_bytes(); + // SAFETY: `output_len` was computed from this exact transformation. + unsafe { + std::ptr::copy_nonoverlapping(encoded.as_ptr(), output_ptr.add(written), encoded.len()); + } + written += encoded.len(); + } + debug_assert_eq!(written, output_len); + RawStr { + ptr: output_ptr, + len: output_len, + } +} + +/// Retain the input when the uppercase mapping expands to preserve one-to-one mappings. +fn simple_uppercase(character: char) -> char { + let mut uppercase = character.to_uppercase(); + let first = uppercase.next().unwrap_or(character); + if uppercase.next().is_some() { + character + } else { + first + } +} + +impl From for FnCallEnum { + fn from(call: UpperFnCall) -> Self { + FnCallEnum::Upper(call) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::{deserialize, infer_types}; + use crate::compile::{STRING_ARENA_CAPACITY, compile}; + use crate::types::VariableValue; + + #[test] + fn test_requires_one_string_argument() { + let expression = deserialize("(UPPER value)").unwrap(); + let inferred_types = infer_types(&expression).unwrap(); + assert_eq!(inferred_types.get("value"), Some(&InferredTypeSet::STRING)); + + for expression in ["(UPPER)", "(UPPER one two)"] { + let expression = deserialize(expression).unwrap(); + assert!(matches!( + infer_types(&expression), + Err(TypeError::InvalidNumberOfArguments { + function: Function::Upper, + expected: 1, + .. + }) + )); + } + } + + #[test] + fn test_uppercases_unicode_without_full_case_expansion_or_mutation() { + let expression = deserialize("(UPPER value)").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + let input_string = String::from("café Straße ı"); + + let output = unsafe { compiled.call(&[VariableValue::some(input_string.as_str())]) }; + + assert_eq!(unsafe { output.as_str() }, Some("CAFÉ STRAßE I")); + assert_eq!(input_string, "café Straße ı"); + } + + #[test] + fn test_null_empty_nested_and_arena_exhaustion() { + let expression = deserialize("(UPPER (UPPER value))").unwrap(); + let variable_types = HashMap::from([("value", VarType::Str)]); + let mut compiled = compile(&expression, &variable_types).unwrap().context(); + + assert_eq!( + unsafe { compiled.call(&[VariableValue::none()]).as_str() }, + None + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("")]).as_str() }, + Some("") + ); + assert_eq!( + unsafe { compiled.call(&[VariableValue::some("Hello")]).as_str() }, + Some("HELLO") + ); + + let too_large = "a".repeat(STRING_ARENA_CAPACITY / 2 + 1); + assert_eq!( + unsafe { + compiled + .call(&[VariableValue::some(too_large.as_str())]) + .as_str() + }, + None + ); + } +} diff --git a/jitexpr/src/lib.rs b/jitexpr/src/lib.rs new file mode 100644 index 0000000000..b6aa1517c3 --- /dev/null +++ b/jitexpr/src/lib.rs @@ -0,0 +1,34 @@ +pub mod ast; +pub mod compile; +pub mod types; + +mod functions; + +#[cfg(test)] +pub(crate) fn typed_expr_from_str( + untyped_expr: &str, + variable_types: &std::collections::HashMap<&str, types::VarType>, +) -> compile::TypedExpr { + let untyped_expr = ast::deserialize(untyped_expr).unwrap(); + let mut context = compile::CompileFnBuilder::new(variable_types); + context + .apply_types(&untyped_expr, ast::InferredTypeSet::ALL) + .unwrap() +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::types::VarType; + + #[test] + fn test_typed_expr_from_str() { + let variable_types = HashMap::from([("value", VarType::U64)]); + + let typed_expr = typed_expr_from_str("(ADD value 1i64)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::U64); + } +} diff --git a/jitexpr/src/types.rs b/jitexpr/src/types.rs new file mode 100644 index 0000000000..2b7466742b --- /dev/null +++ b/jitexpr/src/types.rs @@ -0,0 +1,331 @@ +//! Source types and nullable runtime value representations. + +/// A value type supported by compiled expressions. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Ord, PartialOrd)] +pub enum VarType { + Bool, + F64, + U64, + I64, + Str, + None, +} + +/// The payload of a primitive runtime value. +/// +/// This union is deliberately untagged. The corresponding +/// [`crate::compile::TypedVariable`] identifies the active payload field. +#[repr(C)] +#[derive(Clone, Copy)] +pub union VariablePrimitive { + pub boolean: bool, + pub float: f64, + pub int_u64: u64, + pub int_i64: i64, +} + +impl From for VariablePrimitive { + fn from(value: bool) -> Self { + VariablePrimitive { boolean: value } + } +} + +impl From for VariablePrimitive { + fn from(value: f64) -> Self { + VariablePrimitive { float: value } + } +} + +impl From for VariablePrimitive { + fn from(value: u64) -> Self { + VariablePrimitive { int_u64: value } + } +} + +impl From for VariablePrimitive { + fn from(value: i64) -> Self { + VariablePrimitive { int_i64: value } + } +} + +impl Default for VariablePrimitive { + fn default() -> Self { + VariablePrimitive { int_u64: 0 } + } +} + +/// A nullable primitive value. +/// +/// `value` is meaningful only when `is_present` is true. +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct VariablePrimitiveOpt { + pub value: VariablePrimitive, + pub is_present: bool, +} + +impl VariablePrimitiveOpt { + /// Wraps a present primitive value. + pub fn some(value: impl Into) -> Self { + Self { + value: value.into(), + is_present: true, + } + } + + /// Creates an absent primitive value. + pub fn none() -> Self { + Self::default() + } +} + +impl> From for VariablePrimitiveOpt { + fn from(value: T) -> Self { + Self::some(value) + } +} + +/// A nullable runtime argument or result slot. +/// +/// Primitive values use the [`VariablePrimitiveOpt`] arm. Strings use the +/// nullable `string` arm: a null data pointer represents `None`, while a +/// non-null data pointer and its byte length represent a borrowed `str`. +/// Both arms occupy two machine words on the supported 64-bit targets. +/// +/// string relies on the null ptr optimization to make that happen. +#[repr(C)] +#[derive(Clone, Copy)] +pub union VariableValue<'a> { + pub primitive: VariablePrimitiveOpt, + pub string: Option<&'a str>, +} + +const _: () = { + assert!(std::mem::size_of::() == 8); + assert!(std::mem::offset_of!(VariablePrimitiveOpt, value) == 0); + assert!(std::mem::offset_of!(VariablePrimitiveOpt, is_present) == 8); + assert!(std::mem::size_of::() == 16); + assert!(std::mem::size_of::>() == 16); + assert!(std::mem::size_of::() == 16); + assert!(std::mem::align_of::() == 8); +}; + +impl<'a> VariableValue<'a> { + /// Wraps a present runtime value. + #[inline(always)] + pub fn some(value: impl Into) -> Self { + value.into() + } + + /// Creates an absent runtime value for either arm. + #[inline(always)] + pub fn none() -> Self { + // SAFETY: All-zeroes is both an absent VariablePrimitiveOpt and the + // null niche used by Option<&str>. + unsafe { std::mem::zeroed() } + } + + /// Returns the boolean payload, or `None` when this value is absent. + /// + /// # Safety + /// + /// This value must contain a primitive boolean or be absent. + #[inline(always)] + pub unsafe fn as_bool(self) -> Option { + // SAFETY: Guaranteed by the caller. + let primitive = unsafe { self.primitive }; + if primitive.is_present { + // SAFETY: The caller guarantees that the active payload is `boolean`. + Some(unsafe { primitive.value.boolean }) + } else { + None + } + } + + /// Returns the `f64` payload, or `None` when this value is absent. + /// + /// # Safety + /// + /// This value must contain a primitive `f64` or be absent. + #[inline(always)] + pub unsafe fn as_f64(self) -> Option { + // SAFETY: Guaranteed by the caller. + let primitive = unsafe { self.primitive }; + if primitive.is_present { + // SAFETY: The caller guarantees that the active payload is `float`. + Some(unsafe { primitive.value.float }) + } else { + None + } + } + + /// Returns the `u64` payload, or `None` when this value is absent. + /// + /// # Safety + /// + /// This value must contain a primitive `u64` or be absent. + #[inline(always)] + pub unsafe fn as_u64(self) -> Option { + // SAFETY: Guaranteed by the caller. + let primitive = unsafe { self.primitive }; + if primitive.is_present { + // SAFETY: The caller guarantees that the active payload is `int_u64`. + Some(unsafe { primitive.value.int_u64 }) + } else { + None + } + } + + /// Returns the `i64` payload, or `None` when this value is absent. + /// + /// # Safety + /// + /// This value must contain a primitive `i64` or be absent. + #[inline(always)] + pub unsafe fn as_i64(self) -> Option { + // SAFETY: Guaranteed by the caller. + let primitive = unsafe { self.primitive }; + if primitive.is_present { + // SAFETY: The caller guarantees that the active payload is `int_i64`. + Some(unsafe { primitive.value.int_i64 }) + } else { + None + } + } + + /// Returns the borrowed string payload, or `None` when it is absent. + /// + /// # Safety + /// + /// This value must contain the `string` arm or be the all-zero absent + /// representation returned by [`VariableValue::none`]. + #[inline(always)] + pub unsafe fn as_str(self) -> Option<&'a str> { + // SAFETY: Guaranteed by the caller. + unsafe { self.string } + } +} + +impl Default for VariableValue<'_> { + fn default() -> Self { + Self::none() + } +} + +impl From for VariableValue<'_> { + fn from(value: bool) -> Self { + Self { + primitive: VariablePrimitiveOpt::some(value), + } + } +} + +impl From for VariableValue<'_> { + fn from(value: f64) -> Self { + Self { + primitive: VariablePrimitiveOpt::some(value), + } + } +} + +impl From for VariableValue<'_> { + fn from(value: u64) -> Self { + Self { + primitive: VariablePrimitiveOpt::some(value), + } + } +} + +impl From for VariableValue<'_> { + fn from(value: i64) -> Self { + Self { + primitive: VariablePrimitiveOpt::some(value), + } + } +} + +impl<'a> From<&'a str> for VariableValue<'a> { + fn from(value: &'a str) -> Self { + Self { + string: Some(value), + } + } +} + +impl<'a> From> for VariableValue<'a> { + fn from(value: Option<&'a str>) -> Self { + match value { + Some(value) => Self::from(value), + None => Self::none(), + } + } +} + +impl<'a> From for VariableValue<'a> { + fn from(value: VariablePrimitive) -> Self { + Self { + primitive: VariablePrimitiveOpt::some(value), + } + } +} + +impl<'a> From for VariableValue<'a> { + fn from(value: VariablePrimitiveOpt) -> Self { + Self { primitive: value } + } +} + +#[cfg(test)] +mod tests { + use crate::types::{VariablePrimitive, VariablePrimitiveOpt, VariableValue}; + + #[test] + fn test_runtime_value_layouts() { + assert_eq!(std::mem::size_of::(), 8); + assert_eq!(std::mem::offset_of!(VariablePrimitiveOpt, value), 0); + assert_eq!(std::mem::offset_of!(VariablePrimitiveOpt, is_present), 8); + assert_eq!(std::mem::size_of::(), 16); + assert_eq!(std::mem::size_of::>(), 16); + assert_eq!(std::mem::size_of::(), 16); + assert_eq!(std::mem::align_of::(), 8); + + let text = "hello"; + let words: [usize; 2] = unsafe { std::mem::transmute(VariableValue::some(text)) }; + assert_eq!(words, [text.as_ptr() as usize, text.len()]); + let none_words: [usize; 2] = unsafe { std::mem::transmute(VariableValue::none()) }; + assert_eq!(none_words, [0, 0]); + } + + #[test] + fn test_variable_value_accessors() { + assert_eq!(unsafe { VariableValue::some(true).as_bool() }, Some(true)); + assert_eq!(unsafe { VariableValue::some(1.5f64).as_f64() }, Some(1.5)); + assert_eq!(unsafe { VariableValue::some(7u64).as_u64() }, Some(7)); + assert_eq!(unsafe { VariableValue::some(-3i64).as_i64() }, Some(-3)); + assert_eq!( + unsafe { VariableValue::some(VariablePrimitive { int_u64: 11 }).as_u64() }, + Some(11) + ); + assert_eq!( + unsafe { VariableValue::some("hello").as_str() }, + Some("hello") + ); + + let none = VariableValue::none(); + assert_eq!(unsafe { none.as_bool() }, None); + assert_eq!(unsafe { none.as_f64() }, None); + assert_eq!(unsafe { none.as_u64() }, None); + assert_eq!(unsafe { none.as_i64() }, None); + assert_eq!(unsafe { none.as_str() }, None); + assert_eq!(unsafe { VariableValue::from(None::<&str>).as_str() }, None); + } + + #[test] + fn test_empty_string_is_distinct_from_none() { + let empty = VariableValue::some(""); + let none = VariableValue::none(); + + assert_eq!(unsafe { empty.as_str() }, Some("")); + assert_eq!(unsafe { none.as_str() }, None); + } +}