perf: batch current rolling means per stock

crates/fidc-core/src/platform_expr_strategy.rs

@@ -718,6 +718,8 @@ struct StockExpressionState {
     stock_volume_ma20: f64,
     stock_volume_ma60: f64,
     stock_volume_ma100: f64,
+    current_close_rolling_means: [f64; 7],
+    current_volume_rolling_means: [f64; 5],
     current_series_end: Option<usize>,
     extra_factors: BTreeMap<String, f64>,
     extra_text_factors: BTreeMap<String, String>,
@@ -771,6 +773,20 @@ impl StockRollingRequirements {
         };
         self.fields.contains(&(field, lookback))
     }
+
+    fn lookbacks<const N: usize>(&self, field: StockRollingField) -> [usize; N] {
+        let mut lookbacks = [0; N];
+        for (slot, (_, lookback)) in self
+            .fields
+            .iter()
+            .filter(|(candidate, _)| *candidate == field)
+            .take(N)
+            .enumerate()
+        {
+            lookbacks[slot] = *lookback;
+        }
+        lookbacks
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -1053,6 +1069,8 @@ pub struct PlatformExprStrategy {
     selection_quote_usage: StockFilterQuoteUsage,
     stock_rolling_requirements: StockRollingRequirements,
     stock_current_rolling_mean_required: bool,
+    stock_current_close_lookbacks: [usize; 7],
+    stock_current_volume_lookbacks: [usize; 5],
     stock_snapshot_field_requirements: StockSnapshotFieldRequirements,
     stock_extra_factors_required: bool,
     stock_extra_factor_identifiers: BTreeSet<String>,
@@ -1331,8 +1349,16 @@ impl PlatformExprStrategy {
         let selection_quote_usage =
             Self::selection_quote_usage_for_config(&config, &normalized_stock_filter_expr);
         let stock_rolling_requirements = Self::stock_rolling_requirements_for_config(&config);
-        let stock_current_rolling_mean_required =
-            Self::stock_current_rolling_mean_required_for_config(&config);
+        let stock_current_rolling_requirements =
+            Self::stock_current_rolling_requirements_for_config(&config);
+        let stock_current_close_lookbacks =
+            stock_current_rolling_requirements.lookbacks(StockRollingField::Close);
+        let stock_current_volume_lookbacks =
+            stock_current_rolling_requirements.lookbacks(StockRollingField::Volume);
+        let stock_current_rolling_mean_required = stock_current_close_lookbacks
+            .iter()
+            .chain(&stock_current_volume_lookbacks)
+            .any(|lookback| *lookback > 0);
         let stock_snapshot_field_requirements =
             Self::stock_snapshot_field_requirements_for_config(&config);
         let stock_extra_factors_required =
@@ -1378,6 +1404,8 @@ impl PlatformExprStrategy {
             selection_quote_usage,
             stock_rolling_requirements,
             stock_current_rolling_mean_required,
+            stock_current_close_lookbacks,
+            stock_current_volume_lookbacks,
             stock_snapshot_field_requirements,
             stock_extra_factors_required,
             stock_extra_factor_identifiers,
@@ -3934,6 +3962,9 @@ impl PlatformExprStrategy {
         field: &str,
         lookback: usize,
     ) -> Option<f64> {
+        if let Some(value) = self.precomputed_stock_current_rolling_mean(stock, field, lookback) {
+            return value;
+        }
         ctx.data
             .market_current_numeric_moving_average_with_end_by_symbol_id(
                 date,
@@ -3945,6 +3976,41 @@ impl PlatformExprStrategy {
             )
     }
 
+    fn precomputed_stock_current_rolling_mean(
+        &self,
+        stock: &StockExpressionState,
+        field: &str,
+        lookback: usize,
+    ) -> Option<Option<f64>> {
+        let (lookbacks, values): (&[usize], &[f64]) =
+            if ["close", "prev_close", "stock_close", "price"]
+                .iter()
+                .any(|candidate| field.eq_ignore_ascii_case(candidate))
+            {
+                (
+                    &self.stock_current_close_lookbacks,
+                    &stock.current_close_rolling_means,
+                )
+            } else if ["volume", "stock_volume"]
+                .iter()
+                .any(|candidate| field.eq_ignore_ascii_case(candidate))
+            {
+                (
+                    &self.stock_current_volume_lookbacks,
+                    &stock.current_volume_rolling_means,
+                )
+            } else {
+                return None;
+            };
+        lookbacks
+            .iter()
+            .position(|candidate| *candidate == lookback && *candidate > 0)
+            .map(|index| {
+                let value = values[index];
+                value.is_finite().then_some(value)
+            })
+    }
+
     fn stock_state_at_time(
         &self,
         ctx: &StrategyContext<'_>,
@@ -4134,6 +4200,16 @@ impl PlatformExprStrategy {
                 &volume_lookbacks,
                 false,
             );
+        let current_rolling_means = ctx
+            .data
+            .market_standard_rolling_means_by_symbol_id_with_calendar_index(
+                date,
+                calendar_index,
+                symbol_id,
+                &self.stock_current_close_lookbacks,
+                &self.stock_current_volume_lookbacks,
+                true,
+            );
         let close_rolling = |index: usize| rolling_means.close[index].unwrap_or(f64::NAN);
         let volume_rolling = |index: usize| rolling_means.volume[index].unwrap_or(f64::NAN);
         let stock_ma_short = close_rolling(0);
@@ -4304,6 +4380,12 @@ impl PlatformExprStrategy {
             stock_volume_ma20,
             stock_volume_ma60,
             stock_volume_ma100,
+            current_close_rolling_means: std::array::from_fn(|index| {
+                current_rolling_means.close[index].unwrap_or(f64::NAN)
+            }),
+            current_volume_rolling_means: std::array::from_fn(|index| {
+                current_rolling_means.volume[index].unwrap_or(f64::NAN)
+            }),
             current_series_end,
             extra_factors,
             extra_text_factors: if self.stock_text_factors_required {
@@ -10045,7 +10127,9 @@ impl PlatformExprStrategy {
         requirements
     }
 
-    fn stock_current_rolling_mean_required_for_config(config: &PlatformExprStrategyConfig) -> bool {
+    fn stock_current_rolling_requirements_for_config(
+        config: &PlatformExprStrategyConfig,
+    ) -> StockRollingRequirements {
         let mut requirements = StockRollingRequirements::default();
         for expr in [
             config.prelude.as_str(),
@@ -10062,11 +10146,8 @@ impl PlatformExprStrategy {
                 &compact,
                 "rolling_mean_current",
             );
-            if !requirements.fields.is_empty() {
-                return true;
-            }
         }
-        false
+        requirements
     }
 
     fn stock_snapshot_field_requirements_for_config(
@@ -14504,6 +14585,11 @@ mod tests {
         assert!(!strategy.stock_rolling_requirements.requires("volume", 5));
         assert!(!strategy.stock_rolling_requirements.requires("volume", 100));
         assert!(strategy.stock_current_rolling_mean_required);
+        assert_eq!(
+            strategy.stock_current_close_lookbacks,
+            [5, 10, 0, 0, 0, 0, 0]
+        );
+        assert_eq!(strategy.stock_current_volume_lookbacks, [5, 100, 0, 0, 0]);
 
         let mut signal_only = PlatformExprStrategyConfig::microcap_rotation();
         signal_only.stock_filter_expr = "true".to_string();
@@ -31076,6 +31162,18 @@ mod tests {
                 .expect("current volume rolling mean"),
             1_000.0
         );
+        assert_eq!(
+            strategy.precomputed_stock_current_rolling_mean(&stock, "close", 5),
+            Some(Some(10.0))
+        );
+        assert_eq!(
+            strategy.precomputed_stock_current_rolling_mean(&stock, "volume", 5),
+            Some(Some(1_000.0))
+        );
+        assert_eq!(
+            strategy.precomputed_stock_current_rolling_mean(&stock, "close", 7),
+            None
+        );
 
         let mut cfg = PlatformExprStrategyConfig::microcap_rotation();
         cfg.stock_filter_expr = "stock_ma5 > 0 && stock_volume_ma5 > 0".to_string();