From 7778e3cf4120f648530cbfa25945af9b8dd6a036 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Fri, 14 Aug 2026 14:21:15 +0200 Subject: [PATCH] first stab at compile --- jitexpr/examples/basic.rs | 14 +- jitexpr/src/ast/apply_types.rs | 260 +++++++++++++++++++++++++++++++++ jitexpr/src/ast/mod.rs | 123 +--------------- jitexpr/src/ast/typed_expr.rs | 61 +++++++- jitexpr/src/compile/mod.rs | 68 +++++++++ jitexpr/src/lib.rs | 1 + jitexpr/src/types.rs | 18 +-- 7 files changed, 390 insertions(+), 155 deletions(-) create mode 100644 jitexpr/src/ast/apply_types.rs create mode 100644 jitexpr/src/compile/mod.rs diff --git a/jitexpr/examples/basic.rs b/jitexpr/examples/basic.rs index ff70512c3..8396c2a0c 100644 --- a/jitexpr/examples/basic.rs +++ b/jitexpr/examples/basic.rs @@ -1,9 +1,8 @@ use std::collections::HashMap; use std::error::Error; -use jitexpr::ast::{ - Function, InferredTypeSet, TypedExprAst, UntypedExpr, apply_types, infer_types, -}; +use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, apply_types, infer_types}; +use jitexpr::compile::{CompiledFunction, compile}; use jitexpr::types::VarType; fn main() -> Result<(), Box> { @@ -23,14 +22,7 @@ fn main() -> Result<(), Box> { let variable_types: HashMap<&str, VarType> = std::iter::once(("my_col", VarType::F64)).collect(); - let typed_expr: TypedExprAst = apply_types(&untyped_expr, variable_types); - assert_eq!( - typed_expr, - Function::Add.call_typed_expr(vec![ - TypedExprAst::variable("my_col", VarType::F64), - TypedExprAst::literal(1.0f64), - ]) - ); + let compiled_fn: CompiledFunction = compile(&untyped_expr, &variable_types).unwrap(); // let function = compile(&expression, selected_types)?; diff --git a/jitexpr/src/ast/apply_types.rs b/jitexpr/src/ast/apply_types.rs new file mode 100644 index 000000000..77e6b0745 --- /dev/null +++ b/jitexpr/src/ast/apply_types.rs @@ -0,0 +1,260 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use crate::ast::typed_expr::TypedVariable; +use crate::ast::{Function, Literal, TypedExpr, TypedExprAst, UntypedExpr}; +use crate::types::VarType; + +/// If a variable is missing from variable_types, it will be treated as if its value is None. +pub fn apply_types( + untyped_expr: &UntypedExpr, + variable_types: &HashMap<&str, VarType>, +) -> (TypedExpr, Vec) { + let mut typed_expr = apply_types_aux(untyped_expr, variable_types); + let var_args: Vec = assign_variable_ids(&mut typed_expr); + (typed_expr, var_args) +} + +fn apply_types_aux( + untyped_expr: &UntypedExpr, + variable_types: &HashMap<&str, VarType>, +) -> TypedExpr { + match untyped_expr { + UntypedExpr::Literal(literal) => TypedExpr { + return_type: literal.r#type(), + ast: TypedExprAst::Literal(literal.clone()), + }, + UntypedExpr::Variable(variable_name) => { + if let Some(variable_type) = variable_types.get(variable_name.as_ref()).copied() { + TypedExpr { + return_type: variable_type, + ast: TypedExprAst::variable(variable_name, variable_type), + } + } else { + // a missing column is treated as if it was there with a constant + // None value. + TypedExpr { + return_type: VarType::None, + ast: TypedExprAst::Literal(Literal::None), + } + } + } + UntypedExpr::Call { function, args } => match function { + Function::Add => apply_types_add_aux(args, variable_types), + }, + } +} + +fn apply_types_add_aux(args: &[UntypedExpr], variable_types: &HashMap<&str, VarType>) -> TypedExpr { + let typed_args: Vec = args + .iter() + .map(|arg| apply_types_aux(arg, variable_types)) + .collect(); + + let mut all_u64 = true; + let mut all_i64 = true; + for typed_arg in &typed_args { + match typed_arg.return_type { + VarType::U64 => all_i64 = false, + VarType::I64 => all_u64 = false, + VarType::F64 => { + all_u64 = false; + all_i64 = false; + } + _ => return TypedExpr::none(), + } + } + let return_type = if all_u64 { + VarType::U64 + } else if all_i64 { + VarType::I64 + } else { + VarType::F64 + }; + let typed_args: Vec = typed_args + .into_iter() + .map(|typed_arg| typed_arg.coerce(return_type)) + .collect(); + TypedExpr { + return_type, + ast: Function::Add.call_typed_expr(typed_args), + } +} + +/// Walks the AST and assigns each distinct variable an auto-incremented id +/// (its offset in the input array). Repeated occurrences of the same variable +/// share the same id. +/// +/// Returns the list of input variables in id order. +fn assign_variable_ids(expr: &mut TypedExpr) -> Vec { + let mut name_to_vars: HashMap, TypedVariable> = HashMap::new(); + assign_variable_ids_aux(&mut expr.ast, &mut name_to_vars); + let mut input_vars: Vec = name_to_vars.into_values().collect(); + input_vars.sort_by_key(|var| var.variable_id); + input_vars +} + +fn assign_variable_ids_aux( + ast: &mut TypedExprAst, + name_to_vars: &mut HashMap, TypedVariable>, +) { + match ast { + TypedExprAst::Literal(_) => {} + TypedExprAst::Variable(var) => { + if let Some(typed_var) = name_to_vars.get(&var.variable_name) { + assert_eq!( + typed_var.r#type, var.r#type, + "variable `{}` appears with two different types (`{:?}` and `{:?}`); a typed \ + expr AST must be built with a single explicit type per variable", + var.variable_name, typed_var.r#type, var.r#type, + ); + var.variable_id = typed_var.variable_id; + } else { + var.variable_id = name_to_vars.len(); + name_to_vars.insert(var.variable_name.clone(), var.clone()); + }; + } + TypedExprAst::Coerce { expr, .. } => { + assign_variable_ids_aux(&mut expr.ast, name_to_vars); + } + TypedExprAst::Call { args, .. } => { + for arg in args { + assign_variable_ids_aux(&mut arg.ast, name_to_vars); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_apply_types_sum_simple() { + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("present"), + UntypedExpr::literal(1u64), + ]); + let variable_types = HashMap::from([("present", VarType::U64)]); + + let (typed_expr, _) = apply_types(&untyped_expr, &variable_types); + + assert_eq!( + typed_expr, + Function::Add + .call_typed_expr(vec![ + TypedExprAst::variable("present", VarType::U64).with_type(VarType::U64), + TypedExpr::literal(1u64), + ]) + .with_type(VarType::U64) + ); + } + + #[test] + fn test_apply_types_sum_coercion() { + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("present"), + UntypedExpr::literal(1.2f64), + ]); + let variable_types = HashMap::from([("present", VarType::U64)]); + + let (typed_expr, _) = apply_types(&untyped_expr, &variable_types); + + assert_eq!( + typed_expr, + Function::Add + .call_typed_expr(vec![ + TypedExprAst::variable("present", VarType::U64) + .with_type(VarType::U64) + .coerce(VarType::F64), + TypedExpr::literal(1.2f64), + ]) + .with_type(VarType::F64) + ); + } + + #[test] + fn test_apply_types_sum_variable_missing() { + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("present"), + Function::Add.call_untyped_expr(vec![ + UntypedExpr::literal(1u64), + UntypedExpr::variable("missing"), + ]), + ]); + let variable_types = HashMap::from([("present", VarType::U64)]); + + let (typed_expr, _) = apply_types(&untyped_expr, &variable_types); + + assert_eq!(typed_expr, TypedExpr::none()); + } + + #[test] + fn test_apply_types_to_literal() { + let untyped_expr = UntypedExpr::literal("hello"); + assert_eq!( + apply_types(&untyped_expr, &HashMap::new()).0, + TypedExprAst::literal("hello").with_type(VarType::Str) + ); + } + + #[test] + fn test_assign_variable_ids_two_variables_different_types() { + // add(x, y) with x: U64 and y: F64. Add coerces U64 to F64, so we + // get: Add(Coerce(x as F64), y). ids are assigned in DFS order. + let untyped_expr = Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("x"), UntypedExpr::variable("y")]); + let variable_types = HashMap::from([("x", VarType::U64), ("y", VarType::F64)]); + + let (_typed_expr, var_args) = apply_types(&untyped_expr, &variable_types); + + assert_eq!(var_args.len(), 2); + + assert_eq!(var_args[0].variable_name.as_ref(), "x"); + assert_eq!(var_args[0].r#type, VarType::U64); + assert_eq!(var_args[0].variable_id, 0); + assert_eq!(var_args[1].variable_name.as_ref(), "y"); + assert_eq!(var_args[1].r#type, VarType::F64); + assert_eq!(var_args[1].variable_id, 1); + } + + #[test] + #[should_panic(expected = "appears with two different types")] + fn test_assign_variable_ids_panics_on_inconsistent_types() { + // Manually build a TypedExpr where the variable `x` appears twice with + // two different types (U64 and F64). This should never happen when the + // tree is built via apply_types, so we panic to surface the bug. + let mut typed_expr = Function::Add + .call_typed_expr(vec![ + TypedExprAst::variable("x", VarType::U64).with_type(VarType::U64), + TypedExprAst::variable("x", VarType::F64).with_type(VarType::F64), + ]) + .with_type(VarType::F64); + + assign_variable_ids(&mut typed_expr); + } + + #[test] + fn test_assign_variable_ids_dedups_repeated_variable() { + // add(x, add(y, x)) — `x` appears twice and must be assigned the same id + // (single slot in the input array). Expected DFS traversal: + // x (new, id=0), y (new, id=1), x (already seen, id=0). + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("x"), + Function::Add + .call_untyped_expr(vec![UntypedExpr::variable("y"), UntypedExpr::variable("x")]), + ]); + let variable_types: HashMap<&str, VarType> = + HashMap::from([("x", VarType::U64), ("y", VarType::U64)]); + + let (_typed_expr, var_args) = apply_types(&untyped_expr, &variable_types); + + assert_eq!(var_args.len(), 2); + assert_eq!(var_args[0].variable_name.as_ref(), "x"); + assert_eq!(var_args[0].r#type, VarType::U64); + assert_eq!(var_args[0].variable_id, 0); + assert_eq!(var_args[1].variable_name.as_ref(), "y"); + assert_eq!(var_args[1].r#type, VarType::U64); + assert_eq!(var_args[1].variable_id, 1); + } +} diff --git a/jitexpr/src/ast/mod.rs b/jitexpr/src/ast/mod.rs index 868e57214..0d202f12d 100644 --- a/jitexpr/src/ast/mod.rs +++ b/jitexpr/src/ast/mod.rs @@ -1,18 +1,15 @@ +mod apply_types; mod infer_types; mod literal; mod typed_expr; mod untyped_expr; -use std::collections::HashMap; - +pub use apply_types::apply_types; pub use infer_types::{InferredTypeSet, infer_types}; pub use literal::Literal; -pub use typed_expr::TypedExprAst; +pub use typed_expr::{TypedExpr, TypedExprAst, TypedVariable}; pub use untyped_expr::UntypedExpr; -use crate::ast::typed_expr::TypedExpr; -use crate::types::VarType; - /// A function supported by the first expression-language milestone. #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] pub enum Function { @@ -34,117 +31,3 @@ impl Function { } } } - -/// 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 { - apply_types_aux(untyped_expr, &variable_types) -} - -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) => { - if let Some(variable_type) = variable_types.get(variable_name.as_ref()).copied() { - TypedExpr { - return_type: variable_type, - ast: TypedExprAst::variable(variable_name, variable_type), - } - } else { - // a missing column is treated as if it was there with a constant - // None value. - TypedExpr { - return_type: VarType::None, - ast: TypedExprAst::Literal(Literal::None), - } - } - } - UntypedExpr::Call { function, args } => match function { - Function::Add => apply_types_add_aux(args, variable_types), - }, - } -} - -fn apply_types_add_aux(args: &[UntypedExpr], variable_types: &HashMap<&str, VarType>) -> TypedExpr { - let typed_args: Vec = args - .iter() - .map(|arg| apply_types_aux(arg, variable_types)) - .collect(); - - let mut all_u64 = true; - let mut all_i64 = true; - for typed_arg in &typed_args { - match typed_arg.return_type { - VarType::U64 => all_i64 = false, - VarType::I64 => all_u64 = false, - VarType::F64 => { - all_u64 = false; - all_i64 = false; - } - _ => return TypedExpr::none(), - } - } - let return_type = if all_u64 { - VarType::U64 - } else if all_i64 { - VarType::I64 - } else { - VarType::F64 - }; - let typed_args: Vec = typed_args - .into_iter() - .map(|typed_arg| typed_arg.coerce(return_type)) - .collect(); - TypedExpr { - return_type, - ast: Function::Add.call_typed_expr(typed_args), - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_apply_types_recursively() { - 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, - Function::Add.call_typed_expr(vec![ - TypedExprAst::variable("present", VarType::U64), - Function::Add.call_typed_expr(vec![ - TypedExprAst::literal(1u64), - TypedExprAst::variable("missing", VarType::None), - ]), - ]) - ); - } - - #[test] - fn test_apply_types_to_literal() { - let untyped_expr = UntypedExpr::literal("hello"); - - assert_eq!( - apply_types(&untyped_expr, HashMap::new()), - TypedExprAst::literal("hello") - ); - } -} diff --git a/jitexpr/src/ast/typed_expr.rs b/jitexpr/src/ast/typed_expr.rs index 51b3e71d8..67fcebf5d 100644 --- a/jitexpr/src/ast/typed_expr.rs +++ b/jitexpr/src/ast/typed_expr.rs @@ -5,14 +5,9 @@ use crate::types::VarType; #[derive(Clone, PartialEq)] pub struct TypedVariable { - variable_name: Arc, - r#type: VarType, -} - -impl std::fmt::Debug for TypedVariable { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{{{}:{:?}}}", self.variable_name, self.r#type) - } + pub(super) variable_name: Arc, + pub(super) r#type: VarType, + pub(super) variable_id: usize, //< offset in the input array. } #[derive(Clone, PartialEq)] @@ -42,6 +37,12 @@ impl TypedExpr { ast: TypedExprAst::Literal(Literal::None), } } + + pub fn literal(val: impl Into) -> TypedExpr { + let literal: Literal = val.into(); + let r#type = literal.r#type(); + TypedExprAst::Literal(literal).with_type(r#type) + } } #[derive(Clone, PartialEq)] @@ -59,6 +60,13 @@ pub enum TypedExprAst { } impl TypedExprAst { + pub fn with_type(self, return_type: VarType) -> TypedExpr { + TypedExpr { + return_type, + ast: self, + } + } + pub fn literal(val: impl Into) -> TypedExprAst { TypedExprAst::Literal(val.into()) } @@ -67,12 +75,49 @@ impl TypedExprAst { TypedExprAst::Variable(TypedVariable { variable_name: Arc::from(variable_name.to_string()), r#type, + variable_id: 0, }) } } +// ---------- boilerplate --------- + impl From for TypedExprAst { fn from(literal: Literal) -> Self { TypedExprAst::Literal(literal) } } + +impl std::fmt::Debug for TypedExpr { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "({:?} : {:?})", self.ast, self.return_type) + } +} + +impl std::fmt::Debug for TypedExprAst { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + TypedExprAst::Literal(literal) => write!(f, "{:?}", literal), + TypedExprAst::Variable(variable) => write!(f, "{:?}", variable), + TypedExprAst::Coerce { target_type, expr } => { + write!(f, "coerce({:?} as {:?})", expr, target_type) + } + TypedExprAst::Call { function, args } => { + write!(f, "{:?}(", function)?; + for (i, arg) in args.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{:?}", arg)?; + } + write!(f, ")") + } + } + } +} + +impl std::fmt::Debug for TypedVariable { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{{{}:{:?}}}", self.variable_name, self.r#type) + } +} diff --git a/jitexpr/src/compile/mod.rs b/jitexpr/src/compile/mod.rs new file mode 100644 index 000000000..d09e8766f --- /dev/null +++ b/jitexpr/src/compile/mod.rs @@ -0,0 +1,68 @@ +use std::collections::HashMap; + +use cranelift_jit::JITModule; + +use crate::ast::{TypedExpr, TypedVariable, UntypedExpr, apply_types}; +use crate::types::{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, + pub input_vars: Vec, + pub typed_expr: TypedExpr, +} + +impl CompiledFunction { + pub unsafe fn call(&self, args: &[VariableValue], result: &mut VariableValue) { + debug_assert_eq!(args.len(), self.input_vars.len()); + // SAFETY: Guaranteed by the caller. + (self.entry)(args.as_ptr(), result); + } +} + +type JitEntry = unsafe extern "C" fn(*const VariableValue, *mut VariableValue); + +#[derive(Debug, thiserror::Error)] +pub enum CompileError {} + +pub fn compile( + untyped_expr: &UntypedExpr, + var_types: &HashMap<&str, VarType>, +) -> Result { + let (typed_expr, input_vars) = apply_types(&untyped_expr, var_types); + compile_typed_expr(typed_expr, input_vars) +} + +fn compile_typed_expr( + expression: TypedExpr, + input_vars: Vec, +) -> Result { + todo!(); +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::ast::{Function, UntypedExpr}; + use crate::types::VarType; + + #[test] + fn test_compile_simple() { + let untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::literal(1u64), + UntypedExpr::variable("myfield"), + ]); + let variable_types: HashMap<&str, VarType> = + std::iter::once(("myfield", VarType::U64)).collect(); + let compiled_fn = compile(&untyped_expr, &variable_types).unwrap(); + let input: Box<[VariableValue]> = vec![VariableValue { int_u64: 2u64 }].into_boxed_slice(); + let mut output: VariableValue = VariableValue { int_u64: 0u64 }; + unsafe { compiled_fn.call(&input[..], &mut output) }; + assert_eq!(unsafe { output.int_u64 }, 3u64); + } +} diff --git a/jitexpr/src/lib.rs b/jitexpr/src/lib.rs index 272ad1e56..c37d5cd83 100644 --- a/jitexpr/src/lib.rs +++ b/jitexpr/src/lib.rs @@ -1,2 +1,3 @@ pub mod ast; +pub mod compile; pub mod types; diff --git a/jitexpr/src/types.rs b/jitexpr/src/types.rs index 51fe24aa9..c2d07abf7 100644 --- a/jitexpr/src/types.rs +++ b/jitexpr/src/types.rs @@ -74,21 +74,7 @@ pub struct Signature { pub union VariableValue { pub boolean: bool, pub float: f64, + pub int_u64: u64, + pub int_i64: u64, pub string: *mut StringRef, //< this has to be mut for results. } - -// pub enum NumericalType { -// U64, -// F64, -// } - -// /// The physical column types Tantivy can offer for one field name. -// /// -// /// A name can identify both a boolean column and a numerical column. It can -// /// have at most one numerical representation. -// #[derive(Debug, Copy, Clone, Eq, PartialEq)] -// pub struct AvailableVarTypes { -// pub numerical: Option, -// pub boolean: bool, -// pub string: bool, -// }