typed/untyped

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"

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<dyn Error>> {
     // 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(())
 }
 

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<str>),
+}
 
 impl From<bool> for Literal {
     fn from(value: bool) -> Self {
@@ -6,6 +16,12 @@ impl From<bool> for Literal {
     }
 }
 
+impl From<u64> for Literal {
+    fn from(value: u64) -> Self {
+        Literal::U64(value)
+    }
+}
+
 impl From<i64> for Literal {
     fn from(value: i64) -> Self {
         Literal::I64(value)
@@ -20,18 +36,12 @@ impl From<f64> for Literal {
 
 impl From<String> 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<Literal> for Expr {
-    fn from(literal: Literal) -> Self {
-        Expr::Literal(literal)
+        Literal::String(Arc::from(value.to_string()))
     }
 }

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<Expr> },
-}
-
-impl Expr {
-    pub fn literal(val: impl Into<Literal>) -> 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 {
-        Expr::Call {
+    pub fn call_typed_expr(&self, args: Vec<TypedExpr>) -> TypedExpr {
+        TypedExpr::Call {
+            function: *self,
+            args,
+        }
+    }
+
+    pub fn call_untyped_expr(&self, args: Vec<UntypedExpr>) -> UntypedExpr {
+        UntypedExpr::Call {
             function: *self,
             args,
         }
     }
 }
 
-impl Expr {
-    pub fn list_variable_names(&self) -> HashSet<String> {
-        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!()
 }

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<str>,
+    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<TypedExpr>,
+    },
+}
+
+impl TypedExpr {
+    pub fn literal(val: impl Into<Literal>) -> 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<Literal> for TypedExpr {
+    fn from(literal: Literal) -> Self {
+        TypedExpr::Literal(literal)
+    }
+}

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<str>),
+    Call {
+        function: Function,
+        args: Vec<UntypedExpr>,
+    },
+}
+
+impl UntypedExpr {
+    pub fn literal(val: impl Into<Literal>) -> UntypedExpr {
+        UntypedExpr::Literal(val.into())
+    }
+
+    pub fn variable(variable_name: impl ToString) -> UntypedExpr {
+        UntypedExpr::Variable(Arc::from(variable_name.to_string()))
+    }
+}
+
+impl From<Literal> for UntypedExpr {
+    fn from(literal: Literal) -> Self {
+        UntypedExpr::Literal(literal)
+    }
+}

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<HashMap<&'a 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>(
+    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));
+            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<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)
+        }
+    }
+}
+
+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);
+    }
+}

jitexpr/src/lib.rs

@@ -1,2 +1,5 @@
 pub mod ast;
+mod infer_types;
 pub mod types;
+
+pub use infer_types::{InferredTypeSet, infer_types};

jitexpr/src/types.rs

@@ -7,6 +7,7 @@ pub enum VarType {
     U64,
     F64,
     Str,
+    None,
     // TODO: add other types.
 }