second stab for jitexpr

This commit is contained in:
Paul Masurel
2026-08-14 14:42:27 +02:00
parent 8964bacc35
commit 174fdfb4ab
5 changed files with 31 additions and 19 deletions
+9 -2
View File
@@ -1,24 +1,31 @@
use std::collections::HashMap;
use std::error::Error;
use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, apply_types, infer_types};
use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, infer_types};
use jitexpr::compile::{CompiledFunction, compile};
use jitexpr::types::VarType;
fn main() -> Result<(), Box<dyn Error>> {
// A simple expression that goes:
// my_column + 1
let untyped_expr = Function::Add.call_untyped_expr(vec![
UntypedExpr::variable("my_col"),
UntypedExpr::literal(1.0f64),
]);
let inferred_types = infer_types(&untyped_expr)?;
// Infer types does not return specific types, but instead a set of acceptable
// types for each variables.
let inferred_types: HashMap<&str, InferredTypeSet> = infer_types(&untyped_expr)?;
assert_eq!(
inferred_types.get("my_col").unwrap(),
&InferredTypeSet::NUMERICAL
);
// This is then up to us to decide the actual type for each variable.
// In tantivy, this means picking the first column with a type in inferred_types.
//
// If none match then we should use the VarType::None.
let variable_types: HashMap<&str, VarType> =
std::iter::once(("my_col", VarType::F64)).collect();
+2 -4
View File
@@ -1,15 +1,13 @@
mod apply_types;
mod infer_types;
mod literal;
mod typed_expr;
mod untyped_expr;
pub use apply_types::apply_types;
pub use infer_types::{InferredTypeSet, infer_types};
pub use literal::Literal;
pub use typed_expr::{TypedExpr, TypedExprAst, TypedVariable};
pub use untyped_expr::UntypedExpr;
use crate::compile::{TypedExpr, TypedExprAst};
/// A function supported by the first expression-language milestone.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Function {
@@ -1,8 +1,8 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::ast::typed_expr::TypedVariable;
use crate::ast::{Function, Literal, TypedExpr, TypedExprAst, UntypedExpr};
use super::typed_expr::{TypedExpr, TypedExprAst, TypedVariable};
use crate::ast::{Function, Literal, UntypedExpr};
use crate::types::VarType;
/// If a variable is missing from variable_types, it will be treated as if its value is None.
@@ -25,18 +25,22 @@ fn apply_types_aux(
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 {
let variable_type: VarType = variable_types
.get(variable_name.as_ref())
.copied()
// a missing column is treated as if it was there with a constant
// None value.
.unwrap_or(VarType::None);
if variable_type == VarType::None {
TypedExpr {
return_type: VarType::None,
ast: TypedExprAst::Literal(Literal::None),
}
} else {
TypedExpr {
return_type: variable_type,
ast: TypedExprAst::variable(variable_name, variable_type),
}
}
}
UntypedExpr::Call { function, args } => match function {
+6 -3
View File
@@ -1,14 +1,17 @@
mod apply_types;
mod typed_expr;
use std::collections::HashMap;
use std::mem::{self, size_of};
pub use apply_types::apply_types;
use cranelift::codegen::ir::{MemFlagsData, UserFuncName};
use cranelift::prelude::*;
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{Module, default_libcall_names};
pub use typed_expr::{TypedExpr, TypedExprAst, TypedVariable};
use crate::ast::{
Function, Literal, TypedExpr, TypedExprAst, TypedVariable, UntypedExpr, apply_types,
};
use crate::ast::{Function, Literal, UntypedExpr};
use crate::types::{StringRef, VarType, VariableValue};
/// An expression compiled to native machine code.
@@ -6,8 +6,8 @@ use crate::types::VarType;
#[derive(Clone, PartialEq)]
pub struct TypedVariable {
pub(super) variable_name: Arc<str>,
pub(crate) r#type: VarType,
pub(crate) variable_id: usize, //< offset in the input array.
pub(super) r#type: VarType,
pub(super) variable_id: usize, //< offset in the input array.
}
#[derive(Clone, PartialEq)]