From 483349c367740dc3de7daffb124eb085c8e1bfb6 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Fri, 14 Aug 2026 10:57:27 +0200 Subject: [PATCH] added apply_types --- jitexpr/examples/basic.rs | 14 +++- jitexpr/src/{ => ast}/infer_types.rs | 1 + jitexpr/src/ast/literal.rs | 20 +++++ jitexpr/src/ast/mod.rs | 113 ++++++++++++++++++++++++++- jitexpr/src/ast/typed_expr.rs | 37 ++++++--- jitexpr/src/lib.rs | 3 - jitexpr/src/types.rs | 8 +- 7 files changed, 173 insertions(+), 23 deletions(-) rename jitexpr/src/{ => ast}/infer_types.rs (99%) diff --git a/jitexpr/examples/basic.rs b/jitexpr/examples/basic.rs index 15af87d8e..ff70512c3 100644 --- a/jitexpr/examples/basic.rs +++ b/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> { // A simple expression that goes: @@ -22,7 +23,14 @@ fn main() -> Result<(), Box> { 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)?; diff --git a/jitexpr/src/infer_types.rs b/jitexpr/src/ast/infer_types.rs similarity index 99% rename from jitexpr/src/infer_types.rs rename to jitexpr/src/ast/infer_types.rs index d8086dfe0..c75e202c6 100644 --- a/jitexpr/src/infer_types.rs +++ b/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() diff --git a/jitexpr/src/ast/literal.rs b/jitexpr/src/ast/literal.rs index 00f64bb89..b370efae2 100644 --- a/jitexpr/src/ast/literal.rs +++ b/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), } +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 for Literal { fn from(value: bool) -> Self { Literal::Bool(value) diff --git a/jitexpr/src/ast/mod.rs b/jitexpr/src/ast/mod.rs index d6d0bbc96..3987cde3d 100644 --- a/jitexpr/src/ast/mod.rs +++ b/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::Call { + pub fn call_typed_expr(&self, args: Vec) -> 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 = 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") + ); + } } diff --git a/jitexpr/src/ast/typed_expr.rs b/jitexpr/src/ast/typed_expr.rs index feee76b56..00d1a902d 100644 --- a/jitexpr/src/ast/typed_expr.rs +++ b/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, + }, Call { function: Function, args: Vec, }, } -impl TypedExpr { - pub fn literal(val: impl Into) -> TypedExpr { - TypedExpr::Literal(val.into()) +impl TypedExprAst { + pub fn literal(val: impl Into) -> 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 for TypedExpr { +impl From for TypedExprAst { fn from(literal: Literal) -> Self { - TypedExpr::Literal(literal) + TypedExprAst::Literal(literal) } } diff --git a/jitexpr/src/lib.rs b/jitexpr/src/lib.rs index cc7d2c8c6..272ad1e56 100644 --- a/jitexpr/src/lib.rs +++ b/jitexpr/src/lib.rs @@ -1,5 +1,2 @@ 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 f2f61bc32..51fe24aa9 100644 --- a/jitexpr/src/types.rs +++ b/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.