rust/tantivy
1
0
Fork 0
mirror of https://github.com/quickwit-oss/tantivy.git synced 2026-08-18 12:08:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

add serialization format

Browse Source 
Added command for disassembly
This commit is contained in:
Paul Masurel
2026-08-14 14:51:29 +02:00
parent 174fdfb4ab
commit e0b5cb0842
22 changed files with 2766 additions and 879 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jitexpr/Cargo.toml
+1
View File
@@ -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"
jitexpr/README.md
+9
View File
@@ -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
jitexpr/examples/basic.rs
+11 -104
View File
@@ -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<dyn Error>> {
// 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<dyn Error>> {
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::<Vec<_>>();
// 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<dyn Error>> {
// // 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::<Vec<_>>();
// 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() }
// }
jitexpr/src/FUNCTION.md
+3
View File
@@ -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.
jitexpr/src/ast/infer_types.rs
+180 -127
View File
@@ -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<VarType> 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<HashMap<&'a str, InferredTypeSet>, TypeError> {
pub fn infer_types(expr: &UntypedExpr) -> Result<HashMap<&str, InferredTypeSet>, 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<InferredTypeSet, TypeError> {
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<InferredTypeSet, TypeError> {
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<InferredTypeSet, TypeError> {
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}");
}
}
jitexpr/src/ast/literal.rs
+98
View File
@@ -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
);
}
}
jitexpr/src/ast/mod.rs
+5 -24
View File
@@ -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<TypedExpr>) -> TypedExprAst {
TypedExprAst::Call {
function: *self,
args,
}
}
pub fn call_untyped_expr(&self, args: Vec<UntypedExpr>) -> UntypedExpr {
UntypedExpr::Call {
function: *self,
args,
}
}
}
pub use crate::functions::Function;
jitexpr/src/ast/serialize.rs
+508
View File
@@ -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<UntypedExpr, DeserializeError> {
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<String>) -> 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<Self, Self::Err> {
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<Function, DeserializeError> {
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<UntypedExpr, DeserializeError> {
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<UntypedExpr, DeserializeError> {
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<UntypedExpr, DeserializeError> {
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<UntypedExpr, DeserializeError> {
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<UntypedExpr, DeserializeError> {
if let Some(value) = atom.strip_suffix("u64")
&& let Ok(value) = value.parse::<u64>()
{
return Ok(UntypedExpr::Literal(Literal::U64(value)));
}
if let Some(value) = atom.strip_suffix("i64")
&& let Ok(value) = value.parse::<i64>()
{
return Ok(UntypedExpr::Literal(Literal::I64(value)));
}
if let Some(value) = atom.strip_suffix("f64")
&& let Ok(value) = value.parse::<f64>()
{
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<String, DeserializeError> {
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<char, DeserializeError> {
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<char, DeserializeError> {
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<char> {
self.input[self.offset..].chars().next()
}
fn advance(&mut self) -> Option<char> {
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}"
);
}
}
}
jitexpr/src/ast/untyped_expr.rs
+4 -2
View File
@@ -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<str>),
jitexpr/src/bin/jitexpr-asm.rs
+151
View File
@@ -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<bool> {
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<String, ExpressionError> {
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::<HashMap<_, _>>();
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());
}
}
jitexpr/src/compile/apply_types.rs
-264
View File
@@ -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<TypedVariable>) {
let mut typed_expr = apply_types_aux(untyped_expr, variable_types);
let var_args: Vec<TypedVariable> = 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<TypedExpr> = 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<TypedExpr> = 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<TypedVariable> {
let mut name_to_vars: HashMap<Arc<str>, TypedVariable> = HashMap::new();
assign_variable_ids_aux(&mut expr.ast, &mut name_to_vars);
let mut input_vars: Vec<TypedVariable> = 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<Arc<str>, 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);
}
}
jitexpr/src/compile/compile_fn_builder.rs
+491
View File
@@ -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<TypedVariable>,
regexes: Vec<Regex>,
regex_match_results: Vec<UnsafeCell<StringRef>>,
string_literals: Vec<Arc<str>>,
}
struct LoweredFunction {
module: JITModule,
context: CodegenContext,
function_id: FuncId,
input_vars: Vec<TypedVariable>,
regexes: Box<[Regex]>,
regex_match_results: Box<[UnsafeCell<StringRef>]>,
string_literals: Box<[Arc<str>]>,
expression: Box<TypedExpr>,
}
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<str>) -> 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<TypedExpr, CompileError> {
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<TypedExpr, CompileError> {
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<CompiledFn, CompileError> {
self.lower_typed_expr(expression)?.into_compiled_fn()
}
pub(super) fn compile_typed_expr_to_assembly(
self,
expression: TypedExpr,
) -> Result<String, CompileError> {
self.lower_typed_expr(expression)?.into_assembly()
}
fn lower_typed_expr(self, expression: TypedExpr) -> Result<LoweredFunction, CompileError> {
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: &regex_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<VarType> {
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<CompiledFn, CompileError> {
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<String, CompileError> {
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<i64> {
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<u64> {
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<TypedVariable> {
let mut name_to_vars: HashMap<Arc<str>, TypedVariable> = HashMap::new();
assign_variable_ids_aux(&mut expr.ast, &mut name_to_vars);
let mut input_vars: Vec<TypedVariable> = 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<Arc<str>, 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);
}
}
jitexpr/src/compile/compiled_fn.rs
+48
View File
@@ -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<str>]>,
// 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<StringRef>]>,
pub(crate) input_vars: Vec<TypedVariable>,
// Generated code embeds addresses of StringRef descriptors in this AST.
pub(crate) _typed_expr: Box<TypedExpr>,
}
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()) };
}
}
jitexpr/src/compile/error.rs
+31
View File
@@ -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<cranelift_module::ModuleError>),
#[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<cranelift_module::ModuleError> for CompileError {
fn from(error: cranelift_module::ModuleError) -> Self {
CompileError::Module(Box::new(error))
}
}
jitexpr/src/compile/mod.rs
+91 -312
View File
@@ -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<TypedVariable>,
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<cranelift_module::ModuleError>),
#[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<cranelift_module::ModuleError> 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<CompiledFunction, CompileError> {
let (typed_expr, input_vars) = apply_types(untyped_expr, var_types);
compile_typed_expr(typed_expr, input_vars)
) -> Result<CompiledFn, CompileError> {
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<TypedVariable>,
) -> Result<CompiledFunction, CompileError> {
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<String, CompileError> {
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<StringRef>],
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<Value, CompileError> {
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<Value, CompileError> {
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<Value, CompileError> {
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<StringRef> {
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<StringRef>) {
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());
}
}
jitexpr/src/compile/typed_expr.rs
+92 -41
View File
@@ -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<str>,
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<Literal>) -> TypedExpr {
#[cfg(test)]
pub(crate) fn literal(val: impl Into<Literal>) -> 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<TypedExpr>,
},
Call {
function: Function,
args: Vec<TypedExpr>,
},
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<Literal>) -> 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<Literal> for TypedExprAst {
fn from(literal: Literal) -> Self {
TypedExprAst::Literal(literal)
pub(crate) fn from_call(fn_call: impl Into<FnCallEnum>) -> 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:?}")
}
}
}
jitexpr/src/functions/add.rs
+454
View File
@@ -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<InferredTypeSet, TypeError> {
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<TypedExpr, CompileError> {
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<TypedExpr> = args
.iter()
.map(|arg| context.apply_types(arg, InferredTypeSet::singleton(return_type)))
.collect::<Result<_, _>>()?;
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<cranelift::codegen::ir::Value, CompileError> {
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<AddFnCall> 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);
}
}
jitexpr/src/functions/mod.rs
+142
View File
@@ -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<TypedExpr, CompileError> {
match self {
Function::Add => <AddFnCall as FnCall>::call_with_types(args, target_type_set, context),
Function::RegexpExtract => {
<RegexpExtractFnCall as FnCall>::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<InferredTypeSet, TypeError> {
match self {
Function::Add => <AddFnCall as FnCall>::infer_types(args, target_type, inferred_types),
Function::RegexpExtract => {
<RegexpExtractFnCall as FnCall>::infer_types(args, target_type, inferred_types)
}
}
}
pub fn call_untyped_expr(self, args: Vec<UntypedExpr>) -> 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<cranelift::codegen::ir::Value, CompileError> {
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<FnCallEnum> {
/// 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<InferredTypeSet, TypeError>
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<TypedExpr, CompileError>
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<cranelift::codegen::ir::Value, CompileError>;
}
jitexpr/src/functions/native_function.rs
+32
View File
@@ -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<NativeFunctions, CompileError> {
Ok(NativeFunctions {
regexp_extract: regexp_extract::declare_native_function(module, function, pointer_type)?,
})
}
jitexpr/src/functions/regexp_extract.rs
+363
View File
@@ -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<TypedExpr>,
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<InferredTypeSet, TypeError> {
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<TypedExpr, CompileError> {
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<Value, CompileError> {
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<FuncRef, CompileError> {
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<RegexpExtractFnCall> 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 == "("
));
}
}
jitexpr/src/lib.rs
+31
View File
@@ -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);
}
}
jitexpr/src/types.rs
+21 -5
View File
@@ -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::<VariableValue>(), 8);
}
}