Skip to main content

subtr_actor/collector/stats/
collector.rs

1use std::collections::HashSet;
2use std::marker::PhantomData;
3
4use serde_json::{Map, Value};
5
6use crate::collector::frame_resolution::{
7    FinalStatsFrameAction, StatsFramePersistenceController, StatsFrameResolution,
8};
9use crate::stats::analysis_graph::{graph_with_builtin_analysis_nodes, AnalysisGraph};
10use crate::stats::calculators::ReplayFrameInputBuilder;
11use crate::*;
12
13use super::builtins::{
14    builtin_module_json, builtin_snapshot_config_json, builtin_snapshot_frame_json,
15    builtin_stats_module_names,
16};
17use super::playback::{
18    CapturedStatsData, CapturedStatsFrame, StatsSnapshotData, StatsSnapshotFrame,
19};
20use super::types::{serialize_to_json_value, CollectedStats, CollectedStatsModule};
21
22#[derive(Default)]
23enum SampleMode {
24    #[default]
25    Aggregate,
26    Timeline,
27}
28
29struct BuiltinModuleSelection {
30    module_names: Vec<&'static str>,
31}
32
33impl BuiltinModuleSelection {
34    fn all() -> Self {
35        Self {
36            module_names: builtin_stats_module_names().to_vec(),
37        }
38    }
39
40    fn from_names<I, S>(module_names: I) -> SubtrActorResult<Self>
41    where
42        I: IntoIterator<Item = S>,
43        S: AsRef<str>,
44    {
45        let mut selected = Vec::new();
46        let mut seen = HashSet::new();
47        for module_name in module_names {
48            let module_name = module_name.as_ref();
49            let resolved_name = builtin_stats_module_names()
50                .iter()
51                .copied()
52                .find(|candidate| *candidate == module_name)
53                .ok_or_else(|| {
54                    SubtrActorError::new(SubtrActorErrorVariant::UnknownStatsModuleName(
55                        module_name.to_owned(),
56                    ))
57                })?;
58            if seen.insert(resolved_name) {
59                selected.push(resolved_name);
60            }
61        }
62        Ok(Self {
63            module_names: selected,
64        })
65    }
66
67    fn graph(&self) -> SubtrActorResult<AnalysisGraph> {
68        if self.module_names == builtin_stats_module_names() {
69            return Ok(build_legacy_timeline_graph());
70        }
71        let mut node_names = self.module_names.clone();
72        node_names.push("stats_projection");
73        graph_with_builtin_analysis_nodes(node_names)
74    }
75
76    fn collected_modules(
77        &self,
78        graph: &AnalysisGraph,
79    ) -> SubtrActorResult<Vec<CollectedStatsModule>> {
80        self.module_names
81            .iter()
82            .copied()
83            .map(|module_name| {
84                Ok(CollectedStatsModule {
85                    name: module_name,
86                    value: builtin_module_json(module_name, graph)?,
87                })
88            })
89            .collect()
90    }
91
92    fn modules_json(&self, graph: &AnalysisGraph) -> SubtrActorResult<Map<String, Value>> {
93        let mut modules = Map::new();
94        for module_name in self.module_names.iter().copied() {
95            modules.insert(
96                module_name.to_owned(),
97                builtin_module_json(module_name, graph)?,
98            );
99        }
100        Ok(modules)
101    }
102
103    fn frame_modules_json(
104        &self,
105        graph: &AnalysisGraph,
106        replay_meta: &ReplayMeta,
107    ) -> SubtrActorResult<Map<String, Value>> {
108        let mut modules = Map::new();
109        for module_name in self.module_names.iter().copied() {
110            if let Some(snapshot) = builtin_snapshot_frame_json(module_name, graph, replay_meta)? {
111                modules.insert(module_name.to_owned(), snapshot);
112            }
113            if module_name == "ball_carry" {
114                if let Some(snapshot) =
115                    builtin_snapshot_frame_json("air_dribble", graph, replay_meta)?
116                {
117                    modules.insert("air_dribble".to_owned(), snapshot);
118                }
119            }
120        }
121        Ok(modules)
122    }
123
124    fn snapshot_config_json(&self, graph: &AnalysisGraph) -> SubtrActorResult<Map<String, Value>> {
125        let mut config = Map::new();
126        for module_name in self.module_names.iter().copied() {
127            if let Some(module_config) = builtin_snapshot_config_json(module_name, graph)? {
128                config.insert(module_name.to_owned(), module_config);
129            }
130        }
131        Ok(config)
132    }
133
134    fn snapshot_frame(
135        &self,
136        graph: &AnalysisGraph,
137        replay_meta: &ReplayMeta,
138    ) -> SubtrActorResult<StatsSnapshotFrame> {
139        let frame = graph.state::<FrameInfo>().ok_or_else(|| {
140            SubtrActorError::new(SubtrActorErrorVariant::CallbackError(
141                "missing FrameInfo state while snapshotting stats frame".to_owned(),
142            ))
143        })?;
144        let gameplay = graph.state::<GameplayState>().ok_or_else(|| {
145            SubtrActorError::new(SubtrActorErrorVariant::CallbackError(
146                "missing GameplayState state while snapshotting stats frame".to_owned(),
147            ))
148        })?;
149        let live_play_state = graph.state::<LivePlayState>().cloned().unwrap_or_default();
150        Ok(StatsSnapshotFrame {
151            frame_number: frame.frame_number,
152            time: frame.time,
153            dt: frame.dt,
154            seconds_remaining: frame.seconds_remaining,
155            game_state: gameplay.game_state,
156            ball_has_been_hit: gameplay.ball_has_been_hit,
157            kickoff_countdown_time: gameplay.kickoff_countdown_time,
158            gameplay_phase: live_play_state.gameplay_phase,
159            is_live_play: live_play_state.is_live_play,
160            modules: self.frame_modules_json(graph, replay_meta)?,
161        })
162    }
163}
164
165pub fn builtin_stats_graph_snapshot_json(
166    graph: &AnalysisGraph,
167    replay_meta: Option<&ReplayMeta>,
168) -> SubtrActorResult<Value> {
169    let modules = BuiltinModuleSelection::all();
170    let frame = if let Some(replay_meta) = replay_meta {
171        if graph.state::<FrameInfo>().is_some() && graph.state::<GameplayState>().is_some() {
172            serialize_to_json_value(&modules.snapshot_frame(graph, replay_meta)?)?
173        } else {
174            Value::Null
175        }
176    } else {
177        Value::Null
178    };
179
180    let mut payload = Map::new();
181    payload.insert(
182        "module_names".to_owned(),
183        serialize_to_json_value(&modules.module_names)?,
184    );
185    payload.insert(
186        "config".to_owned(),
187        Value::Object(modules.snapshot_config_json(graph)?),
188    );
189    payload.insert(
190        "modules".to_owned(),
191        Value::Object(modules.modules_json(graph)?),
192    );
193    payload.insert("frame".to_owned(), frame);
194    Ok(Value::Object(payload))
195}
196
197pub trait FrameTransform {
198    type Output;
199
200    fn transform(
201        &mut self,
202        replay_meta: &ReplayMeta,
203        frame: StatsSnapshotFrame,
204    ) -> SubtrActorResult<Self::Output>;
205}
206
207impl<F, T> FrameTransform for F
208where
209    F: FnMut(&ReplayMeta, StatsSnapshotFrame) -> SubtrActorResult<T>,
210{
211    type Output = T;
212
213    fn transform(
214        &mut self,
215        replay_meta: &ReplayMeta,
216        frame: StatsSnapshotFrame,
217    ) -> SubtrActorResult<Self::Output> {
218        self(replay_meta, frame)
219    }
220}
221
222#[derive(Default, Clone, Copy)]
223pub struct IdentityFrameTransform;
224
225impl FrameTransform for IdentityFrameTransform {
226    type Output = StatsSnapshotFrame;
227
228    fn transform(
229        &mut self,
230        _replay_meta: &ReplayMeta,
231        frame: StatsSnapshotFrame,
232    ) -> SubtrActorResult<Self::Output> {
233        Ok(frame)
234    }
235}
236
237pub struct ModuleFrameTransform<F> {
238    transform: F,
239}
240
241impl<F> ModuleFrameTransform<F> {
242    fn new(transform: F) -> Self {
243        Self { transform }
244    }
245}
246
247impl<F, Modules> FrameTransform for ModuleFrameTransform<F>
248where
249    F: FnMut(Map<String, Value>) -> SubtrActorResult<Modules>,
250{
251    type Output = CapturedStatsFrame<Modules>;
252
253    fn transform(
254        &mut self,
255        _replay_meta: &ReplayMeta,
256        frame: StatsSnapshotFrame,
257    ) -> SubtrActorResult<Self::Output> {
258        frame.map_modules(&mut self.transform)
259    }
260}
261
262struct ReplayStatsFrameTransform;
263
264impl FrameTransform for ReplayStatsFrameTransform {
265    type Output = ReplayStatsFrame;
266
267    fn transform(
268        &mut self,
269        replay_meta: &ReplayMeta,
270        frame: StatsSnapshotFrame,
271    ) -> SubtrActorResult<Self::Output> {
272        CapturedStatsData::<StatsSnapshotFrame> {
273            replay_meta: replay_meta.clone(),
274            config: Map::new(),
275            modules: Map::new(),
276            frames: Vec::new(),
277        }
278        .replay_stats_frame(&frame)
279    }
280}
281
282pub struct StatsCollector<T = StatsSnapshotFrame, F = IdentityFrameTransform> {
283    modules: BuiltinModuleSelection,
284    graph: AnalysisGraph,
285    replay_meta: Option<ReplayMeta>,
286    frame_transform: F,
287    captured_frames: Option<Vec<T>>,
288    sample_mode: SampleMode,
289    frame_input_builder: ReplayFrameInputBuilder,
290    last_replay_meta_player_count: Option<usize>,
291    last_sample_time: Option<f32>,
292    frame_persistence: StatsFramePersistenceController,
293    _marker: PhantomData<T>,
294}
295
296impl Default for StatsCollector<StatsSnapshotFrame, IdentityFrameTransform> {
297    fn default() -> Self {
298        Self::new()
299    }
300}
301
302impl StatsCollector<StatsSnapshotFrame, IdentityFrameTransform> {
303    pub fn new() -> Self {
304        Self::with_selection_and_frame_transform(
305            BuiltinModuleSelection::all(),
306            IdentityFrameTransform,
307        )
308        .expect("builtin stats modules should resolve without conflicts")
309    }
310
311    pub fn only_modules<I>(modules: I) -> Self
312    where
313        I: IntoIterator,
314        I::Item: AsRef<str>,
315    {
316        Self::try_only_modules(modules).expect("builtin stats module names should be valid")
317    }
318
319    pub fn try_only_modules<I>(modules: I) -> SubtrActorResult<Self>
320    where
321        I: IntoIterator,
322        I::Item: AsRef<str>,
323    {
324        Self::with_builtin_module_names(modules)
325    }
326
327    pub fn with_builtin_module_names<I, S>(module_names: I) -> SubtrActorResult<Self>
328    where
329        I: IntoIterator<Item = S>,
330        S: AsRef<str>,
331    {
332        Self::with_selection_and_frame_transform(
333            BuiltinModuleSelection::from_names(module_names)?,
334            IdentityFrameTransform,
335        )
336    }
337
338    pub fn get_snapshot_data(self, replay: &boxcars::Replay) -> SubtrActorResult<StatsSnapshotData>
339    where
340        IdentityFrameTransform: FrameTransform<Output = StatsSnapshotFrame>,
341    {
342        self.capture_frames().get_captured_data(replay)
343    }
344
345    /// Collect the legacy full per-frame stats timeline as JSON.
346    ///
347    /// This serializes cumulative team/player partial sums on every captured
348    /// frame. Prefer `StatsTimelineEventCollector` for compact event-backed
349    /// timeline transfer.
350    pub fn get_legacy_stats_timeline_value(
351        self,
352        replay: &boxcars::Replay,
353    ) -> SubtrActorResult<Value> {
354        serialize_to_json_value(&self.get_legacy_replay_stats_timeline(replay)?)
355    }
356
357    /// Collect the legacy full per-frame stats timeline.
358    ///
359    /// This preserves the pre-event-transfer snapshot shape for compatibility
360    /// and parity checks.
361    pub fn get_legacy_replay_stats_timeline(
362        self,
363        replay: &boxcars::Replay,
364    ) -> SubtrActorResult<ReplayStatsTimeline> {
365        self.with_frame_transform(ReplayStatsFrameTransform)
366            .capture_frames()
367            .get_captured_data(replay)?
368            .into_legacy_replay_stats_timeline()
369    }
370
371    #[deprecated(
372        note = "use get_legacy_stats_timeline_value for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
373    )]
374    pub fn get_stats_timeline_value(self, replay: &boxcars::Replay) -> SubtrActorResult<Value> {
375        self.get_legacy_stats_timeline_value(replay)
376    }
377
378    #[deprecated(
379        note = "use get_legacy_stats_timeline_value for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
380    )]
381    pub fn get_replay_stats_timeline(
382        self,
383        replay: &boxcars::Replay,
384    ) -> SubtrActorResult<ReplayStatsTimeline> {
385        self.get_legacy_replay_stats_timeline(replay)
386    }
387
388    pub fn into_snapshot_data(self) -> SubtrActorResult<StatsSnapshotData> {
389        self.into_captured_data()
390    }
391
392    pub fn into_legacy_stats_timeline_value(self) -> SubtrActorResult<Value> {
393        self.into_snapshot_data()?.to_legacy_stats_timeline_value()
394    }
395
396    pub fn into_legacy_replay_stats_timeline(self) -> SubtrActorResult<ReplayStatsTimeline> {
397        self.into_snapshot_data()?
398            .into_legacy_replay_stats_timeline()
399    }
400
401    #[deprecated(
402        note = "use into_legacy_stats_timeline_value for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
403    )]
404    pub fn into_stats_timeline_value(self) -> SubtrActorResult<Value> {
405        self.into_legacy_stats_timeline_value()
406    }
407
408    #[deprecated(
409        note = "use into_legacy_replay_stats_timeline for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
410    )]
411    pub fn into_replay_stats_timeline(self) -> SubtrActorResult<ReplayStatsTimeline> {
412        self.into_legacy_replay_stats_timeline()
413    }
414}
415
416impl<T, F> StatsCollector<T, F> {
417    fn with_selection_and_frame_transform(
418        modules: BuiltinModuleSelection,
419        frame_transform: F,
420    ) -> SubtrActorResult<Self> {
421        Ok(Self {
422            graph: modules.graph()?,
423            modules,
424            replay_meta: None,
425            frame_transform,
426            captured_frames: None,
427            sample_mode: SampleMode::Aggregate,
428            frame_input_builder: ReplayFrameInputBuilder::default(),
429            last_replay_meta_player_count: None,
430            last_sample_time: None,
431            frame_persistence: StatsFramePersistenceController::new(StatsFrameResolution::default()),
432            _marker: PhantomData,
433        })
434    }
435
436    pub fn capture_frames(mut self) -> Self {
437        self.captured_frames = Some(Vec::new());
438        self.sample_mode = SampleMode::Timeline;
439        self
440    }
441
442    pub fn with_frame_transform<U, G>(self, frame_transform: G) -> StatsCollector<U, G> {
443        let StatsCollector {
444            modules,
445            graph,
446            replay_meta,
447            captured_frames,
448            sample_mode,
449            frame_input_builder,
450            last_replay_meta_player_count,
451            last_sample_time,
452            frame_persistence,
453            ..
454        } = self;
455        StatsCollector {
456            modules,
457            graph,
458            replay_meta,
459            frame_transform,
460            captured_frames: captured_frames.map(|_| Vec::new()),
461            sample_mode,
462            frame_input_builder,
463            last_replay_meta_player_count,
464            last_sample_time,
465            frame_persistence,
466            _marker: PhantomData,
467        }
468    }
469
470    pub fn with_module_transform<Modules, G>(
471        self,
472        transform: G,
473    ) -> StatsCollector<CapturedStatsFrame<Modules>, ModuleFrameTransform<G>>
474    where
475        G: FnMut(Map<String, Value>) -> SubtrActorResult<Modules>,
476    {
477        self.with_frame_transform(ModuleFrameTransform::new(transform))
478    }
479
480    pub fn with_frame_resolution(mut self, resolution: StatsFrameResolution) -> Self {
481        self.frame_persistence = StatsFramePersistenceController::new(resolution);
482        self
483    }
484
485    pub fn get_stats(mut self, replay: &boxcars::Replay) -> SubtrActorResult<CollectedStats>
486    where
487        F: FrameTransform<Output = T>,
488    {
489        self.sample_mode = SampleMode::Aggregate;
490        let mut processor = ReplayProcessor::new(replay)?;
491        processor.process(&mut self)?;
492        if self.replay_meta.is_none() {
493            self.replay_meta = Some(processor.get_replay_meta()?);
494        }
495        self.into_stats()
496    }
497
498    pub fn get_captured_data(
499        mut self,
500        replay: &boxcars::Replay,
501    ) -> SubtrActorResult<CapturedStatsData<T>>
502    where
503        F: FrameTransform<Output = T>,
504    {
505        let mut processor = ReplayProcessor::new(replay)?;
506        processor.process(&mut self)?;
507        if self.replay_meta.is_none() {
508            self.replay_meta = Some(processor.get_replay_meta()?);
509        }
510        self.into_captured_data()
511    }
512
513    pub fn into_stats(self) -> SubtrActorResult<CollectedStats> {
514        let replay_meta = self
515            .replay_meta
516            .ok_or_else(|| SubtrActorError::new(SubtrActorErrorVariant::CouldNotBuildReplayMeta))?;
517        Ok(CollectedStats {
518            replay_meta,
519            modules: self.modules.collected_modules(&self.graph)?,
520        })
521    }
522
523    pub fn into_captured_data(self) -> SubtrActorResult<CapturedStatsData<T>> {
524        let replay_meta = self
525            .replay_meta
526            .ok_or_else(|| SubtrActorError::new(SubtrActorErrorVariant::CouldNotBuildReplayMeta))?;
527        Ok(CapturedStatsData {
528            replay_meta: replay_meta.clone(),
529            config: self.modules.snapshot_config_json(&self.graph)?,
530            modules: self.modules.modules_json(&self.graph)?,
531            frames: self.captured_frames.unwrap_or_default(),
532        })
533    }
534
535    fn capture_frame_snapshot(
536        &mut self,
537        replay_meta: &ReplayMeta,
538        frame: StatsSnapshotFrame,
539    ) -> SubtrActorResult<()>
540    where
541        F: FrameTransform<Output = T>,
542    {
543        if let Some(frames) = &mut self.captured_frames {
544            frames.push(self.frame_transform.transform(replay_meta, frame)?);
545        }
546        Ok(())
547    }
548
549    fn replace_last_frame_snapshot(
550        &mut self,
551        replay_meta: &ReplayMeta,
552        frame: StatsSnapshotFrame,
553    ) -> SubtrActorResult<()>
554    where
555        F: FrameTransform<Output = T>,
556    {
557        if let Some(frames) = &mut self.captured_frames {
558            if let Some(last_frame) = frames.last_mut() {
559                *last_frame = self.frame_transform.transform(replay_meta, frame)?;
560            }
561        }
562        Ok(())
563    }
564
565    fn refresh_replay_meta(&mut self, processor: &dyn ProcessorView) -> SubtrActorResult<()> {
566        let player_count = processor.player_count();
567        if self.last_replay_meta_player_count == Some(player_count) {
568            return Ok(());
569        }
570
571        let replay_meta = processor.get_replay_meta()?;
572        self.graph.on_replay_meta(&replay_meta)?;
573        self.replay_meta = Some(replay_meta);
574        self.last_replay_meta_player_count = Some(player_count);
575        Ok(())
576    }
577}
578
579impl<T, F> Collector for StatsCollector<T, F>
580where
581    F: FrameTransform<Output = T>,
582{
583    fn process_frame(
584        &mut self,
585        processor: &dyn ProcessorView,
586        _frame: &boxcars::Frame,
587        frame_number: usize,
588        current_time: f32,
589    ) -> SubtrActorResult<TimeAdvance> {
590        self.refresh_replay_meta(processor)?;
591
592        let dt = self
593            .last_sample_time
594            .map(|last_time| (current_time - last_time).max(0.0))
595            .unwrap_or(0.0);
596        let frame_input = match self.sample_mode {
597            SampleMode::Aggregate => {
598                self.frame_input_builder
599                    .aggregate(processor, frame_number, current_time, dt)
600            }
601            SampleMode::Timeline => {
602                self.frame_input_builder
603                    .timeline(processor, frame_number, current_time, dt)
604            }
605        };
606        self.graph.evaluate_with_state(&frame_input)?;
607
608        if self.captured_frames.is_some() {
609            let replay_meta = self
610                .replay_meta
611                .as_ref()
612                .expect("replay metadata should be initialized before snapshotting")
613                .clone();
614            if let Some(emitted_dt) = self.frame_persistence.on_frame(frame_number, current_time) {
615                let mut frame = self.modules.snapshot_frame(&self.graph, &replay_meta)?;
616                frame.dt = emitted_dt;
617                self.capture_frame_snapshot(&replay_meta, frame)?;
618            }
619        }
620        self.last_sample_time = Some(current_time);
621
622        Ok(TimeAdvance::NextFrame)
623    }
624
625    fn finish_replay(&mut self, _processor: &dyn ProcessorView) -> SubtrActorResult<()> {
626        self.graph.finish()?;
627        let Some(replay_meta) = self.replay_meta.as_ref().cloned() else {
628            return Ok(());
629        };
630        let Some(_) = self.graph.state::<FrameInfo>() else {
631            return Ok(());
632        };
633        let mut final_snapshot = self.modules.snapshot_frame(&self.graph, &replay_meta)?;
634        if self.captured_frames.is_some() {
635            match self
636                .frame_persistence
637                .final_frame_action(final_snapshot.frame_number, final_snapshot.time)
638            {
639                Some(FinalStatsFrameAction::Append { dt }) => {
640                    final_snapshot.dt = dt;
641                    self.capture_frame_snapshot(&replay_meta, final_snapshot)?;
642                }
643                Some(FinalStatsFrameAction::ReplaceLast { dt }) => {
644                    final_snapshot.dt = dt;
645                    self.replace_last_frame_snapshot(&replay_meta, final_snapshot)?;
646                }
647                None => {}
648            }
649        }
650        Ok(())
651    }
652}