Pre-compute map_each args for expensive functions

Author: utkarshgupta137

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))
@@ -322,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(