functions: one work budget across every call of a lifted function

A function lifted over a computed array runs once per element, and only the
six lookups counted their repeated work against LIFTED_SCAN_BUDGET. Criteria
aggregates rescanned their ranges per criterion, CUMIPMT/CUMPRINC stepped
through up to 1e7 periods per element, and array-producing functions such as
MUNIT built a whole result per element before keeping its first cell.

The budget now covers all of it: COUNTIF/SUMIF/AVERAGEIF(S)/MAXIFS/MINIFS
charge their range cells per call, the period loops of every element are
added up before any runs, and array results charge their cells as they are
built. Past the budget the call is #NUM!, as lifted lookups already are.

crates/engine/src/cf.rs

@@ -1,7 +1,7 @@
 //! Conditional formatting evaluation: what extra style, data bar, colour scale or icon a cell
 //! gets.
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 
 use gridcraft_core::{CellRef, RangeRef, Value, compare};
 use gridcraft_model::{CfOperator, CfRule, CfValueKind, Color, CondFormat, PatternType, Sheet, Style, Underline, Workbook, default_icon_thresholds};
@@ -70,31 +70,15 @@ thread_local! {
     static SCANNED: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
 }
 
-/// Most cells one rule's statistics visit, across all of its ranges together: a rule repeating
-/// the same range many times must not multiply the scan.
-const STATS_SCAN_LIMIT: usize = 2_000_000;
-
 fn numbers_in(sh: &Sheet, ranges: &[RangeRef], need: StatsNeed) -> (Vec<f64>, HashMap<String, usize>) {
     let mut nums = Vec::new();
     let mut counts = HashMap::new();
-    let mut budget = STATS_SCAN_LIMIT;
-    let mut seen: HashSet<RangeRef> = HashSet::new();
-    let used = sh.used_range();
     for r in ranges {
-        if budget == 0 {
-            break;
-        }
-        let r = match used.and_then(|u| r.intersection(&u)) {
+        let r = match sh.used_range().and_then(|u| r.intersection(&u)) {
             Some(x) => x,
             None => continue,
         };
-        // The same range listed twice counts its cells once.
-        if !seen.insert(r) {
-            continue;
-        }
-        let cells = r.iter().take(budget);
-        for c in cells {
-            budget -= 1;
+        for c in r.iter().take(2_000_000) {
             #[cfg(test)]
             SCANNED.with(|n| n.set(n.get() + 1));
             let v = sh.value(c);
@@ -294,8 +278,7 @@ impl CfCache {
                     _ => None,
                 },
                 CfRule::ColorScale { stops } => {
-                    // Only two- and three-colour scales paint: evaluate no thresholds for others.
-                    if let Some(n) = v.as_f64().filter(|_| matches!(stops.len(), 2 | 3)) {
+                    if let Some(n) = v.as_f64() {
                         let pts: Vec<(f64, [u8; 3])> = stops
                             .iter()
                             .map(|(k, col)| (threshold(wb, sheet, c, st, k), col.resolve(&wb.theme).unwrap_or([255, 255, 255])))
@@ -577,50 +560,6 @@ mod tests {
         assert_eq!(SCANNED.with(|n| n.get()), 10_000);
     }
 
-    /// A rule listing the same range many times scans it once, and never more than the shared
-    /// limit however many ranges it holds.
-    #[test]
-    fn repeated_ranges_share_one_scan_budget() {
-        let mut s = crate::Session::new();
-        s.new_workbook();
-        s.execute("range.setValues", json!({"range": "A1", "values": [[1], [2], [3]]})).unwrap();
-        s.execute("range.setValues", json!({"range": "XFD1048576", "values": [[4]]})).unwrap();
-        let d = s.doc().unwrap();
-        let sh = d.wb.sheet(0).unwrap();
-        let dense = RangeRef::parse("A1:XFD1048576").unwrap();
-        let small = RangeRef::parse("A1:A3").unwrap();
-        SCANNED.with(|n| n.set(0));
-        let st = RuleStats::build(sh, &vec![small; 1000], StatsNeed::Numbers);
-        assert_eq!(st.sorted, vec![1.0, 2.0, 3.0]);
-        assert_eq!(SCANNED.with(|n| n.get()), 3);
-        let shifted: Vec<RangeRef> = (0..50).map(|i| RangeRef::new(CellRef::new(i, 0), dense.end)).collect();
-        SCANNED.with(|n| n.set(0));
-        RuleStats::build(sh, &shifted, StatsNeed::Numbers);
-        assert!(SCANNED.with(|n| n.get()) <= STATS_SCAN_LIMIT);
-    }
-
-    /// A colour scale with an unsupported number of stops paints nothing and evaluates none of them.
-    #[test]
-    fn colour_scales_with_other_stop_counts_paint_nothing() {
-        let mut s = crate::Session::new();
-        s.new_workbook();
-        s.execute("range.setValues", json!({"range": "A1", "values": [[1], [5]]})).unwrap();
-        let mut wb: Workbook = (*s.doc().unwrap().wb).clone();
-        let stop = (CfValueKind::Formula("=1".into()), Color::default());
-        let rule = |n: usize| CondFormat {
-            ranges: vec![RangeRef::parse("A1:A2").unwrap()],
-            rule: CfRule::ColorScale { stops: vec![stop.clone(); n] },
-            priority: 1,
-            stop_if_true: false,
-        };
-        wb.sheet_mut(0).unwrap().cond_formats.push(rule(4));
-        let mut cache = CfCache::new();
-        assert_eq!(cache.look(&wb, 0, CellRef::parse("A1").unwrap()).and_then(|l| l.scale_fill), None);
-        wb.sheet_mut(0).unwrap().cond_formats = vec![rule(2)];
-        let mut cache = CfCache::new();
-        assert!(cache.look(&wb, 0, CellRef::parse("A1").unwrap()).and_then(|l| l.scale_fill).is_some());
-    }
-
     #[test]
     fn cell_value_operator_names() {
         let mut s = crate::Session::new();

crates/functions/src/lift.rs

@@ -68,26 +68,70 @@ const SCANS: &[&str] = &["VLOOKUP", "HLOOKUP", "LOOKUP", "MATCH", "XLOOKUP", "XM
 /// `XMATCH(SEQUENCE(1048576), -SEQUENCE(1048576))` runs for hours and hangs opening the file.
 /// Past the budget the call is `#NUM!`, as MMULT is past 2e9 multiply-adds. 5e8 keeps the worst
 /// case to a few seconds and still admits 5,000 needles over a 100,000-row table.
+///
+/// The same budget covers all the work of a lifted call added up across its elements: criteria
+/// aggregates rescanning their ranges, period loops, and array results built per element.
 const LIFTED_SCAN_BUDGET: usize = 500_000_000;
 
-/// The length of the longest vector a lifted call of `spec` searches per needle: 0 unless `spec`
-/// is one of the [`SCANS`] lookups and a non-lifted argument is an array.
+/// Criteria aggregates: every call tests each cell of its ranges against its criteria, so a call
+/// lifted over `n` criteria does `n` x (range cells) work, like the [`SCANS`] lookups.
+const CRITERIA: &[&str] = &["COUNTIF", "COUNTIFS", "SUMIF", "SUMIFS", "AVERAGEIF", "AVERAGEIFS", "MAXIFS", "MINIFS"];
+
+/// Cumulative TVM functions that step through every period from `start_period` (argument 3) to
+/// `end_period` (argument 4) on each call.
+const PERIOD_LOOPS: &[&str] = &["CUMIPMT", "CUMPRINC"];
+
+/// The work every lifted call of `spec` repeats over its non-lifted arguments: the longest vector
+/// a [`SCANS`] lookup searches, or every cell a [`CRITERIA`] aggregate tests. 0 for other
+/// functions.
 fn scan_len(spec: &FnSpec, args: &[Arg]) -> usize {
-    if !SCANS.contains(&spec.name) {
+    let scans = SCANS.contains(&spec.name);
+    if !scans && !CRITERIA.contains(&spec.name) {
         return 0;
     }
-    args.iter()
-        .enumerate()
-        .filter(|(i, _)| !spec.is_scalar_arg(*i))
-        .filter_map(|(_, a)| match &a.value {
-            Value::Array(arr) => Some(arr.rows.max(arr.cols)),
-            _ => None,
-        })
-        .max()
-        .unwrap_or(0)
+    let lens = args.iter().enumerate().filter(|(i, _)| !spec.is_scalar_arg(*i)).filter_map(|(_, a)| match &a.value {
+        Value::Array(arr) => Some(if scans { arr.rows.max(arr.cols) } else { arr.data.len() }),
+        _ => None,
+    });
+    if scans { lens.max().unwrap_or(0) } else { lens.fold(0, usize::saturating_add) }
+}
+
+/// The work of one lifted call that depends on its own (lifted) scalar arguments: the number of
+/// periods a [`PERIOD_LOOPS`] function steps through. 0 for other functions.
+fn call_cost(spec: &FnSpec, cell_args: &[Arg]) -> usize {
+    if !PERIOD_LOOPS.contains(&spec.name) {
+        return 0;
+    }
+    let period = |i: usize| match cell_args.get(i).map(|a| &a.value) {
+        Some(Value::Number(n)) if n.is_finite() => Some(n.trunc()),
+        _ => None,
+    };
+    match (period(3), period(4)) {
+        // The function itself refuses spans past 1e7 periods; anything larger costs the cap.
+        (Some(start), Some(end)) if end >= start => ((end - start) + 1.0).min(1e8) as usize,
+        _ => 0,
+    }
 }
 
 pub(crate) fn call_lifted(spec: &FnSpec, args: &[Arg], ctx: &mut dyn Ctx) -> Value {
+    call_lifted_within(spec, args, ctx, LIFTED_SCAN_BUDGET)
+}
+
+/// Sets the lifted arguments of `cell_args` to their elements at (`r`, `c`).
+fn element_args(spec: &FnSpec, args: &[Arg], r: usize, c: usize, cell_args: &mut [Arg]) {
+    for (i, a) in args.iter().enumerate() {
+        if spec.is_scalar_arg(i)
+            && let Value::Array(arr) = &a.value
+            && let Some(slot) = cell_args.get_mut(i)
+        {
+            let v = arr.get_broadcast(r, c);
+            slot.value = if let Value::Array(inner) = &v { inner.data.first().cloned().unwrap_or(Value::Empty) } else { v };
+        }
+    }
+}
+
+/// [`call_lifted`] with every call's work added up against `budget` (see [`LIFTED_SCAN_BUDGET`]).
+fn call_lifted_within(spec: &FnSpec, args: &[Arg], ctx: &mut dyn Ctx, budget: usize) -> Value {
     let rejected = if spec.value_args.is_empty() { None } else { reject_multi_cell_refs(spec, args) };
     let args = rejected.as_deref().unwrap_or(args);
     let mut rows = 0usize;
@@ -106,25 +150,42 @@ pub(crate) fn call_lifted(spec: &FnSpec, args: &[Arg], ctx: &mut dyn Ctx) -> Val
         return (spec.imp)(args, ctx);
     }
     let needles = rows.saturating_mul(cols);
-    if needles > MAX_CELLS || needles.saturating_mul(scan_len(spec, args)) > LIFTED_SCAN_BUDGET {
+    // Every call is charged its repeated scan (at least 1), its own period loop, and the cells of
+    // an array result (built in full, though only its first cell is kept), so lifting a scalar
+    // argument can't multiply what one call may do into hours of work:
+    // `COUNTIF(big_range, SEQUENCE(1e4))`, `CUMIPMT(0.1, 1e7, 1, 1, SEQUENCE(1e4,,1e7,0), 0)`,
+    // `MUNIT(SEQUENCE(1e4,,3000,0))`.
+    let per_call = scan_len(spec, args).max(1);
+    let mut cell_args: Vec<Arg> = args.to_vec();
+    let mut spent = needles.saturating_mul(per_call);
+    if needles > MAX_CELLS || spent > budget {
         return Value::Error(CellError::Num);
     }
-    let mut cell_args: Vec<Arg> = args.to_vec();
+    // Period loops are known before any call runs, so a call over budget does none of them.
+    if PERIOD_LOOPS.contains(&spec.name) {
+        for r in 0..rows {
+            for c in 0..cols {
+                element_args(spec, args, r, c, &mut cell_args);
+                spent = spent.saturating_add(call_cost(spec, &cell_args));
+                if spent > budget {
+                    return Value::Error(CellError::Num);
+                }
+            }
+        }
+    }
     let mut data = Vec::with_capacity(rows * cols);
     for r in 0..rows {
         for c in 0..cols {
-            for (i, a) in args.iter().enumerate() {
-                if spec.is_scalar_arg(i)
-                    && let Value::Array(arr) = &a.value
-                    && let Some(slot) = cell_args.get_mut(i)
-                {
-                    let v = arr.get_broadcast(r, c);
-                    slot.value = if let Value::Array(inner) = &v { inner.data.first().cloned().unwrap_or(Value::Empty) } else { v };
-                }
-            }
+            element_args(spec, args, r, c, &mut cell_args);
             let out = (spec.imp)(&cell_args, ctx);
             data.push(match out {
-                Value::Array(a) => a.data.first().cloned().unwrap_or(Value::Empty),
+                Value::Array(a) => {
+                    spent = spent.saturating_add(a.data.len());
+                    if spent > budget {
+                        return Value::Error(CellError::Num);
+                    }
+                    a.data.first().cloned().unwrap_or(Value::Empty)
+                }
                 v => v,
             });
         }
@@ -209,6 +270,52 @@ mod tests {
         assert_eq!(rows_of(&out), vec![vec![nv(-1.0)]; 1000]);
     }
 
+    /// Security: lifting a criteria aggregate rescans its range for every criterion, so criteria
+    /// x range cells past the budget is `#NUM!` before any scan runs.
+    #[test]
+    fn lifted_criteria_aggregates_stop_past_their_work_budget() {
+        use gridcraft_core::Array;
+        let tall = |n: usize| av(Array::filled(n, 1, nv(-1.0)).into());
+        is_err(ev("COUNTIF", vec![tall(1_000_000), tall(501)]), CellError::Num);
+        is_err(ev("SUMIF", vec![tall(1_000_000), tall(501)]), CellError::Num);
+        is_err(ev("COUNTIFS", vec![tall(1_000_000), tall(501)]), CellError::Num);
+        // Two ranges of 250,000 cells each count together.
+        is_err(ev("SUMIFS", vec![tall(250_000), tall(250_000), tall(1001)]), CellError::Num);
+        let counts = ev("COUNTIF", vec![av(col(&[1.0, 2.0, 2.0])), av(col(&[2.0, 1.0, 5.0]))]);
+        assert_eq!(rows_of(&counts), vec![vec![nv(2.0)], vec![nv(1.0)], vec![nv(0.0)]]);
+        for name in super::CRITERIA.iter().chain(super::PERIOD_LOOPS) {
+            assert!(crate::lookup(name).is_some(), "{name} isn't registered");
+        }
+    }
+
+    /// Security: CUMIPMT and CUMPRINC loop over their periods once per lifted element; the
+    /// periods of all elements are added up first, so an over-budget call runs none of them.
+    #[test]
+    fn lifted_period_loops_stop_past_their_work_budget() {
+        use gridcraft_core::Array;
+        let ends = |len: usize, end: f64| av(Array::filled(len, 1, nv(end)).into());
+        let args = |len: usize, end: f64| vec![n(0.01), n(1e7), n(1000.0), n(1.0), ends(len, end), n(0.0)];
+        // 51 elements of 1e7 periods each is past 5e8.
+        is_err(ev("CUMIPMT", args(51, 1e7)), CellError::Num);
+        is_err(ev("CUMPRINC", args(51, 1e7)), CellError::Num);
+        let out = ev("CUMIPMT", args(2, 12.0));
+        assert_eq!(rows_of(&out).len(), 2);
+        assert!(matches!(rows_of(&out).first().and_then(|r| r.first()), Some(Value::Number(x)) if *x < 0.0));
+    }
+
+    /// Security: an array-producing function lifted over a computed array builds each element's
+    /// whole result before keeping one cell; their cells count against the budget.
+    #[test]
+    fn lifted_array_results_count_against_the_budget() {
+        use gridcraft_core::Array;
+        let spec = crate::lookup("MUNIT").unwrap_or_else(|| panic!("MUNIT isn't registered"));
+        let dims = |n: usize| [av(Array::filled(n, 1, nv(10.0)).into())];
+        // Ten 10x10 identity matrices are 1,000 cells.
+        let out = super::call_lifted_within(spec, &dims(10), &mut TestCtx::default(), 1_010);
+        assert_eq!(rows_of(&out), vec![vec![nv(1.0)]; 10]);
+        is_err(super::call_lifted_within(spec, &dims(11), &mut TestCtx::default(), 1_010), CellError::Num);
+    }
+
     #[test]
     fn value_params_reject_multi_cell_references() {
         let s = FnSpec { value_args: &[true, false], ..spec() };