diff --git a/engine/src/ast/field_expr.rs b/engine/src/ast/field_expr.rs index 36f1f89d..d5b23d4a 100644 --- a/engine/src/ast/field_expr.rs +++ b/engine/src/ast/field_expr.rs @@ -795,6 +795,7 @@ impl Expr for ComparisonExpr { #[allow(clippy::bool_assert_comparison)] mod tests { use super::*; + use crate::ast::ValueExpr; use crate::ast::function_expr::{FunctionCallArgExpr, FunctionCallExpr}; use crate::ast::logical_expr::LogicalExpr; use crate::execution_context::ExecutionContext; @@ -1034,6 +1035,26 @@ mod tests { }, ) .unwrap(); + builder + .add_function( + "concat_fields", + SimpleFunctionDefinition { + params: vec![ + SimpleFunctionParam { + arg_kind: SimpleFunctionArgKind::Field, + val_type: Type::Bytes, + }, + SimpleFunctionParam { + arg_kind: SimpleFunctionArgKind::Field, + val_type: Type::Bytes, + }, + ], + opt_params: vec![], + return_type: Type::Bytes, + implementation: SimpleFunctionImpl::new(concat_function), + }, + ) + .unwrap(); builder .add_function("filter", FilterFunction::new()) .unwrap(); @@ -2207,6 +2228,86 @@ mod tests { assert_eq!(expr.execute_one(ctx), true); } + // The non-mapped argument (the "-cf" literal) must be applied to *every* + // mapped element, even though it is now evaluated only once per call. + #[test] + fn test_map_each_function_non_mapped_arg_applied_to_all_elements() { + let (expr, rest) = FilterParser::new(&SCHEME) + .lex_as::(r#"concat(http.cookies[*], "-cf")"#) + .unwrap(); + assert_eq!(rest, ""); + + let expr = expr.compile(); + let ctx = &mut ExecutionContext::new(&SCHEME); + ctx.set_field_value( + field("http.cookies"), + Array::from_iter(["one", "two", "three"]), + ) + .unwrap(); + + assert_eq!( + expr.execute(ctx), + Ok(LhsValue::Array(Array::from_iter([ + "one-cf", "two-cf", "three-cf" + ]))) + ); + } + + // A non-mapped argument that is expensive to re-evaluate (here a nested + // function call) is evaluated once and reused for every mapped element. + #[test] + fn test_map_each_memoizes_expensive_non_mapped_arg() { + let (expr, rest) = FilterParser::new(&SCHEME) + .lex_as::(r#"concat_fields(http.cookies[*], lowercase(http.host))"#) + .unwrap(); + assert_eq!(rest, ""); + + let expr = expr.compile(); + let ctx = &mut ExecutionContext::new(&SCHEME); + ctx.set_field_value( + field("http.cookies"), + Array::from_iter(["one", "two", "three"]), + ) + .unwrap(); + ctx.set_field_value(field("http.host"), "SUFFIX").unwrap(); + + // `lowercase(http.host)` == "suffix" is appended to every element. + assert_eq!( + expr.execute(ctx), + Ok(LhsValue::Array(Array::from_iter([ + "onesuffix", + "twosuffix", + "threesuffix" + ]))) + ); + } + + // map_each over a Map with no extra args: exercises the fused Map -> Array + // path (no intermediate array allocation) and the empty-args fast path. + #[test] + fn test_map_each_on_map_no_extra_args() { + let (expr, rest) = FilterParser::new(&SCHEME) + .lex_as::(r#"lowercase(http.headers[*])"#) + .unwrap(); + assert_eq!(rest, ""); + + let expr = expr.compile(); + let ctx = &mut ExecutionContext::new(&SCHEME); + let headers = LhsValue::from({ + let mut map = TypedMap::new(); + map.insert(b"0".to_vec().into(), "ONE"); + map.insert(b"1".to_vec().into(), "TWO"); + map.insert(b"2".to_vec().into(), "THREE"); + map + }); + ctx.set_field_value(field("http.headers"), headers).unwrap(); + + assert_eq!( + expr.execute(ctx), + Ok(LhsValue::Array(Array::from_iter(["one", "two", "three"]))) + ); + } + #[test] fn test_map_each_on_array_for_cmp() { let expr = assert_ok!( diff --git a/engine/src/ast/function_expr.rs b/engine/src/ast/function_expr.rs index c0656c91..de5c5178 100644 --- a/engine/src/ast/function_expr.rs +++ b/engine/src/ast/function_expr.rs @@ -2,7 +2,7 @@ use super::ValueExpr; use super::parse::FilterParser; use super::visitor::{Visitor, VisitorMut}; use crate::FunctionRef; -use crate::ast::field_expr::{ComparisonExpr, ComparisonOp, ComparisonOpExpr}; +use crate::ast::field_expr::{ComparisonExpr, ComparisonOp, ComparisonOpExpr, IdentifierExpr}; use crate::ast::index_expr::IndexExpr; use crate::ast::logical_expr::{LogicalExpr, UnaryOp}; use crate::compiler::Compiler; @@ -89,6 +89,20 @@ impl FunctionCallArgExpr { } } + /// Returns `true` if re-evaluating this argument for every mapped element + /// could be costly, i.e. it is a nested function call or a logical + /// sub-expression. Literals and plain field accesses are cheap to evaluate + /// repeatedly, so memoizing them would only add overhead. + fn is_expensive_to_reevaluate(&self) -> bool { + match self { + FunctionCallArgExpr::Literal(_) => false, + FunctionCallArgExpr::Logical(_) => true, + FunctionCallArgExpr::IndexExpr(index_expr) => { + matches!(index_expr.identifier, IdentifierExpr::FunctionCallExpr(_)) + } + } + } + #[allow(dead_code)] pub(crate) fn simplify(self) -> Self { match self { @@ -263,6 +277,16 @@ impl ValueExpr for FunctionCallExpr { let call = function .as_definition() .compile(&mut args.iter().map(|arg| arg.into()), context); + // For `map_each`, only bother evaluating the non-mapped arguments once + // (instead of once per element) when at least one of them is expensive + // to re-evaluate. For trivial arguments (literals / plain field + // accesses) inline re-evaluation is just as cheap and avoids a + // per-call allocation. + let memoize_extra_args = map_each_count > 0 + && args + .iter() + .skip(1) + .any(FunctionCallArgExpr::is_expensive_to_reevaluate); let mut args = args .into_iter() .map(|arg| compiler.compile_function_call_arg_expr(arg)) @@ -278,27 +302,38 @@ impl ValueExpr for FunctionCallExpr { return_type: Type, f: impl Fn(LhsValue<'a>) -> I, ) -> CompiledValueResult<'a> { - let mut first = match first { + let first = match first { Ok(first) => first, Err(_) => { return Err(Type::Array(return_type.into())); } }; - // Extract the values of the map - if let LhsValue::Map(map) = first { - first = LhsValue::Array( - Array::try_from_iter(map.value_type(), map.into_values()).unwrap(), - ); - } - // Retrieve the underlying `Array` - let mut first = match first { - LhsValue::Array(arr) => arr, + let result = match first { + // Map the values straight into the result array. This avoids + // the intermediate `Array` allocation (and per-element type + // re-check) that a separate map -> array conversion followed + // by `filter_map_to` would incur. + LhsValue::Map(map) => { + // Reserve up front for the whole map: `filter_map`'s + // `size_hint` lower bound is 0, and `map_each` rarely + // filters, so this avoids repeated reallocations. + let len = map.len(); + Array::try_from_iter_with_capacity( + return_type, + len, + map.into_values().filter_map(|elem| call(&mut f(elem))), + ) + .unwrap() + } + LhsValue::Array(mut arr) => { + if !arr.is_empty() { + arr = arr.filter_map_to(return_type, |elem| call(&mut f(elem))); + } + arr + } _ => unreachable!(), }; - if !first.is_empty() { - first = first.filter_map_to(return_type, |elem| call(&mut f(elem))); - } - Ok(LhsValue::Array(first)) + Ok(LhsValue::Array(result)) } if args.is_empty() { @@ -311,6 +346,21 @@ impl ValueExpr for FunctionCallExpr { |elem| once(Ok(elem)), ) }) + } else if memoize_extra_args { + CompiledValueExpr::new(move |ctx| { + // At least one non-mapped argument is expensive to + // re-evaluate, so evaluate all of them once per call and + // reuse them (cheaply cloned) for every element instead of + // re-executing the argument expressions for every element. + let extra_args = args.iter().map(|arg| arg.execute(ctx)).collect::>(); + compute( + first.execute(ctx), + &call, + return_type, + #[inline] + |elem| ExactSizeChain::new(once(Ok(elem)), extra_args.iter().cloned()), + ) + }) } else { CompiledValueExpr::new(move |ctx| { compute( diff --git a/engine/src/lhs_types/array.rs b/engine/src/lhs_types/array.rs index b3620f00..86508bbc 100644 --- a/engine/src/lhs_types/array.rs +++ b/engine/src/lhs_types/array.rs @@ -160,22 +160,33 @@ impl<'a> Array<'a> { where F: Fn(LhsValue<'a>) -> Option>, { - let Self { data, .. } = self; - let mut vec = match data { - InnerArray::Owned(vec) => vec, - InnerArray::Borrowed(slice) => slice.to_vec(), - }; let val_type = value_type.into(); - let mut write = 0; - for read in 0..vec.len() { - let elem = &mut vec[read]; - if let Some(elem) = func(std::mem::replace(elem, LhsValue::Bool(false))) { - assert!(elem.get_type() == val_type.into()); - vec[write] = elem; - write += 1; + let Self { data, .. } = self; + let vec = match data { + InnerArray::Owned(mut vec) => { + let mut write = 0; + for read in 0..vec.len() { + let elem = &mut vec[read]; + if let Some(elem) = func(std::mem::replace(elem, LhsValue::Bool(false))) { + assert!(elem.get_type() == val_type.into()); + vec[write] = elem; + write += 1; + } + } + vec.truncate(write); + vec } - } - vec.truncate(write); + InnerArray::Borrowed(slice) => { + let mut vec = Vec::with_capacity(slice.len()); + for elem in slice { + if let Some(elem) = func(elem.as_ref()) { + assert!(elem.get_type() == val_type.into()); + vec.push(elem); + } + } + vec + } + }; Array { val_type, data: InnerArray::Owned(vec), @@ -186,26 +197,36 @@ impl<'a> Array<'a> { pub fn try_from_iter>>( val_type: impl Into, iter: impl IntoIterator, + ) -> Result { + Self::try_from_iter_with_capacity(val_type, 0, iter) + } + + /// Creates a new array from the specified iterator, reserving space for at + /// least `capacity` elements up front. + pub(crate) fn try_from_iter_with_capacity>>( + val_type: impl Into, + capacity: usize, + iter: impl IntoIterator, ) -> Result { let val_type = val_type.into(); - iter.into_iter() - .map(|elem| { - let elem = elem.into(); - let elem_type = elem.get_type(); - if val_type != elem_type.into() { - Err(TypeMismatchError { - expected: Type::from(val_type).into(), - actual: elem_type, - }) - } else { - Ok(elem) - } - }) - .collect::, _>>() - .map(|vec| Array { - val_type, - data: InnerArray::Owned(vec), - }) + let iter = iter.into_iter(); + let capacity = capacity.max(iter.size_hint().0); + let mut vec = Vec::with_capacity(capacity); + for elem in iter { + let elem = elem.into(); + let elem_type = elem.get_type(); + if val_type != elem_type.into() { + return Err(TypeMismatchError { + expected: Type::from(val_type).into(), + actual: elem_type, + }); + } + vec.push(elem); + } + Ok(Array { + val_type, + data: InnerArray::Owned(vec), + }) } /// Creates a new array form the specified vector. diff --git a/ffi/Cargo.toml b/ffi/Cargo.toml index fd588f1c..f1289c24 100644 --- a/ffi/Cargo.toml +++ b/ffi/Cargo.toml @@ -36,5 +36,5 @@ regex-automata.workspace = true [build-dependencies] cbindgen.workspace = true -[target.'cfg(unix)'.dev-dependencies] +[target."cfg(unix)".dev-dependencies] wirefilter-ffi-ctests = { path = "tests/ctests" }