From e0b5cb0842d62a7bb63a51280403184a4b9ba558 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Fri, 14 Aug 2026 14:51:29 +0200 Subject: [PATCH] add serialization format Added command for disassembly --- jitexpr/Cargo.toml | 1 + jitexpr/README.md | 9 + jitexpr/examples/basic.rs | 115 +---- jitexpr/src/FUNCTION.md | 3 + jitexpr/src/ast/infer_types.rs | 307 +++++++------ jitexpr/src/ast/literal.rs | 98 +++++ jitexpr/src/ast/mod.rs | 29 +- jitexpr/src/ast/serialize.rs | 508 ++++++++++++++++++++++ jitexpr/src/ast/untyped_expr.rs | 6 +- jitexpr/src/bin/jitexpr-asm.rs | 151 +++++++ jitexpr/src/compile/apply_types.rs | 264 ----------- jitexpr/src/compile/compile_fn_builder.rs | 491 +++++++++++++++++++++ jitexpr/src/compile/compiled_fn.rs | 48 ++ jitexpr/src/compile/error.rs | 31 ++ jitexpr/src/compile/mod.rs | 403 ++++------------- jitexpr/src/compile/typed_expr.rs | 133 ++++-- jitexpr/src/functions/add.rs | 454 +++++++++++++++++++ jitexpr/src/functions/mod.rs | 142 ++++++ jitexpr/src/functions/native_function.rs | 32 ++ jitexpr/src/functions/regexp_extract.rs | 363 ++++++++++++++++ jitexpr/src/lib.rs | 31 ++ jitexpr/src/types.rs | 26 +- 22 files changed, 2766 insertions(+), 879 deletions(-) create mode 100644 jitexpr/README.md create mode 100644 jitexpr/src/FUNCTION.md create mode 100644 jitexpr/src/ast/serialize.rs create mode 100644 jitexpr/src/bin/jitexpr-asm.rs delete mode 100644 jitexpr/src/compile/apply_types.rs create mode 100644 jitexpr/src/compile/compile_fn_builder.rs create mode 100644 jitexpr/src/compile/compiled_fn.rs create mode 100644 jitexpr/src/compile/error.rs create mode 100644 jitexpr/src/functions/add.rs create mode 100644 jitexpr/src/functions/mod.rs create mode 100644 jitexpr/src/functions/native_function.rs create mode 100644 jitexpr/src/functions/regexp_extract.rs diff --git a/jitexpr/Cargo.toml b/jitexpr/Cargo.toml index aeef2fc5c..86c169c2e 100644 --- a/jitexpr/Cargo.toml +++ b/jitexpr/Cargo.toml @@ -7,4 +7,5 @@ edition = "2024" cranelift = "0.134.3" cranelift-jit = "0.134.3" cranelift-module = "0.134.3" +regex = "1" thiserror = "2.0.1" diff --git a/jitexpr/README.md b/jitexpr/README.md new file mode 100644 index 000000000..2baa59e1d --- /dev/null +++ b/jitexpr/README.md @@ -0,0 +1,9 @@ +This is an expression compiler relying on Cranelift. + + UntypedExpr + ↓ injecting variable types, and type checking + TypedExpr + ↓ lowering + Cranelift IR + ↓ Cranelift code generation + Machine code diff --git a/jitexpr/examples/basic.rs b/jitexpr/examples/basic.rs index 6ecbb7865..2a0ef623d 100644 --- a/jitexpr/examples/basic.rs +++ b/jitexpr/examples/basic.rs @@ -2,13 +2,12 @@ use std::collections::HashMap; use std::error::Error; use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, infer_types}; -use jitexpr::compile::{CompiledFunction, compile}; -use jitexpr::types::VarType; +use jitexpr::compile::{CompiledFn, compile}; +use jitexpr::types::{VarType, VariableValue}; fn main() -> Result<(), Box> { // A simple expression that goes: - // my_column + 1 - + // my_col + 1 let untyped_expr = Function::Add.call_untyped_expr(vec![ UntypedExpr::variable("my_col"), UntypedExpr::literal(1.0f64), @@ -29,107 +28,15 @@ fn main() -> Result<(), Box> { let variable_types: HashMap<&str, VarType> = std::iter::once(("my_col", VarType::F64)).collect(); - let compiled_fn: CompiledFunction = compile(&untyped_expr, &variable_types).unwrap(); + let compiled_fn: CompiledFn = compile(&untyped_expr, &variable_types)?; - // let function = compile(&expression, selected_types)?; + // We use a union to pass typed variables to the function. + // It is up to us to correctly populate it. Not doing so is UB. + let input: Box<[VariableValue]> = vec![VariableValue { float: 1.2f64 }].into_boxed_slice(); + // The initialization does not really matter. + let mut output: VariableValue = VariableValue { int_u64: 0u64 }; + unsafe { compiled_fn.call(&input[..], &mut output) }; + assert_eq!(unsafe { output.float }, 1.2f64 + 1.0f64); - // assert_eq!(evaluate(&function, 100, 4), 25.0); - // assert_eq!(evaluate(&function, 100, 0), 0.0); Ok(()) } - -// fn evaluate(function: &CompiledFunction, request_size: u64, elapsed: u64) -> f64 { -// // NamedInput defines the positional order expected by the generated code. -// let args = function -// .inputs() -// .iter() -// .map(|input| { -// assert_eq!(input.var_type(), VarType::U64); -// match input.name() { -// "request_size" => Variable::from_u64(request_size), -// "elapsed" => Variable::from_u64(elapsed), -// name => panic!("unexpected input `{name}`"), -// } -// }) -// .collect::>(); - -// let mut result = Variable::null(); -// // SAFETY: `args` follows `function.inputs()` and every payload matches its -// // reported VarType. `result` is a writable Variable slot. -// unsafe { function.call(&args, &mut result) }; -// assert!(!result.is_null()); -// // SAFETY: The compiled signature reports F64 and the result is non-null. -// unsafe { result.as_f64() } -// } - -// use jitexpr::ast::Expr; - -// fn main() -> Result<(), Box> { -// // COALESCE(request_size / elapsed, 0.0). Division produces F64 and returns -// // null when its divisor is zero, so COALESCE supplies the fallback. -// let expression = Expr::call( -// Function::Coalesce, -// [ -// Expr::call( -// Function::Divide, -// [Expr::variable("request_size"), Expr::variable("elapsed")], -// ), -// Literal::F64(0.0).into(), -// ], -// ); - -// let argument_names = expression.list_argument_names(); -// println!("referenced fields: {argument_names:?}"); - -// // An integrating crate would obtain these entries by looking up the -// // referenced fields in Tantivy's columnar schema. -// let available_types = HashMap::from([ -// ( -// "request_size".to_string(), -// AvailableVarTypes { -// numerical: Some(NumericalType::U64), -// boolean: false, -// string: false, -// }, -// ), -// ( -// "elapsed".to_string(), -// AvailableVarTypes { -// numerical: Some(NumericalType::U64), -// boolean: false, -// string: false, -// }, -// ), -// ]); - -// let selected_types = infer_types(&expression, &available_types)?; -// let function = compile(&expression, selected_types)?; - -// assert_eq!(evaluate(&function, 100, 4), 25.0); -// assert_eq!(evaluate(&function, 100, 0), 0.0); -// Ok(()) -// } - -// fn evaluate(function: &CompiledFunction, request_size: u64, elapsed: u64) -> f64 { -// // NamedInput defines the positional order expected by the generated code. -// let args = function -// .inputs() -// .iter() -// .map(|input| { -// assert_eq!(input.var_type(), VarType::U64); -// match input.name() { -// "request_size" => Variable::from_u64(request_size), -// "elapsed" => Variable::from_u64(elapsed), -// name => panic!("unexpected input `{name}`"), -// } -// }) -// .collect::>(); - -// let mut result = Variable::null(); -// // SAFETY: `args` follows `function.inputs()` and every payload matches its -// // reported VarType. `result` is a writable Variable slot. -// unsafe { function.call(&args, &mut result) }; -// assert!(!result.is_null()); -// // SAFETY: The compiled signature reports F64 and the result is non-null. -// unsafe { result.as_f64() } -// } diff --git a/jitexpr/src/FUNCTION.md b/jitexpr/src/FUNCTION.md new file mode 100644 index 000000000..c3c0b2ee1 --- /dev/null +++ b/jitexpr/src/FUNCTION.md @@ -0,0 +1,3 @@ + +Regexp: +only 3 args supported. group is not optional. if not passed, set 0 during the conversion from proto to rust object. diff --git a/jitexpr/src/ast/infer_types.rs b/jitexpr/src/ast/infer_types.rs index c75e202c6..1ee42044b 100644 --- a/jitexpr/src/ast/infer_types.rs +++ b/jitexpr/src/ast/infer_types.rs @@ -2,50 +2,115 @@ use std::collections::HashMap; use std::collections::hash_map::Entry; use crate::ast::{Function, Literal, UntypedExpr}; +use crate::types::VarType; #[derive(Default, Copy, Clone, Debug, Eq, PartialEq)] pub struct InferredTypeSet { - string: bool, - numerical: bool, - boolean: bool, + pub string: bool, + pub i64: bool, + pub u64: bool, + pub f64: bool, + pub boolean: bool, } impl InferredTypeSet { pub const NONE: InferredTypeSet = InferredTypeSet { string: false, - numerical: false, + i64: false, + u64: false, + f64: false, boolean: false, }; pub const ALL: InferredTypeSet = InferredTypeSet { string: true, - numerical: true, + i64: true, + u64: true, + f64: true, boolean: true, }; pub const NUMERICAL: InferredTypeSet = InferredTypeSet { - numerical: true, + i64: true, + u64: true, + f64: true, boolean: false, string: false, }; - pub const STRING: InferredTypeSet = InferredTypeSet { - numerical: false, - boolean: false, - string: true, + pub const I64: InferredTypeSet = InferredTypeSet { + i64: true, + ..Self::NONE }; - fn is_none(self) -> bool { + 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 } - fn intersect(self, target_inferred_type: InferredTypeSet) -> InferredTypeSet { + 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, - numerical: self.numerical && target_inferred_type.numerical, + 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 { @@ -54,8 +119,14 @@ impl std::fmt::Display for InferredTypeSet { if self.string { types.push("string"); } - if self.numerical { - types.push("numerical"); + if self.i64 { + types.push("i64"); + } + if self.u64 { + types.push("u64"); + } + if self.f64 { + types.push("f64"); } if self.boolean { types.push("boolean"); @@ -66,10 +137,17 @@ impl std::fmt::Display for InferredTypeSet { #[derive(Debug, thiserror::Error)] pub enum TypeError { - #[error("function `{function:?}` returns a number, expected `{expected}`")] + #[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 { @@ -79,23 +157,20 @@ pub enum TypeError { } /// Infer the accepted types for the different variables present in the formula. -pub fn infer_types<'a>( - expr: &'a UntypedExpr, -) -> Result, TypeError> { +pub fn infer_types(expr: &UntypedExpr) -> Result, TypeError> { let mut inferred_type_res = HashMap::default(); infer_types_aux(expr, InferredTypeSet::ALL, &mut inferred_type_res)?; Ok(inferred_type_res) } -fn infer_types_aux<'a>( +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(literal)); + let literal_type: InferredTypeSet = target_inferred_type.intersect(literal.types()); if literal_type.is_none() { return Err(TypeError::InvalidLiteralType { literal: literal.clone(), @@ -104,122 +179,61 @@ fn infer_types_aux<'a>( } Ok(literal_type) } - UntypedExpr::Variable(variable_name) => match inferred_types_res.entry(&*variable_name) { - Entry::Occupied(mut occupied_entry) => { - let inferred_types = occupied_entry.get().intersect(target_inferred_type); - occupied_entry.insert(inferred_types); - Ok(inferred_types) + 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) + } } - Entry::Vacant(vacant_entry) => { - vacant_entry.insert_entry(target_inferred_type); - Ok(target_inferred_type) - } - }, - UntypedExpr::Call { function, args } => infer_types_function_aux( - *function, - &args[..], - target_inferred_type, - inferred_types_res, - ), - } -} - -fn infer_types_function_aux<'a>( - function: Function, - args: &'a [UntypedExpr], - target_inferred_type: InferredTypeSet, - inferred_types_res: &mut HashMap<&'a str, InferredTypeSet>, -) -> Result { - match function { - Function::Add => { - // This is valid for all functions taking a bunch of number and returning a number. - if !target_inferred_type.numerical { - return Err(TypeError::WrongFunctionReturnType { - function, - expected: target_inferred_type, - }); - } - for arg in args { - infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types_res)?; - } - Ok(InferredTypeSet::NUMERICAL) + } + UntypedExpr::Call { function, args } => { + function.infer_types(args, target_inferred_type, inferred_types_res) } } } -fn literal_types<'a>(literal: &'a Literal) -> InferredTypeSet { - match literal { - Literal::None => InferredTypeSet::ALL, - Literal::Bool(_) => InferredTypeSet { - boolean: true, - ..Default::default() - }, - Literal::I64(_) | Literal::U64(_) | Literal::F64(_) => InferredTypeSet::NUMERICAL, - Literal::String(_) => InferredTypeSet::STRING, +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 std::assert_matches; - use super::*; - use crate::ast::{Function, Literal, UntypedExpr}; - - #[test] - fn test_infer_types_add_string_and_float_returns_error() { - // add(1.0, "hello") should fail because a string cannot be numerical. - let expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::literal(1.0), - UntypedExpr::literal("hello"), - ]); - let err = infer_types(&expr).unwrap_err(); - assert_matches!( - err, - TypeError::InvalidLiteralType { - literal: Literal::String(_), - expected: InferredTypeSet { - string: false, - numerical: true, - boolean: false, - }, - } - ); - } - - #[test] - fn test_infer_types_add_literal_and_variable() { - // add(1, a) should infer that `a` is numerical. - let expr = Function::Add - .call_untyped_expr(vec![UntypedExpr::literal(1i64), UntypedExpr::variable("a")]); - let inferred_types = infer_types(&expr).unwrap(); - let a_types = inferred_types.get("a").unwrap(); - assert!(a_types.numerical); - assert!(!a_types.string); - assert!(!a_types.boolean); - } - - #[test] - fn test_infer_types_add_heterogenous_literals() { - let expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::literal(1.2f64), - UntypedExpr::literal(2u64), - ]); - assert!(infer_types(&expr).is_ok()); - } - - #[test] - fn test_infer_types_add_two_variables() { - // add(a, b) should infer that both `a` and `b` are numerical. - let expr = Function::Add - .call_untyped_expr(vec![UntypedExpr::variable("a"), UntypedExpr::variable("b")]); - let inferred_types = infer_types(&expr).unwrap(); - - let a_types = inferred_types.get("a").unwrap(); - assert_eq!(a_types, &InferredTypeSet::NUMERICAL); - let b_types = inferred_types.get("b").unwrap(); - assert_eq!(b_types, &InferredTypeSet::NUMERICAL); - } #[test] fn test_infer_types_bare_variable_accepts_all() { @@ -229,4 +243,43 @@ mod tests { 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 index b370efae2..1eb0d56a2 100644 --- a/jitexpr/src/ast/literal.rs +++ b/jitexpr/src/ast/literal.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use crate::ast::InferredTypeSet; use crate::types::VarType; /// A literal supported by the first expression-language milestone. @@ -18,6 +19,36 @@ impl Literal { matches!(self, Literal::None) } + pub fn types(&self) -> InferredTypeSet { + match self { + Literal::None => InferredTypeSet::ALL, + Literal::Bool(_) => InferredTypeSet::BOOLEAN, + Literal::I64(value) => InferredTypeSet { + i64: true, + u64: *value >= 0, + f64: (*value as f64) as i128 == *value as i128, + ..InferredTypeSet::NONE + }, + Literal::U64(value) => InferredTypeSet { + i64: *value <= i64::MAX as u64, + u64: true, + f64: (*value as f64) as i128 == *value as i128, + ..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, @@ -65,3 +96,70 @@ impl From<&str> for Literal { 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_require_exact_float_representation() { + let integer_types = InferredTypeSet { + i64: true, + u64: true, + ..InferredTypeSet::NONE + }; + + assert_eq!(Literal::I64((1 << 53) + 1).types(), integer_types); + assert_eq!(Literal::I64(i64::MAX).types(), integer_types); + assert_eq!(Literal::U64(u64::MAX).types(), InferredTypeSet::U64); + 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 index 25ed8431f..5b2607d08 100644 --- a/jitexpr/src/ast/mod.rs +++ b/jitexpr/src/ast/mod.rs @@ -1,31 +1,12 @@ mod infer_types; mod literal; +mod serialize; mod untyped_expr; -pub use infer_types::{InferredTypeSet, infer_types}; +pub use infer_types::{InferredTypeSet, TypeError, infer_types}; +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; -use crate::compile::{TypedExpr, TypedExprAst}; - -/// A function supported by the first expression-language milestone. -#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -pub enum Function { - Add, -} - -impl Function { - pub fn call_typed_expr(&self, args: Vec) -> TypedExprAst { - TypedExprAst::Call { - function: *self, - args, - } - } - - pub fn call_untyped_expr(&self, args: Vec) -> UntypedExpr { - UntypedExpr::Call { - function: *self, - args, - } - } -} +pub use crate::functions::Function; diff --git a/jitexpr/src/ast/serialize.rs b/jitexpr/src/ast/serialize.rs new file mode 100644 index 000000000..01fdfc30a --- /dev/null +++ b/jitexpr/src/ast/serialize.rs @@ -0,0 +1,508 @@ +//! 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::Add => "ADD", + Function::RegexpExtract => "REGEXP_EXTRACT", + } +} + +fn parse_function(name: &str, offset: usize) -> Result { + match name { + "ADD" => Ok(Function::Add), + "REGEXP_EXTRACT" => Ok(Function::RegexpExtract), + _ 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_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 index 113169742..0127639a9 100644 --- a/jitexpr/src/ast/untyped_expr.rs +++ b/jitexpr/src/ast/untyped_expr.rs @@ -2,8 +2,10 @@ use std::sync::Arc; use crate::ast::{Function, Literal}; -/// An expression independent from its protobuf representation. -#[derive(Clone, Debug, PartialEq)] +/// An expression AST. +/// +/// The expression at this point is untyped and not necessarily valid. +#[derive(Clone, PartialEq)] pub enum UntypedExpr { Literal(Literal), Variable(Arc), diff --git a/jitexpr/src/bin/jitexpr-asm.rs b/jitexpr/src/bin/jitexpr-asm.rs new file mode 100644 index 000000000..034539e7f --- /dev/null +++ b/jitexpr/src/bin/jitexpr-asm.rs @@ -0,0 +1,151 @@ +//! Prints 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. + +use std::collections::HashMap; +use std::io::{self, BufRead, Write}; +use std::process::ExitCode; + +use jitexpr::ast::{DeserializeError, InferredTypeSet, TypeError, deserialize, infer_types}; +use jitexpr::compile::{CompileError, compile_to_assembly}; +use jitexpr::types::VarType; + +#[derive(Debug, thiserror::Error)] +enum ExpressionError { + #[error(transparent)] + Deserialize(#[from] DeserializeError), + #[error(transparent)] + Type(#[from] TypeError), + #[error(transparent)] + Compile(#[from] CompileError), +} + +fn main() -> ExitCode { + let stdin = io::stdin(); + let stdout = io::stdout(); + let stderr = io::stderr(); + match process_lines(stdin.lock(), stdout.lock(), stderr.lock()) { + 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, +) -> io::Result { + let mut all_succeeded = true; + let mut wrote_assembly = false; + + for (line_index, line) in input.lines().enumerate() { + let line = line?; + if line.trim().is_empty() { + continue; + } + + match compile_line(&line) { + Ok(assembly) => { + if wrote_assembly { + writeln!(output)?; + } + output.write_all(assembly.as_bytes())?; + if !assembly.ends_with('\n') { + writeln!(output)?; + } + wrote_assembly = true; + } + Err(error) => { + writeln!(errors, "line {}: {error}", line_index + 1)?; + all_succeeded = false; + } + } + } + + Ok(all_succeeded) +} + +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::>(); + Ok(compile_to_assembly(&expression, &variable_types)?) +} + +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 assembly = compile_line("(ADD 1i64 my_col)").unwrap(); + + assert!(assembly.contains("block0:")); + assert!(!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).unwrap(); + + assert!(!all_succeeded); + let output = String::from_utf8(output).unwrap(); + assert_eq!(output.matches("block0:").count(), 2); + 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).unwrap(); + + assert!(all_succeeded); + assert!(output.is_empty()); + assert!(errors.is_empty()); + } +} diff --git a/jitexpr/src/compile/apply_types.rs b/jitexpr/src/compile/apply_types.rs deleted file mode 100644 index 352d4b492..000000000 --- a/jitexpr/src/compile/apply_types.rs +++ /dev/null @@ -1,264 +0,0 @@ -use std::collections::HashMap; -use std::sync::Arc; - -use super::typed_expr::{TypedExpr, TypedExprAst, TypedVariable}; -use crate::ast::{Function, Literal, UntypedExpr}; -use crate::types::VarType; - -/// If a variable is missing from variable_types, it will be treated as if its value is None. -pub fn apply_types( - untyped_expr: &UntypedExpr, - variable_types: &HashMap<&str, VarType>, -) -> (TypedExpr, Vec) { - let mut typed_expr = apply_types_aux(untyped_expr, variable_types); - let var_args: Vec = assign_variable_ids(&mut typed_expr); - (typed_expr, var_args) -} - -fn apply_types_aux( - untyped_expr: &UntypedExpr, - variable_types: &HashMap<&str, VarType>, -) -> TypedExpr { - match untyped_expr { - UntypedExpr::Literal(literal) => TypedExpr { - return_type: literal.r#type(), - ast: TypedExprAst::Literal(literal.clone()), - }, - UntypedExpr::Variable(variable_name) => { - let variable_type: VarType = variable_types - .get(variable_name.as_ref()) - .copied() - // a missing column is treated as if it was there with a constant - // None value. - .unwrap_or(VarType::None); - if variable_type == VarType::None { - TypedExpr { - return_type: VarType::None, - ast: TypedExprAst::Literal(Literal::None), - } - } else { - TypedExpr { - return_type: variable_type, - ast: TypedExprAst::variable(variable_name, variable_type), - } - } - } - UntypedExpr::Call { function, args } => match function { - Function::Add => apply_types_add_aux(args, variable_types), - }, - } -} - -fn apply_types_add_aux(args: &[UntypedExpr], variable_types: &HashMap<&str, VarType>) -> TypedExpr { - let typed_args: Vec = args - .iter() - .map(|arg| apply_types_aux(arg, variable_types)) - .collect(); - - let mut all_u64 = true; - let mut all_i64 = true; - for typed_arg in &typed_args { - match typed_arg.return_type { - VarType::U64 => all_i64 = false, - VarType::I64 => all_u64 = false, - VarType::F64 => { - all_u64 = false; - all_i64 = false; - } - _ => return TypedExpr::none(), - } - } - let return_type = if all_u64 { - VarType::U64 - } else if all_i64 { - VarType::I64 - } else { - VarType::F64 - }; - let typed_args: Vec = typed_args - .into_iter() - .map(|typed_arg| typed_arg.coerce(return_type)) - .collect(); - TypedExpr { - return_type, - ast: Function::Add.call_typed_expr(typed_args), - } -} - -/// Walks the AST and assigns each distinct variable an auto-incremented id -/// (its offset in the input array). Repeated occurrences of the same variable -/// share the same id. -/// -/// Returns the list of input variables in id order. -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::Call { args, .. } => { - for arg in args { - assign_variable_ids_aux(&mut arg.ast, name_to_vars); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_apply_types_sum_simple() { - let untyped_expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::variable("present"), - UntypedExpr::literal(1u64), - ]); - let variable_types = HashMap::from([("present", VarType::U64)]); - - let (typed_expr, _) = apply_types(&untyped_expr, &variable_types); - - assert_eq!( - typed_expr, - Function::Add - .call_typed_expr(vec![ - TypedExprAst::variable("present", VarType::U64).with_type(VarType::U64), - TypedExpr::literal(1u64), - ]) - .with_type(VarType::U64) - ); - } - - #[test] - fn test_apply_types_sum_coercion() { - let untyped_expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::variable("present"), - UntypedExpr::literal(1.2f64), - ]); - let variable_types = HashMap::from([("present", VarType::U64)]); - - let (typed_expr, _) = apply_types(&untyped_expr, &variable_types); - - assert_eq!( - typed_expr, - Function::Add - .call_typed_expr(vec![ - TypedExprAst::variable("present", VarType::U64) - .with_type(VarType::U64) - .coerce(VarType::F64), - TypedExpr::literal(1.2f64), - ]) - .with_type(VarType::F64) - ); - } - - #[test] - fn test_apply_types_sum_variable_missing() { - let untyped_expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::variable("present"), - Function::Add.call_untyped_expr(vec![ - UntypedExpr::literal(1u64), - UntypedExpr::variable("missing"), - ]), - ]); - let variable_types = HashMap::from([("present", VarType::U64)]); - - let (typed_expr, _) = apply_types(&untyped_expr, &variable_types); - - assert_eq!(typed_expr, TypedExpr::none()); - } - - #[test] - fn test_apply_types_to_literal() { - let untyped_expr = UntypedExpr::literal("hello"); - assert_eq!( - apply_types(&untyped_expr, &HashMap::new()).0, - TypedExprAst::literal("hello").with_type(VarType::Str) - ); - } - - #[test] - fn test_assign_variable_ids_two_variables_different_types() { - // add(x, y) with x: U64 and y: F64. Add coerces U64 to F64, so we - // get: Add(Coerce(x as F64), y). ids are assigned in DFS order. - 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 (_typed_expr, var_args) = apply_types(&untyped_expr, &variable_types); - - 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() { - // Manually build a TypedExpr where the variable `x` appears twice with - // two different types (U64 and F64). This should never happen when the - // tree is built via apply_types, so we panic to surface the bug. - let mut typed_expr = Function::Add - .call_typed_expr(vec![ - TypedExprAst::variable("x", VarType::U64).with_type(VarType::U64), - TypedExprAst::variable("x", VarType::F64).with_type(VarType::F64), - ]) - .with_type(VarType::F64); - - assign_variable_ids(&mut typed_expr); - } - - #[test] - fn test_assign_variable_ids_dedups_repeated_variable() { - // add(x, add(y, x)) — `x` appears twice and must be assigned the same id - // (single slot in the input array). Expected DFS traversal: - // x (new, id=0), y (new, id=1), x (already seen, id=0). - 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 (_typed_expr, var_args) = apply_types(&untyped_expr, &variable_types); - - 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); - } -} diff --git a/jitexpr/src/compile/compile_fn_builder.rs b/jitexpr/src/compile/compile_fn_builder.rs new file mode 100644 index 000000000..c96a75a92 --- /dev/null +++ b/jitexpr/src/compile/compile_fn_builder.rs @@ -0,0 +1,491 @@ +use std::cell::UnsafeCell; +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 regex::Regex; + +use super::compiled_fn::JitEntry; +use super::{ + CompileError, CompiledFn, LoweringContext, TypedExpr, TypedExprAst, TypedLiteral, + TypedVariable, lower_expr, +}; +use crate::ast::{InferredTypeSet, Literal, UntypedExpr}; +use crate::functions::{declare_native_functions, register_jit_symbols}; +use crate::types::{StringRef, VarType}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct RegexRef(usize); + +impl RegexRef { + pub(crate) fn index(self) -> usize { + self.0 + } +} + +pub(crate) struct CompileFnBuilder<'types, 'names> { + variable_types: &'types HashMap<&'names str, VarType>, + input_vars: Vec, + regexes: Vec, + regex_match_results: Vec>, + string_literals: Vec>, +} + +struct LoweredFunction { + module: JITModule, + context: CodegenContext, + function_id: FuncId, + input_vars: Vec, + regexes: Box<[Regex]>, + regex_match_results: Box<[UnsafeCell]>, + string_literals: Box<[Arc]>, + expression: Box, +} + +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(), + regexes: Vec::new(), + regex_match_results: Vec::new(), + string_literals: Vec::new(), + } + } + + pub(crate) fn variable_types(&self) -> &HashMap<&'names str, VarType> { + self.variable_types + } + + pub(crate) fn register_regex(&mut self, regex: Regex) -> RegexRef { + let regex_ref = RegexRef(self.regexes.len()); + self.regexes.push(regex); + self.regex_match_results + .push(UnsafeCell::new(StringRef::new(""))); + regex_ref + } + + pub(crate) fn register_string_literal(&mut self, value: Arc) -> StringRef { + let string_ref = StringRef::new(&value); + self.string_literals.push(value); + string_ref + } + + /// 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(self.register_string_literal(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, + regexes, + regex_match_results, + string_literals, + .. + } = self; + let regexes = regexes.into_boxed_slice(); + let regex_match_results = regex_match_results.into_boxed_slice(); + let string_literals = string_literals.into_boxed_slice(); + let expression = Box::new(expression); + + let mut jit_builder = + JITBuilder::with_flags(&[("opt_level", "speed")], default_libcall_names())?; + 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: the first two arguments are the + // input and output slots, and the third points to CompiledFn::regexes. + let mut signature = module.make_signature(); + signature.params.push(AbiParam::new(pointer_type)); + signature.params.push(AbiParam::new(pointer_type)); + signature.params.push(AbiParam::new(pointer_type)); + 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 result_ptr = builder.block_params(entry_block)[1]; + let regexes_ptr = builder.block_params(entry_block)[2]; + let mut lowering_context = LoweringContext { + args_ptr, + regexes_ptr, + pointer_type, + regex_match_results: ®ex_match_results, + native_functions: &native_functions, + }; + let value = lower_expr(&expression, &mut lowering_context, &mut builder)?; + builder + .ins() + .store(MemFlagsData::trusted(), value, result_ptr, 0); + builder.ins().return_(&[]); + builder.finalize(target_config); + } + + Ok(LoweredFunction { + module, + context, + function_id, + input_vars, + regexes, + regex_match_results, + string_literals, + 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, + regexes, + regex_match_results, + string_literals, + 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, + _string_literals: string_literals, + regexes, + _regex_match_results: regex_match_results, + 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(string_ref)) = typed_expr.ast else { + panic!("expected a typed string literal"); + }; + assert_eq!(unsafe { string_ref.as_str() }, "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); + } +} diff --git a/jitexpr/src/compile/compiled_fn.rs b/jitexpr/src/compile/compiled_fn.rs new file mode 100644 index 000000000..49878f4c4 --- /dev/null +++ b/jitexpr/src/compile/compiled_fn.rs @@ -0,0 +1,48 @@ +use std::cell::UnsafeCell; +use std::sync::Arc; + +use cranelift_jit::JITModule; +use regex::Regex; + +use super::{TypedExpr, TypedVariable}; +use crate::types::{StringRef, VariableValue}; + +pub(crate) type JitEntry = + unsafe extern "C" fn(*const VariableValue, *mut VariableValue, *const Regex); + +/// 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, + // Typed string-literal descriptors borrow their bytes from these values. + pub(crate) _string_literals: Box<[Arc]>, + // Generated code selects a compiled regex by its index in this array. + pub(crate) regexes: Box<[Regex]>, + // Each regex call site owns stable storage for the StringRef descriptor it + // returns. UnsafeCell makes the mutation performed by the Rust helper + // explicit and prevents CompiledFn from being shared between threads. + pub(crate) _regex_match_results: Box<[UnsafeCell]>, + pub(crate) input_vars: Vec, + // Generated code embeds addresses of StringRef descriptors in this AST. + pub(crate) _typed_expr: Box, +} + +impl CompiledFn { + /// Evaluate the compiled expression. + /// + /// # Safety + /// + /// `args` must follow `input_vars` exactly: every slot must contain the + /// union member corresponding to that variable's type. `result` must be a + /// valid writable slot and any referenced strings must remain alive for + /// the duration of this call. A string result descriptor remains valid + /// until the next call to this `CompiledFn`. + pub unsafe fn call(&self, args: &[VariableValue], result: &mut VariableValue) { + debug_assert_eq!(args.len(), self.input_vars.len()); + // SAFETY: Guaranteed by the caller. + unsafe { (self.entry)(args.as_ptr(), result, self.regexes.as_ptr()) }; + } +} diff --git a/jitexpr/src/compile/error.rs b/jitexpr/src/compile/error.rs new file mode 100644 index 000000000..ac099a8d7 --- /dev/null +++ b/jitexpr/src/compile/error.rs @@ -0,0 +1,31 @@ +use crate::ast::{Function, 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("input variable {variable_id} has an address offset that is too large")] + InputOffsetOverflow { variable_id: usize }, + #[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, + }, +} + +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 index 7bf99d7db..3e68def8f 100644 --- a/jitexpr/src/compile/mod.rs +++ b/jitexpr/src/compile/mod.rs @@ -1,161 +1,84 @@ -mod apply_types; +mod compile_fn_builder; +mod compiled_fn; +mod error; mod typed_expr; +use std::cell::UnsafeCell; use std::collections::HashMap; -use std::mem::{self, size_of}; +use std::mem::size_of; -pub use apply_types::apply_types; -use cranelift::codegen::ir::{MemFlagsData, UserFuncName}; +pub(crate) use compile_fn_builder::{CompileFnBuilder, RegexRef}; +pub use compiled_fn::CompiledFn; +use cranelift::codegen::ir::MemFlagsData; use cranelift::prelude::*; -use cranelift_jit::{JITBuilder, JITModule}; -use cranelift_module::{Module, default_libcall_names}; -pub use typed_expr::{TypedExpr, TypedExprAst, TypedVariable}; +pub use error::CompileError; +pub(crate) use typed_expr::{TypedExpr, TypedExprAst, TypedLiteral, TypedVariable}; -use crate::ast::{Function, Literal, UntypedExpr}; +use crate::ast::UntypedExpr; +use crate::functions::NativeFunctions; use crate::types::{StringRef, VarType, VariableValue}; -/// An expression compiled to native machine code. -/// -/// This object owns the JIT module containing its executable memory. -pub struct CompiledFunction { - pub(crate) entry: JitEntry, - pub(crate) _module: JITModule, - // String literals are addressed directly by the generated code. Keeping - // their descriptors here gives those addresses the lifetime of the JIT. - pub(crate) _literal_strings: Box<[StringRef]>, - pub input_vars: Vec, - pub typed_expr: TypedExpr, -} - -impl CompiledFunction { - /// Evaluate the compiled expression. - /// - /// # Safety - /// - /// `args` must follow `input_vars` exactly: every slot must contain the - /// union member corresponding to that variable's type. `result` must be a - /// valid writable slot and any referenced strings must remain alive for - /// the duration of this call. - pub unsafe fn call(&self, args: &[VariableValue], result: &mut VariableValue) { - debug_assert_eq!(args.len(), self.input_vars.len()); - // SAFETY: Guaranteed by the caller. - unsafe { (self.entry)(args.as_ptr(), result) }; - } -} - -type JitEntry = unsafe extern "C" fn(*const VariableValue, *mut VariableValue); - -#[derive(Debug, thiserror::Error)] -pub enum CompileError { - #[error("JIT compilation failed: {0}")] - Module(#[source] Box), - #[error("input variable {variable_id} has an address offset that is too large")] - InputOffsetOverflow { variable_id: usize }, - #[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, - }, -} - -impl From for CompileError { - fn from(error: cranelift_module::ModuleError) -> Self { - CompileError::Module(Box::new(error)) - } -} - pub fn compile( untyped_expr: &UntypedExpr, var_types: &HashMap<&str, VarType>, -) -> Result { - let (typed_expr, input_vars) = apply_types(untyped_expr, var_types); - compile_typed_expr(typed_expr, input_vars) +) -> Result { + let mut builder = CompileFnBuilder::new(var_types); + let typed_expr = builder.build_typed_expr(untyped_expr)?; + builder.compile_typed_expr(typed_expr) } -fn compile_typed_expr( - expression: TypedExpr, - input_vars: Vec, -) -> Result { - let jit_builder = JITBuilder::new(default_libcall_names())?; - let mut module = JITModule::new(jit_builder); - let target_config = module.target_config(); - let pointer_type = target_config.pointer_type(); +/// 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) +} - // The native entry point mirrors JitEntry: both arguments are pointers and - // the expression result is written into the second one. - let mut signature = module.make_signature(); - signature.params.push(AbiParam::new(pointer_type)); - signature.params.push(AbiParam::new(pointer_type)); - let function_id = module.declare_anonymous_function(&signature)?; +pub(crate) struct LoweringContext<'a> { + args_ptr: Value, + regexes_ptr: Value, + pointer_type: Type, + regex_match_results: &'a [UnsafeCell], + native_functions: &'a NativeFunctions, +} - let mut context = module.make_context(); - context.func.signature = signature; - context.func.name = UserFuncName::user(0, function_id.as_u32()); - - let mut function_builder_context = FunctionBuilderContext::new(); - let literal_strings = collect_literal_strings(&expression).into_boxed_slice(); - let mut next_literal_string = 0; - { - 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 result_ptr = builder.block_params(entry_block)[1]; - let value = lower_expr( - &expression, - args_ptr, - pointer_type, - &literal_strings, - &mut next_literal_string, - &mut builder, - )?; - debug_assert_eq!(next_literal_string, literal_strings.len()); - builder - .ins() - .store(MemFlagsData::trusted(), value, result_ptr, 0); - builder.ins().return_(&[]); - builder.finalize(target_config); +impl LoweringContext<'_> { + pub(crate) fn lower_expr( + &mut self, + expression: &TypedExpr, + builder: &mut FunctionBuilder<'_>, + ) -> Result { + lower_expr(expression, self, builder) } - module.define_function(function_id, &mut context)?; - module.finalize_definitions()?; + pub(crate) fn pointer_type(&self) -> Type { + self.pointer_type + } - 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 - // `CompiledFunction`, so its executable allocation outlives `entry`. - let entry = unsafe { mem::transmute::<*const u8, JitEntry>(code) }; + pub(crate) fn regexes_ptr(&self) -> Value { + self.regexes_ptr + } - Ok(CompiledFunction { - entry, - _module: module, - _literal_strings: literal_strings, - input_vars, - typed_expr: expression, - }) + pub(crate) fn native_functions(&self) -> &NativeFunctions { + self.native_functions + } + + pub(crate) fn regex_match_result(&self, regex_ref: RegexRef) -> *mut StringRef { + self.regex_match_results[regex_ref.index()].get() + } } +/// Produce cranelift IR from the function fn lower_expr( expression: &TypedExpr, - args_ptr: Value, - pointer_type: Type, - literal_strings: &[StringRef], - next_literal_string: &mut usize, + context: &mut LoweringContext<'_>, builder: &mut FunctionBuilder<'_>, ) -> Result { match &expression.ast { - TypedExprAst::Literal(literal) => Ok(lower_literal( - literal, - pointer_type, - literal_strings, - next_literal_string, - builder, - )), + TypedExprAst::Literal(literal) => Ok(lower_literal(literal, context, builder)), TypedExprAst::Variable(variable) => { let byte_offset = variable .variable_id @@ -165,56 +88,37 @@ fn lower_expr( variable_id: variable.variable_id, })?; Ok(builder.ins().load( - cranelift_type(variable.r#type, pointer_type), + cranelift_type(variable.r#type, context.pointer_type), MemFlagsData::trusted(), - args_ptr, + context.args_ptr, byte_offset, )) } TypedExprAst::Coerce { target_type, expr } => { let source_type = expr.return_type; - let value = lower_expr( - expr, - args_ptr, - pointer_type, - literal_strings, - next_literal_string, - builder, - )?; + let value = lower_expr(expr, context, builder)?; lower_coercion(value, source_type, *target_type, builder) } - TypedExprAst::Call { function, args } => match function { - Function::Add => lower_add( - args, - expression.return_type, - args_ptr, - pointer_type, - literal_strings, - next_literal_string, - builder, - ), - }, + TypedExprAst::FnCall(fn_call) => fn_call.lower(expression.return_type, context, builder), } } fn lower_literal( - literal: &Literal, - pointer_type: Type, - literal_strings: &[StringRef], - next_literal_string: &mut usize, + literal: &TypedLiteral, + context: &mut LoweringContext<'_>, builder: &mut FunctionBuilder<'_>, ) -> Value { match literal { - Literal::None => builder.ins().iconst(types::I64, 0), - Literal::Bool(value) => builder.ins().iconst(types::I8, i64::from(*value)), - Literal::U64(value) => builder.ins().iconst(types::I64, *value as i64), - Literal::I64(value) => builder.ins().iconst(types::I64, *value), - Literal::F64(value) => builder.ins().f64const(Ieee64::with_bits(value.to_bits())), - Literal::String(_) => { - let string_ref = &literal_strings[*next_literal_string]; - *next_literal_string += 1; + TypedLiteral::None => builder.ins().iconst(types::I64, 0), + TypedLiteral::Bool(value) => builder.ins().iconst(types::I8, i64::from(*value)), + TypedLiteral::U64(value) => builder.ins().iconst(types::I64, *value as i64), + TypedLiteral::I64(value) => builder.ins().iconst(types::I64, *value), + TypedLiteral::F64(value) => builder.ins().f64const(Ieee64::with_bits(value.to_bits())), + TypedLiteral::String(string_ref) => { let string_ref_ptr = (string_ref as *const StringRef) as usize; - builder.ins().iconst(pointer_type, string_ref_ptr as i64) + builder + .ins() + .iconst(context.pointer_type, string_ref_ptr as i64) } } } @@ -242,65 +146,6 @@ fn lower_coercion( Ok(coerced) } -fn lower_add( - args: &[TypedExpr], - return_type: VarType, - args_ptr: Value, - pointer_type: Type, - literal_strings: &[StringRef], - next_literal_string: &mut usize, - 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(Ieee64::with_bits(0)), - _ => { - return Err(CompileError::UnsupportedFunctionType { - function: Function::Add, - return_type, - }); - } - }; - - for arg in args { - let value = lower_expr( - arg, - args_ptr, - pointer_type, - literal_strings, - next_literal_string, - builder, - )?; - sum = match return_type { - VarType::U64 | VarType::I64 => builder.ins().iadd(sum, value), - VarType::F64 => builder.ins().fadd(sum, value), - _ => unreachable!("the return type was checked above"), - }; - } - Ok(sum) -} - -fn collect_literal_strings(expression: &TypedExpr) -> Vec { - let mut literal_strings = Vec::new(); - collect_literal_strings_aux(expression, &mut literal_strings); - literal_strings -} - -fn collect_literal_strings_aux(expression: &TypedExpr, literal_strings: &mut Vec) { - match &expression.ast { - TypedExprAst::Literal(Literal::String(value)) => { - literal_strings.push(StringRef::new(value)); - } - TypedExprAst::Literal(_) | TypedExprAst::Variable(_) => {} - TypedExprAst::Coerce { expr, .. } => collect_literal_strings_aux(expr, literal_strings), - TypedExprAst::Call { args, .. } => { - for arg in args { - collect_literal_strings_aux(arg, literal_strings); - } - } - } -} - fn cranelift_type(var_type: VarType, pointer_type: Type) -> Type { match var_type { VarType::Bool => types::I8, @@ -318,87 +163,6 @@ mod tests { use crate::ast::{Function, UntypedExpr}; use crate::types::VarType; - #[test] - fn test_compile_signed_add() { - let untyped_expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::literal(-4i64), - UntypedExpr::variable("myfield"), - ]); - let variable_types = HashMap::from([("myfield", VarType::I64)]); - let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); - let input = [VariableValue { int_i64: -8 }]; - let mut output = VariableValue { int_i64: 0 }; - - unsafe { compiled_fn.call(&input, &mut output) }; - - assert_eq!(unsafe { output.int_i64 }, -12); - } - - #[test] - fn test_compile_add_coerces_integers_to_float() { - let untyped_expr = 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 compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); - let input = [VariableValue { int_u64: 10 }]; - let mut output = VariableValue { float: 0.0 }; - - unsafe { compiled_fn.call(&input, &mut output) }; - - assert_eq!(unsafe { output.float }, 8.5); - } - - #[test] - fn test_compile_add_loads_multiple_variable_slots() { - 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 compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); - let input = [VariableValue { int_u64: 10 }, VariableValue { float: 0.5 }]; - let mut output = VariableValue { float: 0.0 }; - - unsafe { compiled_fn.call(&input, &mut output) }; - - assert_eq!(unsafe { output.float }, 10.5); - } - - #[test] - fn test_compile_u64_to_float_coercion_is_unsigned() { - let untyped_expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::variable("x"), - UntypedExpr::literal(0.0f64), - ]); - let variable_types = HashMap::from([("x", VarType::U64)]); - let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); - let input = [VariableValue { int_u64: u64::MAX }]; - let mut output = VariableValue { float: 0.0 }; - - unsafe { compiled_fn.call(&input, &mut output) }; - - assert_eq!(unsafe { output.float }, u64::MAX as f64); - } - - #[test] - fn test_compile_reuses_repeated_variable_slot() { - let untyped_expr = Function::Add.call_untyped_expr(vec![ - UntypedExpr::variable("x"), - UntypedExpr::variable("x"), - UntypedExpr::literal(1u64), - ]); - let variable_types = HashMap::from([("x", VarType::U64)]); - let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); - let input = [VariableValue { int_u64: 4 }]; - let mut output = VariableValue { int_u64: 0 }; - - unsafe { compiled_fn.call(&input, &mut output) }; - - assert_eq!(compiled_fn.input_vars.len(), 1); - assert_eq!(unsafe { output.int_u64 }, 9); - } - #[test] fn test_compile_bool_variable() { let untyped_expr = UntypedExpr::variable("flag"); @@ -428,13 +192,28 @@ mod tests { } #[test] - fn test_compile_empty_add_uses_zero_identity() { - let untyped_expr = Function::Add.call_untyped_expr(Vec::new()); - let compiled_fn = compile(&untyped_expr, &HashMap::new()).unwrap(); - let mut output = VariableValue { int_u64: 1 }; + fn test_compile_to_assembly() { + let untyped_expr = UntypedExpr::variable("value"); + let variable_types = HashMap::from([("value", VarType::F64)]); - unsafe { compiled_fn.call(&[], &mut output) }; + let assembly = compile_to_assembly(&untyped_expr, &variable_types).unwrap(); - assert_eq!(unsafe { output.int_u64 }, 0); + 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()); } } diff --git a/jitexpr/src/compile/typed_expr.rs b/jitexpr/src/compile/typed_expr.rs index 67fcebf5d..4045a4edc 100644 --- a/jitexpr/src/compile/typed_expr.rs +++ b/jitexpr/src/compile/typed_expr.rs @@ -1,90 +1,148 @@ use std::sync::Arc; -use crate::ast::{Function, Literal}; -use crate::types::VarType; +#[cfg(test)] +use crate::ast::Literal; +use crate::functions::FnCallEnum; +use crate::types::{StringRef, VarType}; #[derive(Clone, PartialEq)] -pub struct TypedVariable { +pub(crate) struct TypedVariable { pub(super) variable_name: Arc, pub(super) r#type: VarType, pub(super) variable_id: usize, //< offset in the input array. } #[derive(Clone, PartialEq)] -pub struct TypedExpr { - pub return_type: VarType, - pub ast: TypedExprAst, +pub(crate) struct TypedExpr { + pub(crate) return_type: VarType, + pub(crate) ast: TypedExprAst, } impl TypedExpr { - pub fn coerce(self, target_type: VarType) -> 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: TypedExprAst::Coerce { - target_type, - expr: Box::new(self), - }, + ast, } } } - pub fn none() -> TypedExpr { + pub(crate) fn none() -> TypedExpr { TypedExpr { return_type: VarType::None, - ast: TypedExprAst::Literal(Literal::None), + ast: TypedExprAst::Literal(TypedLiteral::None), } } - pub fn literal(val: impl Into) -> TypedExpr { + #[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(StringRef), +} + +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 enum TypedExprAst { - Literal(Literal), +pub(crate) enum TypedExprAst { + Literal(TypedLiteral), Variable(TypedVariable), Coerce { target_type: VarType, expr: Box, }, - Call { - function: Function, - args: Vec, - }, + FnCall(FnCallEnum), } impl TypedExprAst { - pub fn with_type(self, return_type: VarType) -> TypedExpr { + #[cfg(test)] + pub(crate) fn with_type(self, return_type: VarType) -> TypedExpr { TypedExpr { return_type, ast: self, } } - pub fn literal(val: impl Into) -> TypedExprAst { - TypedExprAst::Literal(val.into()) - } - - pub fn variable(variable_name: impl ToString, r#type: VarType) -> TypedExprAst { + 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, }) } -} -// ---------- boilerplate --------- - -impl From for TypedExprAst { - fn from(literal: Literal) -> Self { - TypedExprAst::Literal(literal) + pub(crate) fn from_call(fn_call: impl Into) -> TypedExprAst { + TypedExprAst::FnCall(fn_call.into()) } } @@ -102,15 +160,8 @@ impl std::fmt::Debug for TypedExprAst { TypedExprAst::Coerce { target_type, expr } => { write!(f, "coerce({:?} as {:?})", expr, target_type) } - TypedExprAst::Call { function, args } => { - write!(f, "{:?}(", function)?; - for (i, arg) in args.iter().enumerate() { - if i > 0 { - write!(f, ", ")?; - } - write!(f, "{:?}", arg)?; - } - write!(f, ")") + TypedExprAst::FnCall(fn_call) => { + write!(f, "{fn_call:?}") } } } diff --git a/jitexpr/src/functions/add.rs b/jitexpr/src/functions/add.rs new file mode 100644 index 000000000..88d7ce7a4 --- /dev/null +++ b/jitexpr/src/functions/add.rs @@ -0,0 +1,454 @@ +// 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 actually automatically coerce +// 1f64 to 1i64 (because we have detected that the conversion was lossless), and the operation will +// run over integer. + +use std::collections::HashMap; + +use cranelift::prelude::{FunctionBuilder, InstBuilder, types}; + +use crate::ast::{Function, InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{CompileError, CompileFnBuilder, 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 { + 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 { + 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 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, + }); + } + }; + + for arg in &self.args { + let value = context.lower_expr(arg, builder)?; + sum = match return_type { + VarType::U64 | VarType::I64 => builder.ins().iadd(sum, value), + VarType::F64 => builder.ins().fadd(sum, value), + _ => unreachable!("the return type was checked above"), + }; + } + Ok(sum) + } +} + +fn with_float_fallback(inferred_types: InferredTypeSet) -> InferredTypeSet { + if inferred_types.is_none() { + InferredTypeSet::F64 + } else { + inferred_types + } +} + +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 + } +} + +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 typed_expr = + crate::typed_expr_from_str("(ADD 9223372036854775808u64)", &variable_types); + + assert_eq!(typed_expr.return_type, VarType::U64); + } + + #[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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_i64: -8 }]; + let mut output = VariableValue { int_i64: 0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { output.int_i64 }, -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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: 41 }]; + let mut output = VariableValue { int_u64: 0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { output.int_u64 }, 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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: 39 }]; + let mut output = VariableValue { int_u64: 0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { output.int_u64 }, 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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: 10 }]; + let mut output = VariableValue { float: 0.0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { output.float }, 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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: 10 }, VariableValue { float: 0.5 }]; + let mut output = VariableValue { float: 0.0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { output.float }, 10.5); + } + + #[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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: u64::MAX }]; + let mut output = VariableValue { float: 0.0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { output.float }, 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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: 4 }]; + let mut output = VariableValue { int_u64: 0 }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(compiled.input_vars.len(), 1); + assert_eq!(unsafe { output.int_u64 }, 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 compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { int_u64: 4 }]; + let mut output = VariableValue::default(); + unsafe { compiled.call(&input, &mut output) }; + assert_eq!(compiled.input_vars.len(), 1); + assert_eq!(unsafe { output.float }, 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 compiled = compile(&expression, &HashMap::new()).unwrap(); + let mut output = VariableValue { int_i64: 1 }; + + unsafe { compiled.call(&[], &mut output) }; + + assert_eq!(unsafe { output.int_i64 }, 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 compiled = compile(&expression, &HashMap::new()).unwrap(); + let mut output = VariableValue { int_i64: 1 }; + unsafe { compiled.call(&[], &mut output) }; + assert_eq!(unsafe { output.float }, 2.2f64); + } +} diff --git a/jitexpr/src/functions/mod.rs b/jitexpr/src/functions/mod.rs new file mode 100644 index 000000000..b26039e8c --- /dev/null +++ b/jitexpr/src/functions/mod.rs @@ -0,0 +1,142 @@ +mod add; +mod native_function; +mod regexp_extract; + +use std::collections::HashMap; + +use cranelift::frontend::FunctionBuilder; + +pub(crate) use self::add::AddFnCall; +pub(crate) use self::native_function::{ + NativeFunctions, declare_native_functions, register_jit_symbols, +}; +pub(crate) use self::regexp_extract::RegexpExtractFnCall; +use crate::ast::{InferredTypeSet, TypeError, UntypedExpr}; +use crate::compile::{CompileError, CompileFnBuilder, 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 { + /// Adds zero or more numerical expressions. + Add, + /// Extracts a capture group from a string using a constant regular expression. + RegexpExtract, +} + +impl Function { + pub(crate) fn call_with_types( + self, + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + match self { + Function::Add => ::call_with_types(args, target_type_set, context), + Function::RegexpExtract => { + ::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::Add => ::infer_types(args, target_type, inferred_types), + Function::RegexpExtract => { + ::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 { + Add(AddFnCall), + RegexpExtract(RegexpExtractFnCall), +} + +impl FnCallEnum { + pub(crate) fn args_mut(&mut self) -> &mut [TypedExpr] { + match self { + FnCallEnum::Add(call) => call.args_mut(), + FnCallEnum::RegexpExtract(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::Add(call) => call.emit_cranelift_ir(return_type, context, builder), + FnCallEnum::RegexpExtract(call) => { + call.emit_cranelift_ir(return_type, context, builder) + } + } + } +} + +/// 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 { + /// 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; + + /// Builds the typed call after concrete variable types have been supplied. + /// + /// `target_type_set` communicates the result types preferred by the parent call. The + /// implementation selects a concrete result type, applies compatible target types to its + /// arguments through `context`, and registers any compilation resources owned by the call. + 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. + fn args_mut(&mut self) -> &mut [TypedExpr]; + + /// 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; +} diff --git a/jitexpr/src/functions/native_function.rs b/jitexpr/src/functions/native_function.rs new file mode 100644 index 000000000..42afc6829 --- /dev/null +++ b/jitexpr/src/functions/native_function.rs @@ -0,0 +1,32 @@ +use cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type}; +use cranelift_jit::{JITBuilder, JITModule}; + +use super::regexp_extract; +use crate::compile::CompileError; + +/// References to native functions imported into the current Cranelift function. +pub(crate) struct NativeFunctions { + regexp_extract: FuncRef, +} + +impl NativeFunctions { + pub(crate) fn regexp_extract(&self) -> FuncRef { + self.regexp_extract + } +} + +/// Registers the process symbols that native calls may reference from generated code. +pub(crate) fn register_jit_symbols(jit_builder: &mut JITBuilder) { + regexp_extract::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 { + Ok(NativeFunctions { + regexp_extract: regexp_extract::declare_native_function(module, function, pointer_type)?, + }) +} diff --git a/jitexpr/src/functions/regexp_extract.rs b/jitexpr/src/functions/regexp_extract.rs new file mode 100644 index 000000000..20ee3ab9c --- /dev/null +++ b/jitexpr/src/functions/regexp_extract.rs @@ -0,0 +1,363 @@ +// RegexpExtract extracts a regular-expression match from a string. +// +// It takes three arguments: the input string, a regular-expression pattern literal, and a u64 +// capture index literal. 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 cranelift::codegen::ir::{FuncRef, Function as CraneliftFunction, Type, Value, 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, LoweringContext, RegexRef, TypedExpr, TypedExprAst, +}; +use crate::functions::{FnCall, FnCallEnum}; +use crate::types::{StringRef, VarType}; + +const SYMBOL: &str = "jitexpr_regexp_extract"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct RegexpExtractFnCall { + regex_ref: RegexRef, + haystack: Box, + capture_index: u64, +} + +impl FnCall for RegexpExtractFnCall { + 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 args.len() != 3 { + 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)?; + crate::ast::infer_types_aux(&args[2], InferredTypeSet::NUMERICAL, inferred_types)?; + Ok(InferredTypeSet::STRING) + } + + fn call_with_types( + args: &[UntypedExpr], + target_type_set: InferredTypeSet, + context: &mut CompileFnBuilder<'_, '_>, + ) -> Result { + assert_eq!(args.len(), 3, "Expected 3 args for regexp_extract"); + + let haystack = context.apply_types(&args[0], target_type_set)?; + assert_eq!(haystack.return_type, VarType::Str); + + let UntypedExpr::Literal(Literal::String(pattern)) = &args[1] else { + panic!("regexp_extract pattern must be a string literal"); + }; + let regex = Regex::new(pattern).map_err(|source| CompileError::InvalidRegex { + pattern: pattern.to_string(), + source, + })?; + let regex_ref = context.register_regex(regex); + + let UntypedExpr::Literal(Literal::U64(capture_index)) = &args[2] else { + panic!("regexp_extract capture index must be a u64 literal"); + }; + + Ok(TypedExpr { + return_type: VarType::Str, + ast: TypedExprAst::from_call(RegexpExtractFnCall { + regex_ref, + haystack: Box::new(haystack), + capture_index: *capture_index, + }), + }) + } + + fn args_mut(&mut self) -> &mut [TypedExpr] { + std::slice::from_mut(&mut self.haystack) + } + + /// 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.lower_expr(&self.haystack, builder)?; + let regex_index = builder + .ins() + .iconst(context.pointer_type(), self.regex_ref.index() as i64); + let capture_index = builder.ins().iconst(types::I64, self.capture_index as i64); + let match_result = context.regex_match_result(self.regex_ref); + let match_result = builder + .ins() + .iconst(context.pointer_type(), match_result as usize as i64); + let call = builder.ins().call( + context.native_functions().regexp_extract(), + &[ + context.regexes_ptr(), + regex_index, + haystack, + capture_index, + match_result, + ], + ); + Ok(builder.inst_results(call)[0]) + } +} + +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), 3)); + signature.params.push(AbiParam::new(types::I64)); + signature.params.push(AbiParam::new(pointer_type)); + signature.returns.push(AbiParam::new(pointer_type)); + let function_id = module.declare_function(SYMBOL, Linkage::Import, &signature)?; + Ok(module.declare_func_in_func(function_id, function)) +} + +/// Runtime implementation called by generated code for `RegexpExtract`. +/// +/// The JIT only forwards opaque pointers. All knowledge of `Regex` and +/// `StringRef`, including construction of the borrowed result descriptor, +/// stays in this Rust function. +unsafe extern "C" fn regexp_extract( + regexes: *const Regex, + regex_index: usize, + haystack: *const StringRef, + capture_index: u64, + match_result: *mut StringRef, +) -> *mut StringRef { + if haystack.is_null() { + return std::ptr::null_mut(); + } + let Ok(capture_index) = usize::try_from(capture_index) else { + return std::ptr::null_mut(); + }; + // SAFETY: Generated code passes CompiledFn::regexes and an index + // assigned while constructing that same array. + let regex = unsafe { &*regexes.add(regex_index) }; + // SAFETY: The contract of CompiledFn::call requires live StringRef + // input descriptors whose backing bytes contain valid UTF-8. + let haystack = unsafe { (*haystack).as_str() }; + let Some(regex_match) = regex + .captures(haystack) + .and_then(|captures| captures.get(capture_index)) + else { + return std::ptr::null_mut(); + }; + + // SAFETY: Generated code passes a dedicated UnsafeCell-backed result slot. + // The descriptor borrows bytes from the input rather than copying them. + unsafe { match_result.write(StringRef::new(regex_match.as_str())) }; + match_result +} + +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]+)" 0u64)"#).unwrap(); + + let inferred_types = infer_types(&expression).unwrap(); + + assert_eq!( + inferred_types.get("message"), + Some(&InferredTypeSet::STRING) + ); + } + + #[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 compiled = compile(&expression, &variable_types).unwrap(); + let haystack = "prefix user-123 suffix"; + let mut haystack_ref = StringRef::new(haystack); + let input = [VariableValue { + string: &mut haystack_ref, + }]; + let mut output = VariableValue { + string: std::ptr::null_mut(), + }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(compiled.regexes.len(), 1); + assert_eq!(compiled.regexes[0].as_str(), r"([a-z]+)-(\d+)"); + let output = unsafe { output.string }; + assert!(!output.is_null()); + let extracted = unsafe { (*output).as_str() }; + 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 compiled = compile(&expression, &variable_types).unwrap(); + let mut haystack = StringRef::new("user-123"); + let input = [VariableValue { + string: &mut haystack, + }]; + let mut output = VariableValue { + string: std::ptr::null_mut(), + }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { (*output.string).as_str() }, "123"); + } + + #[test] + fn test_compile_returns_null_without_capture() { + let expression = + ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)-(\\d+)" 0u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let compiled = compile(&expression, &variable_types).unwrap(); + let mut haystack = StringRef::new("no digits here"); + let input = [VariableValue { + string: &mut haystack, + }]; + let mut output = VariableValue { + string: std::ptr::dangling_mut(), + }; + + unsafe { compiled.call(&input, &mut output) }; + + assert!(unsafe { output.string }.is_null()); + } + + #[test] + fn test_compile_propagates_null_haystack() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT message "([a-z]+)" 0u64)"#).unwrap(); + let variable_types = HashMap::from([("message", VarType::Str)]); + let compiled = compile(&expression, &variable_types).unwrap(); + let input = [VariableValue { + string: std::ptr::null_mut(), + }]; + let mut output = VariableValue { + string: std::ptr::dangling_mut(), + }; + + unsafe { compiled.call(&input, &mut output) }; + + assert!(unsafe { output.string }.is_null()); + } + + #[test] + fn test_compile_distinguishes_empty_capture_from_null() { + let expression = ast::deserialize(r#"(REGEXP_EXTRACT "b" "(a*)b" 1u64)"#).unwrap(); + let compiled = compile(&expression, &HashMap::new()).unwrap(); + let mut output = VariableValue { + string: std::ptr::null_mut(), + }; + + unsafe { compiled.call(&[], &mut output) }; + + let output = unsafe { output.string }; + assert!(!output.is_null()); + assert_eq!(unsafe { (*output).as_str() }, ""); + } + + #[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 compiled = compile(&expression, &variable_types).unwrap(); + let mut haystack = StringRef::new("prefix user-123 suffix"); + let input = [VariableValue { + string: &mut haystack, + }]; + let mut output = VariableValue { + string: std::ptr::null_mut(), + }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(unsafe { (*output.string).as_str() }, "user-123"); + } + + #[test] + fn test_compile_nested_calls_use_distinct_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 compiled = compile(&expression, &variable_types).unwrap(); + let mut haystack = StringRef::new("id=user-123!"); + let input = [VariableValue { + string: &mut haystack, + }]; + let mut output = VariableValue { + string: std::ptr::null_mut(), + }; + + unsafe { compiled.call(&input, &mut output) }; + + assert_eq!(compiled.regexes.len(), 2); + assert_eq!(unsafe { (*output.string).as_str() }, "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/lib.rs b/jitexpr/src/lib.rs index c37d5cd83..b6aa1517c 100644 --- a/jitexpr/src/lib.rs +++ b/jitexpr/src/lib.rs @@ -1,3 +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 index ceb19a614..18cb913de 100644 --- a/jitexpr/src/types.rs +++ b/jitexpr/src/types.rs @@ -14,15 +14,15 @@ pub enum VarType { /// A borrowed UTF-8 string descriptor passed opaquely through generated code. /// /// The pointer can either refer to: -/// - an input str, if it is representing an arg -/// or if it is a return value that is a slice of an input str (e.g. a regex group). +/// - an input str, if it is representing an arg or if it is a return value that is a slice of an +/// input str (e.g. a regex group). /// - a literal from the original expression -/// - the arena passed to the function if the function "constructs" a new string -/// (e.g. when calling uppercase). +/// - the arena passed to the function if the function "constructs" a new string (e.g. when calling +/// uppercase). /// /// Either way, its lifetime / ownership is controlled by the called of the function. #[repr(C)] -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct StringRef { data: *const u8, len: usize, @@ -78,3 +78,19 @@ pub union VariableValue { pub int_i64: i64, pub string: *mut StringRef, //< this has to be mut for results. } + +impl Default for VariableValue { + fn default() -> VariableValue { + VariableValue { int_u64: 0u64 } + } +} + +#[cfg(test)] +mod tests { + use crate::types::VariableValue; + + #[test] + fn test_variable_value_size() { + assert_eq!(std::mem::size_of::(), 8); + } +}