Avoid cloning in Array::filter_map_to (for map_each)

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::<FunctionCallExpr>(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::<FunctionCallExpr>(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::<FunctionCallExpr>(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!(

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::<Vec<_>>();
+                    compute(
+                        first.execute(ctx),
+                        &call,
+                        return_type,
+                        #[inline]
+                        |elem| ExactSizeChain::new(once(Ok(elem)), extra_args.iter().cloned()),
+                    )
+                })
             } else {
                 CompiledValueExpr::new(move |ctx| {
                     compute(

engine/src/lhs_types/array.rs

@@ -160,22 +160,33 @@ impl<'a> Array<'a> {
     where
         F: Fn(LhsValue<'a>) -> Option<LhsValue<'a>>,
     {
-        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<V: Into<LhsValue<'a>>>(
         val_type: impl Into<CompoundType>,
         iter: impl IntoIterator<Item = V>,
+    ) -> Result<Self, TypeMismatchError> {
+        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<V: Into<LhsValue<'a>>>(
+        val_type: impl Into<CompoundType>,
+        capacity: usize,
+        iter: impl IntoIterator<Item = V>,
     ) -> Result<Self, TypeMismatchError> {
         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::<Result<Vec<_>, _>>()
-            .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.

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" }