added apply_types

jitexpr/examples/basic.rs

@@ -1,9 +1,10 @@
 use std::collections::HashMap;
 use std::error::Error;
 
-use jitexpr::ast::{Function, TypedExpr, UntypedExpr, apply_types};
+use jitexpr::ast::{
+    Function, InferredTypeSet, TypedExprAst, UntypedExpr, apply_types, infer_types,
+};
 use jitexpr::types::VarType;
-use jitexpr::{InferredTypeSet, infer_types};
 
 fn main() -> Result<(), Box<dyn Error>> {
     // A simple expression that goes:
@@ -22,7 +23,14 @@ fn main() -> Result<(), Box<dyn Error>> {
     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 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 function = compile(&expression, selected_types)?;
 

jitexpr/src/ast/infer_types.rs

@@ -149,6 +149,7 @@ fn infer_types_function_aux<'a>(
 
 fn literal_types<'a>(literal: &'a Literal) -> InferredTypeSet {
     match literal {
+        Literal::None => InferredTypeSet::ALL,
         Literal::Bool(_) => InferredTypeSet {
             boolean: true,
             ..Default::default()

jitexpr/src/ast/literal.rs

@@ -1,8 +1,11 @@
 use std::sync::Arc;
 
+use crate::types::VarType;
+
 /// A literal supported by the first expression-language milestone.
 #[derive(Clone, Debug, PartialEq)]
 pub enum Literal {
+    None,
     Bool(bool),
     U64(u64),
     I64(i64),
@@ -10,6 +13,23 @@ pub enum Literal {
     String(Arc<str>),
 }
 
+impl Literal {
+    pub fn is_none(&self) -> bool {
+        matches!(self, Literal::None)
+    }
+
+    pub fn r#type(&self) -> VarType {
+        match self {
+            Literal::None => VarType::None,
+            Literal::Bool(_) => VarType::Bool,
+            Literal::U64(_) => VarType::U64,
+            Literal::I64(_) => VarType::I64,
+            Literal::F64(_) => VarType::F64,
+            Literal::String(_) => VarType::Str,
+        }
+    }
+}
+
 impl From<bool> for Literal {
     fn from(value: bool) -> Self {
         Literal::Bool(value)

jitexpr/src/ast/mod.rs

@@ -1,13 +1,16 @@
+mod infer_types;
 mod literal;
 mod typed_expr;
 mod untyped_expr;
 
 use std::collections::HashMap;
 
+pub use infer_types::{InferredTypeSet, infer_types};
 pub use literal::Literal;
-pub use typed_expr::TypedExpr;
+pub use typed_expr::TypedExprAst;
 pub use untyped_expr::UntypedExpr;
 
+use crate::ast::typed_expr::TypedExpr;
 use crate::types::VarType;
 
 /// A function supported by the first expression-language milestone.
@@ -17,8 +20,8 @@ pub enum Function {
 }
 
 impl Function {
-    pub fn call_typed_expr(&self, args: Vec<TypedExpr>) -> TypedExpr {
-        TypedExpr::Call {
+    pub fn call_typed_expr(&self, args: Vec<TypedExpr>) -> TypedExprAst {
+        TypedExprAst::Call {
             function: *self,
             args,
         }
@@ -37,5 +40,107 @@ pub fn apply_types(
     untyped_expr: &UntypedExpr,
     variable_types: HashMap<&str, VarType>,
 ) -> TypedExpr {
-    todo!()
+    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
+    };
+    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

@@ -15,31 +15,50 @@ impl std::fmt::Debug for TypedVariable {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
-pub enum TypedExpr {
+#[derive(Clone, PartialEq)]
+pub struct TypedExpr {
+    pub return_type: VarType,
+    pub ast: TypedExprAst,
+}
+
+impl TypedExpr {
+    pub fn none() -> TypedExpr {
+        TypedExpr {
+            return_type: VarType::None,
+            ast: TypedExprAst::Literal(Literal::None),
+        }
+    }
+}
+
+#[derive(Clone, PartialEq)]
+pub enum TypedExprAst {
     Literal(Literal),
     Variable(TypedVariable),
+    Coerce {
+        target_type: VarType,
+        expr: Box<TypedExpr>,
+    },
     Call {
         function: Function,
         args: Vec<TypedExpr>,
     },
 }
 
-impl TypedExpr {
-    pub fn literal(val: impl Into<Literal>) -> TypedExpr {
-        TypedExpr::Literal(val.into())
+impl TypedExprAst {
+    pub fn literal(val: impl Into<Literal>) -> TypedExprAst {
+        TypedExprAst::Literal(val.into())
     }
 
-    pub fn variable(variable_name: impl ToString, r#type: VarType) -> TypedExpr {
-        TypedExpr::Variable(TypedVariable {
+    pub fn variable(variable_name: impl ToString, r#type: VarType) -> TypedExprAst {
+        TypedExprAst::Variable(TypedVariable {
             variable_name: Arc::from(variable_name.to_string()),
             r#type,
         })
     }
 }
 
-impl From<Literal> for TypedExpr {
+impl From<Literal> for TypedExprAst {
     fn from(literal: Literal) -> Self {
-        TypedExpr::Literal(literal)
+        TypedExprAst::Literal(literal)
     }
 }

jitexpr/src/lib.rs

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

jitexpr/src/types.rs

@@ -1,14 +1,14 @@
 //! Source types and nullable runtime value representations.
 
 /// A value type supported by compiled expressions.
-#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Ord, PartialOrd)]
 pub enum VarType {
     Bool,
-    U64,
     F64,
+    U64,
+    I64,
     Str,
-    None,
-    // TODO: add other types.
+    None, // TODO: add other types.
 }
 
 /// A borrowed UTF-8 string descriptor passed opaquely through generated code.