diff --git a/jitexpr/Cargo.toml b/jitexpr/Cargo.toml index e2738c172..aeef2fc5c 100644 --- a/jitexpr/Cargo.toml +++ b/jitexpr/Cargo.toml @@ -7,3 +7,4 @@ edition = "2024" cranelift = "0.134.3" cranelift-jit = "0.134.3" cranelift-module = "0.134.3" +thiserror = "2.0.1" diff --git a/jitexpr/examples/basic.rs b/jitexpr/examples/basic.rs index e47c5500c..15af87d8e 100644 --- a/jitexpr/examples/basic.rs +++ b/jitexpr/examples/basic.rs @@ -1,42 +1,33 @@ use std::collections::HashMap; use std::error::Error; -use jitexpr::ast::{Expr, Function, Literal}; +use jitexpr::ast::{Function, TypedExpr, UntypedExpr, apply_types}; +use jitexpr::types::VarType; +use jitexpr::{InferredTypeSet, infer_types}; fn main() -> Result<(), Box> { // A simple expression that goes: // my_column + 1 - let expression = Function::Add.call_expr(vec![Expr::variable("my_col"), Expr::literal(1.0f64)]); - - 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 untyped_expr = Function::Add.call_untyped_expr(vec![ + UntypedExpr::variable("my_col"), + UntypedExpr::literal(1.0f64), ]); - let selected_types = infer_types(&expression, &available_types)?; - let function = compile(&expression, selected_types)?; + let inferred_types = infer_types(&untyped_expr)?; + assert_eq!( + inferred_types.get("my_col").unwrap(), + &InferredTypeSet::NUMERICAL + ); - assert_eq!(evaluate(&function, 100, 4), 25.0); - assert_eq!(evaluate(&function, 100, 0), 0.0); + let variable_types: HashMap<&str, VarType> = + std::iter::once(("my_col", VarType::F64)).collect(); + + let typed_expr: TypedExpr = apply_types(&untyped_expr, variable_types); + + // let function = compile(&expression, selected_types)?; + + // assert_eq!(evaluate(&function, 100, 4), 25.0); + // assert_eq!(evaluate(&function, 100, 0), 0.0); Ok(()) } diff --git a/jitexpr/src/ast/boilerplate.rs b/jitexpr/src/ast/literal.rs similarity index 52% rename from jitexpr/src/ast/boilerplate.rs rename to jitexpr/src/ast/literal.rs index 26dafa98e..00f64bb89 100644 --- a/jitexpr/src/ast/boilerplate.rs +++ b/jitexpr/src/ast/literal.rs @@ -1,4 +1,14 @@ -use crate::ast::{Expr, Literal}; +use std::sync::Arc; + +/// A literal supported by the first expression-language milestone. +#[derive(Clone, Debug, PartialEq)] +pub enum Literal { + Bool(bool), + U64(u64), + I64(i64), + F64(f64), + String(Arc), +} impl From for Literal { fn from(value: bool) -> Self { @@ -6,6 +16,12 @@ impl From for Literal { } } +impl From for Literal { + fn from(value: u64) -> Self { + Literal::U64(value) + } +} + impl From for Literal { fn from(value: i64) -> Self { Literal::I64(value) @@ -20,18 +36,12 @@ impl From for Literal { impl From for Literal { fn from(value: String) -> Self { - Literal::String(value) + Literal::String(Arc::from(value)) } } impl From<&str> for Literal { fn from(value: &str) -> Self { - Literal::String(value.to_owned()) - } -} - -impl From for Expr { - fn from(literal: Literal) -> Self { - Expr::Literal(literal) + Literal::String(Arc::from(value.to_string())) } } diff --git a/jitexpr/src/ast/mod.rs b/jitexpr/src/ast/mod.rs index 7180738a5..d6d0bbc96 100644 --- a/jitexpr/src/ast/mod.rs +++ b/jitexpr/src/ast/mod.rs @@ -1,33 +1,14 @@ -use std::collections::HashSet; +mod literal; +mod typed_expr; +mod untyped_expr; -mod boilerplate; +use std::collections::HashMap; -/// A literal supported by the first expression-language milestone. -#[derive(Clone, Debug, PartialEq)] -pub enum Literal { - Bool(bool), - I64(i64), - F64(f64), - String(String), -} +pub use literal::Literal; +pub use typed_expr::TypedExpr; +pub use untyped_expr::UntypedExpr; -/// An expression independent from its protobuf representation. -#[derive(Clone, Debug, PartialEq)] -pub enum Expr { - Literal(Literal), - Variable(String), - Call { function: Function, args: Vec }, -} - -impl Expr { - pub fn literal(val: impl Into) -> Expr { - Expr::Literal(val.into()) - } - - pub fn variable(variable_name: impl ToString) -> Expr { - Expr::Variable(variable_name.to_string()) - } -} +use crate::types::VarType; /// A function supported by the first expression-language milestone. #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] @@ -36,28 +17,25 @@ pub enum Function { } impl Function { - pub fn call_expr(&self, args: Vec) -> Expr { - Expr::Call { + pub fn call_typed_expr(&self, args: Vec) -> TypedExpr { + TypedExpr::Call { + function: *self, + args, + } + } + + pub fn call_untyped_expr(&self, args: Vec) -> UntypedExpr { + UntypedExpr::Call { function: *self, args, } } } -impl Expr { - pub fn list_variable_names(&self) -> HashSet { - let mut names = HashSet::new(); - match self { - Expr::Literal(_) => {} - Expr::Variable(name) => { - names.insert(name.clone()); - } - Expr::Call { args, .. } => { - for arg in args { - names.extend(arg.list_variable_names()); - } - } - } - names - } +/// 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 { + todo!() } diff --git a/jitexpr/src/ast/typed_expr.rs b/jitexpr/src/ast/typed_expr.rs new file mode 100644 index 000000000..feee76b56 --- /dev/null +++ b/jitexpr/src/ast/typed_expr.rs @@ -0,0 +1,45 @@ +use std::sync::Arc; + +use crate::ast::{Function, Literal}; +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) + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum TypedExpr { + Literal(Literal), + Variable(TypedVariable), + Call { + function: Function, + args: Vec, + }, +} + +impl TypedExpr { + pub fn literal(val: impl Into) -> TypedExpr { + TypedExpr::Literal(val.into()) + } + + pub fn variable(variable_name: impl ToString, r#type: VarType) -> TypedExpr { + TypedExpr::Variable(TypedVariable { + variable_name: Arc::from(variable_name.to_string()), + r#type, + }) + } +} + +impl From for TypedExpr { + fn from(literal: Literal) -> Self { + TypedExpr::Literal(literal) + } +} diff --git a/jitexpr/src/ast/untyped_expr.rs b/jitexpr/src/ast/untyped_expr.rs new file mode 100644 index 000000000..113169742 --- /dev/null +++ b/jitexpr/src/ast/untyped_expr.rs @@ -0,0 +1,30 @@ +use std::sync::Arc; + +use crate::ast::{Function, Literal}; + +/// An expression independent from its protobuf representation. +#[derive(Clone, Debug, PartialEq)] +pub enum UntypedExpr { + Literal(Literal), + Variable(Arc), + Call { + function: Function, + args: Vec, + }, +} + +impl UntypedExpr { + pub fn literal(val: impl Into) -> UntypedExpr { + UntypedExpr::Literal(val.into()) + } + + pub fn variable(variable_name: impl ToString) -> UntypedExpr { + UntypedExpr::Variable(Arc::from(variable_name.to_string())) + } +} + +impl From for UntypedExpr { + fn from(literal: Literal) -> Self { + UntypedExpr::Literal(literal) + } +} diff --git a/jitexpr/src/infer_types.rs b/jitexpr/src/infer_types.rs new file mode 100644 index 000000000..d8086dfe0 --- /dev/null +++ b/jitexpr/src/infer_types.rs @@ -0,0 +1,231 @@ +use std::collections::HashMap; +use std::collections::hash_map::Entry; + +use crate::ast::{Function, Literal, UntypedExpr}; + +#[derive(Default, Copy, Clone, Debug, Eq, PartialEq)] +pub struct InferredTypeSet { + string: bool, + numerical: bool, + boolean: bool, +} + +impl InferredTypeSet { + pub const NONE: InferredTypeSet = InferredTypeSet { + string: false, + numerical: false, + boolean: false, + }; + + pub const ALL: InferredTypeSet = InferredTypeSet { + string: true, + numerical: true, + boolean: true, + }; + + pub const NUMERICAL: InferredTypeSet = InferredTypeSet { + numerical: true, + boolean: false, + string: false, + }; + + pub const STRING: InferredTypeSet = InferredTypeSet { + numerical: false, + boolean: false, + string: true, + }; + + fn is_none(self) -> bool { + self == Self::NONE + } + + fn intersect(self, target_inferred_type: InferredTypeSet) -> InferredTypeSet { + InferredTypeSet { + string: self.string && target_inferred_type.string, + numerical: self.numerical && target_inferred_type.numerical, + boolean: self.boolean && target_inferred_type.boolean, + } + } +} + +impl std::fmt::Display for InferredTypeSet { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let mut types = Vec::new(); + if self.string { + types.push("string"); + } + if self.numerical { + types.push("numerical"); + } + if self.boolean { + types.push("boolean"); + } + write!(f, "{{{}}}", types.join(", ")) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum TypeError { + #[error("function `{function:?}` returns a number, expected `{expected}`")] + WrongFunctionReturnType { + function: Function, + expected: InferredTypeSet, + }, + #[error("expected `{expected}` , got `{literal:?}`")] + InvalidLiteralType { + literal: Literal, + expected: InferredTypeSet, + }, +} + +/// Infer the accepted types for the different variables present in the formula. +pub fn infer_types<'a>( + expr: &'a UntypedExpr, +) -> Result, TypeError> { + let mut inferred_type_res = HashMap::default(); + infer_types_aux(expr, InferredTypeSet::ALL, &mut inferred_type_res)?; + Ok(inferred_type_res) +} + +fn infer_types_aux<'a>( + expr: &'a UntypedExpr, + target_inferred_type: InferredTypeSet, + inferred_types_res: &mut HashMap<&'a str, InferredTypeSet>, +) -> Result { + match expr { + UntypedExpr::Literal(literal) => { + let literal_type: InferredTypeSet = + target_inferred_type.intersect(literal_types(literal)); + if literal_type.is_none() { + return Err(TypeError::InvalidLiteralType { + literal: literal.clone(), + expected: target_inferred_type, + }); + } + Ok(literal_type) + } + UntypedExpr::Variable(variable_name) => match inferred_types_res.entry(&*variable_name) { + Entry::Occupied(mut occupied_entry) => { + let inferred_types = occupied_entry.get().intersect(target_inferred_type); + occupied_entry.insert(inferred_types); + Ok(inferred_types) + } + Entry::Vacant(vacant_entry) => { + vacant_entry.insert_entry(target_inferred_type); + Ok(target_inferred_type) + } + }, + UntypedExpr::Call { function, args } => infer_types_function_aux( + *function, + &args[..], + target_inferred_type, + inferred_types_res, + ), + } +} + +fn infer_types_function_aux<'a>( + function: Function, + args: &'a [UntypedExpr], + target_inferred_type: InferredTypeSet, + inferred_types_res: &mut HashMap<&'a str, InferredTypeSet>, +) -> Result { + match function { + Function::Add => { + // This is valid for all functions taking a bunch of number and returning a number. + if !target_inferred_type.numerical { + return Err(TypeError::WrongFunctionReturnType { + function, + expected: target_inferred_type, + }); + } + for arg in args { + infer_types_aux(arg, InferredTypeSet::NUMERICAL, inferred_types_res)?; + } + Ok(InferredTypeSet::NUMERICAL) + } + } +} + +fn literal_types<'a>(literal: &'a Literal) -> InferredTypeSet { + match literal { + Literal::Bool(_) => InferredTypeSet { + boolean: true, + ..Default::default() + }, + Literal::I64(_) | Literal::U64(_) | Literal::F64(_) => InferredTypeSet::NUMERICAL, + Literal::String(_) => InferredTypeSet::STRING, + } +} + +#[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() { + // A lone variable should accept all types. + let expr = UntypedExpr::variable("a"); + let inferred_types = infer_types(&expr).unwrap(); + let a_types = inferred_types.get("a").unwrap(); + assert_eq!(a_types, &InferredTypeSet::ALL); + } +} diff --git a/jitexpr/src/lib.rs b/jitexpr/src/lib.rs index 272ad1e56..cc7d2c8c6 100644 --- a/jitexpr/src/lib.rs +++ b/jitexpr/src/lib.rs @@ -1,2 +1,5 @@ pub mod ast; +mod infer_types; pub mod types; + +pub use infer_types::{InferredTypeSet, infer_types}; diff --git a/jitexpr/src/types.rs b/jitexpr/src/types.rs index f66a9a01c..f2f61bc32 100644 --- a/jitexpr/src/types.rs +++ b/jitexpr/src/types.rs @@ -7,6 +7,7 @@ pub enum VarType { U64, F64, Str, + None, // TODO: add other types. }