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