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 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 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}