first stab at compile

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<dyn Error>> {
@@ -23,14 +22,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     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)?;
 

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<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) => {
+            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<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/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<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),
-    }
-}
-
-#[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")
-        );
-    }
-}

jitexpr/src/ast/typed_expr.rs

@@ -5,14 +5,9 @@ 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)
-    }
+    pub(super) variable_name: Arc<str>,
+    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<Literal>) -> 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<Literal>) -> 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<Literal> 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)
+    }
+}

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<TypedVariable>,
+    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<CompiledFunction, CompileError> {
+    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<TypedVariable>,
+) -> Result<CompiledFunction, CompileError> {
+    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);
+    }
+}

jitexpr/src/lib.rs

@@ -1,2 +1,3 @@
 pub mod ast;
+pub mod compile;
 pub mod types;

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<NumericalType>,
-//     pub boolean: bool,
-//     pub string: bool,
-// }