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 trait FrameTransform {
163 type Output;
164
165 fn transform(
166 &mut self,
167 replay_meta: &ReplayMeta,
168 frame: StatsSnapshotFrame,
169 ) -> SubtrActorResult<Self::Output>;
170}
171
172impl<F, T> FrameTransform for F
173where
174 F: FnMut(&ReplayMeta, StatsSnapshotFrame) -> SubtrActorResult<T>,
175{
176 type Output = T;
177
178 fn transform(
179 &mut self,
180 replay_meta: &ReplayMeta,
181 frame: StatsSnapshotFrame,
182 ) -> SubtrActorResult<Self::Output> {
183 self(replay_meta, frame)
184 }
185}
186
187#[derive(Default, Clone, Copy)]
188pub struct IdentityFrameTransform;
189
190impl FrameTransform for IdentityFrameTransform {
191 type Output = StatsSnapshotFrame;
192
193 fn transform(
194 &mut self,
195 _replay_meta: &ReplayMeta,
196 frame: StatsSnapshotFrame,
197 ) -> SubtrActorResult<Self::Output> {
198 Ok(frame)
199 }
200}
201
202pub struct ModuleFrameTransform<F> {
203 transform: F,
204}
205
206impl<F> ModuleFrameTransform<F> {
207 fn new(transform: F) -> Self {
208 Self { transform }
209 }
210}
211
212impl<F, Modules> FrameTransform for ModuleFrameTransform<F>
213where
214 F: FnMut(Map<String, Value>) -> SubtrActorResult<Modules>,
215{
216 type Output = CapturedStatsFrame<Modules>;
217
218 fn transform(
219 &mut self,
220 _replay_meta: &ReplayMeta,
221 frame: StatsSnapshotFrame,
222 ) -> SubtrActorResult<Self::Output> {
223 frame.map_modules(&mut self.transform)
224 }
225}
226
227struct ReplayStatsFrameTransform;
228
229impl FrameTransform for ReplayStatsFrameTransform {
230 type Output = ReplayStatsFrame;
231
232 fn transform(
233 &mut self,
234 replay_meta: &ReplayMeta,
235 frame: StatsSnapshotFrame,
236 ) -> SubtrActorResult<Self::Output> {
237 CapturedStatsData::<StatsSnapshotFrame> {
238 replay_meta: replay_meta.clone(),
239 config: Map::new(),
240 modules: Map::new(),
241 frames: Vec::new(),
242 }
243 .replay_stats_frame(&frame)
244 }
245}
246
247pub struct StatsCollector<T = StatsSnapshotFrame, F = IdentityFrameTransform> {
248 modules: BuiltinModuleSelection,
249 graph: AnalysisGraph,
250 replay_meta: Option<ReplayMeta>,
251 last_replay_meta_player_count: Option<usize>,
252 frame_transform: F,
253 captured_frames: Option<Vec<T>>,
254 sample_mode: SampleMode,
255 last_sample_time: Option<f32>,
256 frame_persistence: StatsFramePersistenceController,
257 last_demolish_count: usize,
258 last_boost_pad_event_count: usize,
259 last_touch_event_count: usize,
260 last_player_stat_event_count: usize,
261 last_goal_event_count: usize,
262 _marker: PhantomData<T>,
263}
264
265impl Default for StatsCollector<StatsSnapshotFrame, IdentityFrameTransform> {
266 fn default() -> Self {
267 Self::new()
268 }
269}
270
271impl StatsCollector<StatsSnapshotFrame, IdentityFrameTransform> {
272 pub fn new() -> Self {
273 Self::with_selection_and_frame_transform(
274 BuiltinModuleSelection::all(),
275 IdentityFrameTransform,
276 )
277 .expect("builtin stats modules should resolve without conflicts")
278 }
279
280 pub fn only_modules<I>(modules: I) -> Self
281 where
282 I: IntoIterator,
283 I::Item: AsRef<str>,
284 {
285 Self::try_only_modules(modules).expect("builtin stats module names should be valid")
286 }
287
288 pub fn try_only_modules<I>(modules: I) -> SubtrActorResult<Self>
289 where
290 I: IntoIterator,
291 I::Item: AsRef<str>,
292 {
293 Self::with_builtin_module_names(modules)
294 }
295
296 pub fn with_builtin_module_names<I, S>(module_names: I) -> SubtrActorResult<Self>
297 where
298 I: IntoIterator<Item = S>,
299 S: AsRef<str>,
300 {
301 Self::with_selection_and_frame_transform(
302 BuiltinModuleSelection::from_names(module_names)?,
303 IdentityFrameTransform,
304 )
305 }
306
307 pub fn get_snapshot_data(self, replay: &boxcars::Replay) -> SubtrActorResult<StatsSnapshotData>
308 where
309 IdentityFrameTransform: FrameTransform<Output = StatsSnapshotFrame>,
310 {
311 self.capture_frames().get_captured_data(replay)
312 }
313
314 pub fn get_legacy_stats_timeline_value(
320 self,
321 replay: &boxcars::Replay,
322 ) -> SubtrActorResult<Value> {
323 serialize_to_json_value(&self.get_legacy_replay_stats_timeline(replay)?)
324 }
325
326 pub fn get_legacy_replay_stats_timeline(
331 self,
332 replay: &boxcars::Replay,
333 ) -> SubtrActorResult<ReplayStatsTimeline> {
334 self.with_frame_transform(ReplayStatsFrameTransform)
335 .capture_frames()
336 .get_captured_data(replay)?
337 .into_legacy_replay_stats_timeline()
338 }
339
340 #[deprecated(
341 note = "use get_legacy_stats_timeline_value for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
342 )]
343 pub fn get_stats_timeline_value(self, replay: &boxcars::Replay) -> SubtrActorResult<Value> {
344 self.get_legacy_stats_timeline_value(replay)
345 }
346
347 #[deprecated(
348 note = "use get_legacy_stats_timeline_value for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
349 )]
350 pub fn get_replay_stats_timeline(
351 self,
352 replay: &boxcars::Replay,
353 ) -> SubtrActorResult<ReplayStatsTimeline> {
354 self.get_legacy_replay_stats_timeline(replay)
355 }
356
357 pub fn into_snapshot_data(self) -> SubtrActorResult<StatsSnapshotData> {
358 self.into_captured_data()
359 }
360
361 pub fn into_legacy_stats_timeline_value(self) -> SubtrActorResult<Value> {
362 self.into_snapshot_data()?.to_legacy_stats_timeline_value()
363 }
364
365 pub fn into_legacy_replay_stats_timeline(self) -> SubtrActorResult<ReplayStatsTimeline> {
366 self.into_snapshot_data()?
367 .into_legacy_replay_stats_timeline()
368 }
369
370 #[deprecated(
371 note = "use into_legacy_stats_timeline_value for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
372 )]
373 pub fn into_stats_timeline_value(self) -> SubtrActorResult<Value> {
374 self.into_legacy_stats_timeline_value()
375 }
376
377 #[deprecated(
378 note = "use into_legacy_replay_stats_timeline for full partial-sum snapshots, or StatsTimelineEventCollector for compact event-backed timelines"
379 )]
380 pub fn into_replay_stats_timeline(self) -> SubtrActorResult<ReplayStatsTimeline> {
381 self.into_legacy_replay_stats_timeline()
382 }
383}
384
385impl<T, F> StatsCollector<T, F> {
386 fn with_selection_and_frame_transform(
387 modules: BuiltinModuleSelection,
388 frame_transform: F,
389 ) -> SubtrActorResult<Self> {
390 Ok(Self {
391 graph: modules.graph()?,
392 modules,
393 replay_meta: None,
394 last_replay_meta_player_count: None,
395 frame_transform,
396 captured_frames: None,
397 sample_mode: SampleMode::Aggregate,
398 last_sample_time: None,
399 frame_persistence: StatsFramePersistenceController::new(StatsFrameResolution::default()),
400 last_demolish_count: 0,
401 last_boost_pad_event_count: 0,
402 last_touch_event_count: 0,
403 last_player_stat_event_count: 0,
404 last_goal_event_count: 0,
405 _marker: PhantomData,
406 })
407 }
408
409 pub fn capture_frames(mut self) -> Self {
410 self.captured_frames = Some(Vec::new());
411 self.sample_mode = SampleMode::Timeline;
412 self
413 }
414
415 pub fn with_frame_transform<U, G>(self, frame_transform: G) -> StatsCollector<U, G> {
416 let StatsCollector {
417 modules,
418 graph,
419 replay_meta,
420 last_replay_meta_player_count,
421 captured_frames,
422 sample_mode,
423 last_sample_time,
424 frame_persistence,
425 last_demolish_count,
426 last_boost_pad_event_count,
427 last_touch_event_count,
428 last_player_stat_event_count,
429 last_goal_event_count,
430 ..
431 } = self;
432 StatsCollector {
433 modules,
434 graph,
435 replay_meta,
436 last_replay_meta_player_count,
437 frame_transform,
438 captured_frames: captured_frames.map(|_| Vec::new()),
439 sample_mode,
440 last_sample_time,
441 frame_persistence,
442 last_demolish_count,
443 last_boost_pad_event_count,
444 last_touch_event_count,
445 last_player_stat_event_count,
446 last_goal_event_count,
447 _marker: PhantomData,
448 }
449 }
450
451 pub fn with_module_transform<Modules, G>(
452 self,
453 transform: G,
454 ) -> StatsCollector<CapturedStatsFrame<Modules>, ModuleFrameTransform<G>>
455 where
456 G: FnMut(Map<String, Value>) -> SubtrActorResult<Modules>,
457 {
458 self.with_frame_transform(ModuleFrameTransform::new(transform))
459 }
460
461 pub fn with_frame_resolution(mut self, resolution: StatsFrameResolution) -> Self {
462 self.frame_persistence = StatsFramePersistenceController::new(resolution);
463 self
464 }
465
466 pub fn get_stats(mut self, replay: &boxcars::Replay) -> SubtrActorResult<CollectedStats>
467 where
468 F: FrameTransform<Output = T>,
469 {
470 self.sample_mode = SampleMode::Aggregate;
471 let mut processor = ReplayProcessor::new(replay)?;
472 processor.process(&mut self)?;
473 if self.replay_meta.is_none() {
474 self.replay_meta = Some(processor.get_replay_meta()?);
475 }
476 self.into_stats()
477 }
478
479 pub fn get_captured_data(
480 mut self,
481 replay: &boxcars::Replay,
482 ) -> SubtrActorResult<CapturedStatsData<T>>
483 where
484 F: FrameTransform<Output = T>,
485 {
486 let mut processor = ReplayProcessor::new(replay)?;
487 processor.process(&mut self)?;
488 if self.replay_meta.is_none() {
489 self.replay_meta = Some(processor.get_replay_meta()?);
490 }
491 self.into_captured_data()
492 }
493
494 pub fn into_stats(self) -> SubtrActorResult<CollectedStats> {
495 let replay_meta = self
496 .replay_meta
497 .ok_or_else(|| SubtrActorError::new(SubtrActorErrorVariant::CouldNotBuildReplayMeta))?;
498 Ok(CollectedStats {
499 replay_meta,
500 modules: self.modules.collected_modules(&self.graph)?,
501 })
502 }
503
504 pub fn into_captured_data(self) -> SubtrActorResult<CapturedStatsData<T>> {
505 let replay_meta = self
506 .replay_meta
507 .ok_or_else(|| SubtrActorError::new(SubtrActorErrorVariant::CouldNotBuildReplayMeta))?;
508 Ok(CapturedStatsData {
509 replay_meta: replay_meta.clone(),
510 config: self.modules.snapshot_config_json(&self.graph)?,
511 modules: self.modules.modules_json(&self.graph)?,
512 frames: self.captured_frames.unwrap_or_default(),
513 })
514 }
515
516 fn capture_frame_snapshot(
517 &mut self,
518 replay_meta: &ReplayMeta,
519 frame: StatsSnapshotFrame,
520 ) -> SubtrActorResult<()>
521 where
522 F: FrameTransform<Output = T>,
523 {
524 if let Some(frames) = &mut self.captured_frames {
525 frames.push(self.frame_transform.transform(replay_meta, frame)?);
526 }
527 Ok(())
528 }
529
530 fn replace_last_frame_snapshot(
531 &mut self,
532 replay_meta: &ReplayMeta,
533 frame: StatsSnapshotFrame,
534 ) -> SubtrActorResult<()>
535 where
536 F: FrameTransform<Output = T>,
537 {
538 if let Some(frames) = &mut self.captured_frames {
539 if let Some(last_frame) = frames.last_mut() {
540 *last_frame = self.frame_transform.transform(replay_meta, frame)?;
541 }
542 }
543 Ok(())
544 }
545
546 fn refresh_replay_meta(&mut self, processor: &ReplayProcessor) -> SubtrActorResult<()> {
547 let player_count = processor.player_count();
548 if self.last_replay_meta_player_count == Some(player_count) {
549 return Ok(());
550 }
551
552 let replay_meta = processor.get_replay_meta()?;
553 self.graph.on_replay_meta(&replay_meta)?;
554 self.replay_meta = Some(replay_meta);
555 self.last_replay_meta_player_count = Some(player_count);
556 Ok(())
557 }
558}
559
560impl<T, F> Collector for StatsCollector<T, F>
561where
562 F: FrameTransform<Output = T>,
563{
564 fn process_frame(
565 &mut self,
566 processor: &ReplayProcessor,
567 _frame: &boxcars::Frame,
568 frame_number: usize,
569 current_time: f32,
570 ) -> SubtrActorResult<TimeAdvance> {
571 self.refresh_replay_meta(processor)?;
572
573 let dt = self
574 .last_sample_time
575 .map(|last_time| (current_time - last_time).max(0.0))
576 .unwrap_or(0.0);
577 let frame_input = match self.sample_mode {
578 SampleMode::Aggregate => FrameInput::aggregate(
579 processor,
580 frame_number,
581 current_time,
582 dt,
583 self.last_demolish_count,
584 self.last_boost_pad_event_count,
585 self.last_touch_event_count,
586 self.last_player_stat_event_count,
587 self.last_goal_event_count,
588 ),
589 SampleMode::Timeline => FrameInput::timeline(processor, frame_number, current_time, dt),
590 };
591 self.graph.evaluate_with_state(&frame_input)?;
592
593 if self.captured_frames.is_some() {
594 let replay_meta = self
595 .replay_meta
596 .as_ref()
597 .expect("replay metadata should be initialized before snapshotting")
598 .clone();
599 if let Some(emitted_dt) = self.frame_persistence.on_frame(frame_number, current_time) {
600 let mut frame = self.modules.snapshot_frame(&self.graph, &replay_meta)?;
601 frame.dt = emitted_dt;
602 self.capture_frame_snapshot(&replay_meta, frame)?;
603 }
604 }
605
606 self.last_sample_time = Some(current_time);
607 if matches!(self.sample_mode, SampleMode::Aggregate) {
608 self.last_demolish_count = processor.demolishes.len();
609 self.last_boost_pad_event_count = processor.boost_pad_events.len();
610 self.last_touch_event_count = processor.touch_events.len();
611 self.last_player_stat_event_count = processor.player_stat_events.len();
612 self.last_goal_event_count = processor.goal_events.len();
613 }
614
615 Ok(TimeAdvance::NextFrame)
616 }
617
618 fn finish_replay(&mut self, _processor: &ReplayProcessor) -> SubtrActorResult<()> {
619 self.graph.finish()?;
620 let Some(replay_meta) = self.replay_meta.as_ref().cloned() else {
621 return Ok(());
622 };
623 let Some(_) = self.graph.state::<FrameInfo>() else {
624 return Ok(());
625 };
626 let mut final_snapshot = self.modules.snapshot_frame(&self.graph, &replay_meta)?;
627 if self.captured_frames.is_some() {
628 match self
629 .frame_persistence
630 .final_frame_action(final_snapshot.frame_number, final_snapshot.time)
631 {
632 Some(FinalStatsFrameAction::Append { dt }) => {
633 final_snapshot.dt = dt;
634 self.capture_frame_snapshot(&replay_meta, final_snapshot)?;
635 }
636 Some(FinalStatsFrameAction::ReplaceLast { dt }) => {
637 final_snapshot.dt = dt;
638 self.replace_last_frame_snapshot(&replay_meta, final_snapshot)?;
639 }
640 None => {}
641 }
642 }
643 Ok(())
644 }
645}