Added doc/unit test

jitexpr/README.md

@@ -6,4 +6,7 @@ This is an expression compiler relying on Cranelift.
       ↓ lowering
   Cranelift IR
       ↓ Cranelift code generation
-  Machine code
+  Machine code (or assembly)
+
+The project does not rely on cranelifts function call abstraction.
+Instead it just manipulates expression, so everything is always inlined.

jitexpr/src/FUNCTION.md

@@ -1,3 +0,0 @@
-
-Regexp:
-only 3 args supported. group is not optional. if not passed, set 0 during the conversion from proto to rust object.

jitexpr/src/bin/jitexpr-asm.rs

@@ -7,9 +7,10 @@
 use std::collections::HashMap;
 use std::io::{self, BufRead, Write};
 use std::process::ExitCode;
+use std::time::{Duration, Instant};
 
 use jitexpr::ast::{DeserializeError, InferredTypeSet, TypeError, deserialize, infer_types};
-use jitexpr::compile::{CompileError, compile_to_assembly};
+use jitexpr::compile::{CompileError, compile, compile_to_assembly};
 use jitexpr::types::VarType;
 
 #[derive(Debug, thiserror::Error)]
@@ -22,6 +23,11 @@ enum ExpressionError {
     Compile(#[from] CompileError),
 }
 
+struct LineOutput {
+    assembly: String,
+    codegen_duration: Duration,
+}
+
 fn main() -> ExitCode {
     let stdin = io::stdin();
     let stdout = io::stdout();
@@ -52,12 +58,17 @@ fn process_lines(
         }
 
         match compile_line(&line) {
-            Ok(assembly) => {
+            Ok(line_output) => {
                 if wrote_assembly {
                     writeln!(output)?;
                 }
-                output.write_all(assembly.as_bytes())?;
-                if !assembly.ends_with('\n') {
+                writeln!(
+                    output,
+                    "Code generation time: {:?}",
+                    line_output.codegen_duration
+                )?;
+                output.write_all(line_output.assembly.as_bytes())?;
+                if !line_output.assembly.ends_with('\n') {
                     writeln!(output)?;
                 }
                 wrote_assembly = true;
@@ -72,14 +83,24 @@ fn process_lines(
     Ok(all_succeeded)
 }
 
-fn compile_line(line: &str) -> Result<String, ExpressionError> {
+fn compile_line(line: &str) -> Result<LineOutput, ExpressionError> {
     let expression = deserialize(line)?;
     let inferred_types = infer_types(&expression)?;
     let variable_types = inferred_types
         .into_iter()
         .map(|(name, inferred_type)| (name, concrete_type(inferred_type)))
         .collect::<HashMap<_, _>>();
-    Ok(compile_to_assembly(&expression, &variable_types)?)
+
+    let codegen_start = Instant::now();
+    let compiled_fn = compile(&expression, &variable_types)?;
+    let codegen_duration = codegen_start.elapsed();
+    drop(compiled_fn);
+
+    let assembly = compile_to_assembly(&expression, &variable_types)?;
+    Ok(LineOutput {
+        assembly,
+        codegen_duration,
+    })
 }
 
 fn concrete_type(inferred_type: InferredTypeSet) -> VarType {
@@ -108,10 +129,10 @@ mod tests {
 
     #[test]
     fn test_compile_line_infers_variable_type() {
-        let assembly = compile_line("(ADD 1i64 my_col)").unwrap();
+        let line_output = compile_line("(ADD 1i64 my_col)").unwrap();
 
-        assert!(assembly.contains("block0:"));
-        assert!(!assembly.trim().is_empty());
+        assert!(line_output.assembly.contains("block0:"));
+        assert!(!line_output.assembly.trim().is_empty());
     }
 
     #[test]
@@ -133,6 +154,7 @@ mod tests {
 
         assert!(!all_succeeded);
         let output = String::from_utf8(output).unwrap();
+        assert_eq!(output.matches("Code generation time:").count(), 2);
         assert_eq!(output.matches("block0:").count(), 2);
         assert!(String::from_utf8(errors).unwrap().contains("line 2:"));
     }

jitexpr/src/compile/compile_fn_builder.rs

@@ -13,8 +13,7 @@ use regex::Regex;
 
 use super::compiled_fn::JitEntry;
 use super::{
-    CompileError, CompiledFn, LoweringContext, TypedExpr, TypedExprAst, TypedLiteral,
-    TypedVariable, lower_expr,
+    CompileError, CompiledFn, LoweringContext, TypedExpr, TypedExprAst, TypedLiteral, TypedVariable,
 };
 use crate::ast::{InferredTypeSet, Literal, UntypedExpr};
 use crate::functions::{declare_native_functions, register_jit_symbols};
@@ -266,7 +265,7 @@ impl<'types, 'names> CompileFnBuilder<'types, 'names> {
                 regex_match_results: &regex_match_results,
                 native_functions: &native_functions,
             };
-            let value = lower_expr(&expression, &mut lowering_context, &mut builder)?;
+            let value = lowering_context.compile_expr(&expression, &mut builder)?;
             builder
                 .ins()
                 .store(MemFlagsData::trusted(), value, result_ptr, 0);

jitexpr/src/compile/mod.rs

@@ -46,12 +46,35 @@ pub(crate) struct LoweringContext<'a> {
 }
 
 impl LoweringContext<'_> {
-    pub(crate) fn lower_expr(
+    pub(crate) fn compile_expr(
         &mut self,
         expression: &TypedExpr,
         builder: &mut FunctionBuilder<'_>,
     ) -> Result<Value, CompileError> {
-        lower_expr(expression, self, builder)
+        match &expression.ast {
+            TypedExprAst::Literal(literal) => Ok(lower_literal(literal, self, builder)),
+            TypedExprAst::Variable(variable) => {
+                let byte_offset = variable
+                    .variable_id
+                    .checked_mul(size_of::<VariableValue>())
+                    .and_then(|offset| i32::try_from(offset).ok())
+                    .ok_or(CompileError::InputOffsetOverflow {
+                        variable_id: variable.variable_id,
+                    })?;
+                Ok(builder.ins().load(
+                    cranelift_type(variable.r#type, self.pointer_type),
+                    MemFlagsData::trusted(),
+                    self.args_ptr,
+                    byte_offset,
+                ))
+            }
+            TypedExprAst::Coerce { target_type, expr } => {
+                let source_type = expr.return_type;
+                let value = self.compile_expr(expr, builder)?;
+                lower_coercion(value, source_type, *target_type, builder)
+            }
+            TypedExprAst::FnCall(fn_call) => fn_call.lower(expression.return_type, self, builder),
+        }
     }
 
     pub(crate) fn pointer_type(&self) -> Type {
@@ -71,38 +94,6 @@ impl LoweringContext<'_> {
     }
 }
 
-/// Produce cranelift IR from the function
-fn lower_expr(
-    expression: &TypedExpr,
-    context: &mut LoweringContext<'_>,
-    builder: &mut FunctionBuilder<'_>,
-) -> Result<Value, CompileError> {
-    match &expression.ast {
-        TypedExprAst::Literal(literal) => Ok(lower_literal(literal, context, builder)),
-        TypedExprAst::Variable(variable) => {
-            let byte_offset = variable
-                .variable_id
-                .checked_mul(size_of::<VariableValue>())
-                .and_then(|offset| i32::try_from(offset).ok())
-                .ok_or(CompileError::InputOffsetOverflow {
-                    variable_id: variable.variable_id,
-                })?;
-            Ok(builder.ins().load(
-                cranelift_type(variable.r#type, context.pointer_type),
-                MemFlagsData::trusted(),
-                context.args_ptr,
-                byte_offset,
-            ))
-        }
-        TypedExprAst::Coerce { target_type, expr } => {
-            let source_type = expr.return_type;
-            let value = lower_expr(expr, context, builder)?;
-            lower_coercion(value, source_type, *target_type, builder)
-        }
-        TypedExprAst::FnCall(fn_call) => fn_call.lower(expression.return_type, context, builder),
-    }
-}
-
 fn lower_literal(
     literal: &TypedLiteral,
     context: &mut LoweringContext<'_>,

jitexpr/src/functions/add.rs

@@ -108,7 +108,7 @@ impl FnCall for AddFnCall {
         };
 
         for arg in &self.args {
-            let value = context.lower_expr(arg, builder)?;
+            let value = context.compile_expr(arg, builder)?;
             sum = match return_type {
                 VarType::U64 | VarType::I64 => builder.ins().iadd(sum, value),
                 VarType::F64 => builder.ins().fadd(sum, value),
@@ -268,10 +268,25 @@ mod tests {
     #[test]
     fn test_call_with_types_uses_u64_when_i64_is_not_possible() {
         let variable_types = HashMap::new();
+        let expression = crate::ast::deserialize("(ADD 9223372036854775808u64)").unwrap();
         let typed_expr =
             crate::typed_expr_from_str("(ADD 9223372036854775808u64)", &variable_types);
-
         assert_eq!(typed_expr.return_type, VarType::U64);
+        assert_eq!(
+            typed_expr,
+            TypedExpr {
+                return_type: VarType::U64,
+                ast: TypedExprAst::from_call(AddFnCall {
+                    args: vec![TypedExpr::literal(9223372036854775808u64)].into_boxed_slice(),
+                }),
+            }
+        );
+
+        let compiled = compile(&expression, &variable_types).unwrap();
+        let mut output = VariableValue { int_u64: 0 };
+        unsafe { compiled.call(&[], &mut output) };
+
+        assert_eq!(unsafe { output.int_u64 }, 9223372036854775808u64);
     }
 
     #[test]

jitexpr/src/functions/mod.rs

@@ -112,9 +112,11 @@ pub(crate) trait FnCall: std::fmt::Debug + Into<FnCallEnum> {
 
     /// Builds the typed call after concrete variable types have been supplied.
     ///
-    /// `target_type_set` communicates the result types preferred by the parent call. The
+    /// `target_type_set` communicates the result types set accepted by the parent call. The
     /// implementation selects a concrete result type, applies compatible target types to its
     /// arguments through `context`, and registers any compilation resources owned by the call.
+    ///
+    /// The type of the returned is given to the caller in the TypedExpr object.
     fn call_with_types(
         args: &[UntypedExpr],
         target_type_set: InferredTypeSet,
@@ -127,6 +129,8 @@ pub(crate) trait FnCall: std::fmt::Debug + Into<FnCallEnum> {
     ///
     /// This is only used, to assign and deduplicate variable input slots. Compile-time
     /// configuration stored directly on a call does not need to be returned.
+    ///
+    /// Today this is only used as a cheap visitor to allocate variable ids.
     fn args_mut(&mut self) -> &mut [TypedExpr];
 
     /// Emits Cranelift IR for an already typed call and returns its result SSA value.

jitexpr/src/functions/regexp_extract.rs

@@ -1,7 +1,11 @@
 // RegexpExtract extracts a regular-expression match from a string.
 //
-// It takes three arguments: the input string, a regular-expression pattern literal, and a u64
-// capture index literal. Capture index 0 returns the full match, while indexes 1 and above return
+// It takes three arguments:
+// - string: the input string
+// - const string: a regular-expression pattern literal. This one CANNOT be the result of another
+//   expression
+// - const u64: capture index literal. Capture index 0 returns the full match, while indexes 1 and
+//   above return
 // the corresponding explicit capture group.
 //
 // It returns None when the input is None, the pattern does not match, or the requested capture
@@ -104,7 +108,7 @@ impl FnCall for RegexpExtractFnCall {
     ) -> Result<Value, CompileError> {
         debug_assert_eq!(return_type, VarType::Str);
 
-        let haystack = context.lower_expr(&self.haystack, builder)?;
+        let haystack = context.compile_expr(&self.haystack, builder)?;
         let regex_index = builder
             .ins()
             .iconst(context.pointer_type(), self.regex_ref.index() as i64);