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 graph_with_builtin_analysis_nodes(self.module_names.iter().copied())
68 }
69
70 fn collected_modules(
71 &self,
72 graph: &AnalysisGraph,
73 ) -> SubtrActorResult<Vec<CollectedStatsModule>> {
74 self.module_names
75 .iter()
76 .copied()
77 .map(|module_name| {
78 Ok(CollectedStatsModule {
79 name: module_name,
80 value: builtin_module_json(module_name, graph)?,
81 })
82 })
83 .collect()
84 }
85
86 fn modules_json(&self, graph: &AnalysisGraph) -> SubtrActorResult<Map<String, Value>> {
87 let mut modules = Map::new();
88 for module_name in self.module_names.iter().copied() {
89 modules.insert(
90 module_name.to_owned(),
91 builtin_module_json(module_name, graph)?,
92 );
93 }
94 Ok(modules)
95 }
96
97 fn frame_modules_json(
98 &self,
99 graph: &AnalysisGraph,
100 replay_meta: &ReplayMeta,
101 ) -> SubtrActorResult<Map<String, Value>> {
102 let mut modules = Map::new();
103 for module_name in self.module_names.iter().copied() {
104 if let Some(snapshot) = builtin_snapshot_frame_json(module_name, graph, replay_meta)? {
105 modules.insert(module_name.to_owned(), snapshot);
106 }
107 if module_name == "ball_carry" {
108 if let Some(snapshot) =
109 builtin_snapshot_frame_json("air_dribble", graph, replay_meta)?
110 {
111 modules.insert("air_dribble".to_owned(), snapshot);
112 }
113 }
114 }
115 Ok(modules)
116 }
117
118 fn snapshot_config_json(&self, graph: &AnalysisGraph) -> SubtrActorResult<Map<String, Value>> {
119 let mut config = Map::new();
120 for module_name in self.module_names.iter().copied() {
121 if let Some(module_config) = builtin_snapshot_config_json(module_name, graph)? {
122 config.insert(module_name.to_owned(), module_config);
123 }
124 }
125 Ok(config)
126 }
127
128 fn snapshot_frame(
129 &self,
130 graph: &AnalysisGraph,
131 replay_meta: &ReplayMeta,
132 ) -> SubtrActorResult<StatsSnapshotFrame> {
133 let frame = graph.state::<FrameInfo>().ok_or_else(|| {
134 SubtrActorError::new(SubtrActorErrorVariant::CallbackError(
135 "missing FrameInfo state while snapshotting stats frame".to_owned(),
136 ))
137 })?;
138 let gameplay = graph.state::<GameplayState>().ok_or_else(|| {
139 SubtrActorError::new(SubtrActorErrorVariant::CallbackError(
140 "missing GameplayState state while snapshotting stats frame".to_owned(),
141 ))
142 })?;
143 let live_play_state = graph.state::<LivePlayState>().cloned().unwrap_or_default();
144 Ok(StatsSnapshotFrame {
145 frame_number: frame.frame_number,
146 time: frame.time,
147 dt: frame.dt,
148 seconds_remaining: frame.seconds_remaining,
149 game_state: gameplay.game_state,
150 gameplay_phase: live_play_state.gameplay_phase,
151 is_live_play: live_play_state.is_live_play,
152 modules: self.frame_modules_json(graph, replay_meta)?,
153 })
154 }
155}
156
157pub trait FrameTransform {
158 type Output;
159
160 fn transform(
161 &mut self,
162 replay_meta: &ReplayMeta,
163 frame: StatsSnapshotFrame,
164 ) -> SubtrActorResult<Self::Output>;
165}
166
167impl<F, T> FrameTransform for F
168where
169 F: FnMut(&ReplayMeta, StatsSnapshotFrame) -> SubtrActorResult<T>,
170{
171 type Output = T;
172
173 fn transform(
174 &mut self,
175 replay_meta: &ReplayMeta,
176 frame: StatsSnapshotFrame,
177 ) -> SubtrActorResult<Self::Output> {
178 self(replay_meta, frame)
179 }
180}
181
182#[derive(Default, Clone, Copy)]
183pub struct IdentityFrameTransform;
184
185impl FrameTransform for IdentityFrameTransform {
186 type Output = StatsSnapshotFrame;
187
188 fn transform(
189 &mut self,
190 _replay_meta: &ReplayMeta,
191 frame: StatsSnapshotFrame,
192 ) -> SubtrActorResult<Self::Output> {
193 Ok(frame)
194 }
195}
196
197pub struct ModuleFrameTransform<F> {
198 transform: F,
199}
200
201impl<F> ModuleFrameTransform<F> {
202 fn new(transform: F) -> Self {
203 Self { transform }
204 }
205}
206
207impl<F, Modules> FrameTransform for ModuleFrameTransform<F>
208where
209 F: FnMut(Map<String, Value>) -> SubtrActorResult<Modules>,
210{
211 type Output = CapturedStatsFrame<Modules>;
212
213 fn transform(
214 &mut self,
215 _replay_meta: &ReplayMeta,
216 frame: StatsSnapshotFrame,
217 ) -> SubtrActorResult<Self::Output> {
218 frame.map_modules(&mut self.transform)
219 }
220}
221
222struct ReplayStatsFrameTransform;
223
224impl FrameTransform for ReplayStatsFrameTransform {
225 type Output = ReplayStatsFrame;
226
227 fn transform(
228 &mut self,
229 replay_meta: &ReplayMeta,
230 frame: StatsSnapshotFrame,
231 ) -> SubtrActorResult<Self::Output> {
232 CapturedStatsData::<StatsSnapshotFrame> {
233 replay_meta: replay_meta.clone(),
234 config: Map::new(),
235 modules: Map::new(),
236 frames: Vec::new(),
237 }
238 .replay_stats_frame(&frame)
239 }
240}
241
242pub struct StatsCollector<T = StatsSnapshotFrame, F = IdentityFrameTransform> {
243 modules: BuiltinModuleSelection,
244 graph: AnalysisGraph,
245 replay_meta: Option<ReplayMeta>,
246 last_replay_meta_player_count: Option<usize>,
247 frame_transform: F,
248 captured_frames: Option<Vec<T>>,
249 sample_mode: SampleMode,
250 last_sample_time: Option<f32>,
251 frame_persistence: StatsFramePersistenceController,
252 last_demolish_count: usize,
253 last_boost_pad_event_count: usize,
254 last_touch_event_count: usize,
255 last_player_stat_event_count: usize,
256 last_goal_event_count: usize,
257 _marker: PhantomData<T>,
258}
259
260impl Default for StatsCollector<StatsSnapshotFrame, IdentityFrameTransform> {
261 fn default() -> Self {
262 Self::new()
263 }
264}
265
266impl StatsCollector<StatsSnapshotFrame, IdentityFrameTransform> {
267 pub fn new() -> Self {
268 Self::with_selection_and_frame_transform(
269 BuiltinModuleSelection::all(),
270 IdentityFrameTransform,
271 )
272 .expect("builtin stats modules should resolve without conflicts")
273 }
274
275 pub fn only_modules<I>(modules: I) -> Self
276 where
277 I: IntoIterator,
278 I::Item: AsRef<str>,
279 {
280 Self::try_only_modules(modules).expect("builtin stats module names should be valid")
281 }
282
283 pub fn try_only_modules<I>(modules: I) -> SubtrActorResult<Self>
284 where
285 I: IntoIterator,
286 I::Item: AsRef<str>,
287 {
288 Self::with_builtin_module_names(modules)
289 }
290
291 pub fn with_builtin_module_names<I, S>(module_names: I) -> SubtrActorResult<Self>
292 where
293 I: IntoIterator<Item = S>,
294 S: AsRef<str>,
295 {
296 Self::with_selection_and_frame_transform(
297 BuiltinModuleSelection::from_names(module_names)?,
298 IdentityFrameTransform,
299 )
300 }
301
302 pub fn get_snapshot_data(self, replay: &boxcars::Replay) -> SubtrActorResult<StatsSnapshotData>
303 where
304 IdentityFrameTransform: FrameTransform<Output = StatsSnapshotFrame>,
305 {
306 self.capture_frames().get_captured_data(replay)
307 }
308
309 pub fn get_stats_timeline_value(self, replay: &boxcars::Replay) -> SubtrActorResult<Value> {
310 serialize_to_json_value(&self.get_replay_stats_timeline(replay)?)
311 }
312
313 pub fn get_replay_stats_timeline(
314 self,
315 replay: &boxcars::Replay,
316 ) -> SubtrActorResult<ReplayStatsTimeline> {
317 self.with_frame_transform(ReplayStatsFrameTransform)
318 .capture_frames()
319 .get_captured_data(replay)?
320 .into_replay_stats_timeline()
321 }
322
323 pub fn into_snapshot_data(self) -> SubtrActorResult<StatsSnapshotData> {
324 self.into_captured_data()
325 }
326
327 pub fn into_stats_timeline_value(self) -> SubtrActorResult<Value> {
328 self.into_snapshot_data()?.to_stats_timeline_value()
329 }
330
331 pub fn into_replay_stats_timeline(self) -> SubtrActorResult<ReplayStatsTimeline> {
332 self.into_snapshot_data()?.into_stats_timeline()
333 }
334}
335
336impl<T, F> StatsCollector<T, F> {
337 fn with_selection_and_frame_transform(
338 modules: BuiltinModuleSelection,
339 frame_transform: F,
340 ) -> SubtrActorResult<Self> {
341 Ok(Self {
342 graph: modules.graph()?,
343 modules,
344 replay_meta: None,
345 last_replay_meta_player_count: None,
346 frame_transform,
347 captured_frames: None,
348 sample_mode: SampleMode::Aggregate,
349 last_sample_time: None,
350 frame_persistence: StatsFramePersistenceController::new(StatsFrameResolution::default()),
351 last_demolish_count: 0,
352 last_boost_pad_event_count: 0,
353 last_touch_event_count: 0,
354 last_player_stat_event_count: 0,
355 last_goal_event_count: 0,
356 _marker: PhantomData,
357 })
358 }
359
360 pub fn capture_frames(mut self) -> Self {
361 self.captured_frames = Some(Vec::new());
362 self.sample_mode = SampleMode::Timeline;
363 self
364 }
365
366 pub fn with_frame_transform<U, G>(self, frame_transform: G) -> StatsCollector<U, G> {
367 let StatsCollector {
368 modules,
369 graph,
370 replay_meta,
371 last_replay_meta_player_count,
372 captured_frames,
373 sample_mode,
374 last_sample_time,
375 frame_persistence,
376 last_demolish_count,
377 last_boost_pad_event_count,
378 last_touch_event_count,
379 last_player_stat_event_count,
380 last_goal_event_count,
381 ..
382 } = self;
383 StatsCollector {
384 modules,
385 graph,
386 replay_meta,
387 last_replay_meta_player_count,
388 frame_transform,
389 captured_frames: captured_frames.map(|_| Vec::new()),
390 sample_mode,
391 last_sample_time,
392 frame_persistence,
393 last_demolish_count,
394 last_boost_pad_event_count,
395 last_touch_event_count,
396 last_player_stat_event_count,
397 last_goal_event_count,
398 _marker: PhantomData,
399 }
400 }
401
402 pub fn with_module_transform<Modules, G>(
403 self,
404 transform: G,
405 ) -> StatsCollector<CapturedStatsFrame<Modules>, ModuleFrameTransform<G>>
406 where
407 G: FnMut(Map<String, Value>) -> SubtrActorResult<Modules>,
408 {
409 self.with_frame_transform(ModuleFrameTransform::new(transform))
410 }
411
412 pub fn with_frame_resolution(mut self, resolution: StatsFrameResolution) -> Self {
413 self.frame_persistence = StatsFramePersistenceController::new(resolution);
414 self
415 }
416
417 pub fn get_stats(mut self, replay: &boxcars::Replay) -> SubtrActorResult<CollectedStats>
418 where
419 F: FrameTransform<Output = T>,
420 {
421 self.sample_mode = SampleMode::Aggregate;
422 let mut processor = ReplayProcessor::new(replay)?;
423 processor.process(&mut self)?;
424 if self.replay_meta.is_none() {
425 self.replay_meta = Some(processor.get_replay_meta()?);
426 }
427 self.into_stats()
428 }
429
430 pub fn get_captured_data(
431 mut self,
432 replay: &boxcars::Replay,
433 ) -> SubtrActorResult<CapturedStatsData<T>>
434 where
435 F: FrameTransform<Output = T>,
436 {
437 let mut processor = ReplayProcessor::new(replay)?;
438 processor.process(&mut self)?;
439 if self.replay_meta.is_none() {
440 self.replay_meta = Some(processor.get_replay_meta()?);
441 }
442 self.into_captured_data()
443 }
444
445 pub fn into_stats(self) -> SubtrActorResult<CollectedStats> {
446 let replay_meta = self
447 .replay_meta
448 .ok_or_else(|| SubtrActorError::new(SubtrActorErrorVariant::CouldNotBuildReplayMeta))?;
449 Ok(CollectedStats {
450 replay_meta,
451 modules: self.modules.collected_modules(&self.graph)?,
452 })
453 }
454
455 pub fn into_captured_data(self) -> SubtrActorResult<CapturedStatsData<T>> {
456 let replay_meta = self
457 .replay_meta
458 .ok_or_else(|| SubtrActorError::new(SubtrActorErrorVariant::CouldNotBuildReplayMeta))?;
459 Ok(CapturedStatsData {
460 replay_meta: replay_meta.clone(),
461 config: self.modules.snapshot_config_json(&self.graph)?,
462 modules: self.modules.modules_json(&self.graph)?,
463 frames: self.captured_frames.unwrap_or_default(),
464 })
465 }
466
467 fn capture_frame_snapshot(
468 &mut self,
469 replay_meta: &ReplayMeta,
470 frame: StatsSnapshotFrame,
471 ) -> SubtrActorResult<()>
472 where
473 F: FrameTransform<Output = T>,
474 {
475 if let Some(frames) = &mut self.captured_frames {
476 frames.push(self.frame_transform.transform(replay_meta, frame)?);
477 }
478 Ok(())
479 }
480
481 fn replace_last_frame_snapshot(
482 &mut self,
483 replay_meta: &ReplayMeta,
484 frame: StatsSnapshotFrame,
485 ) -> SubtrActorResult<()>
486 where
487 F: FrameTransform<Output = T>,
488 {
489 if let Some(frames) = &mut self.captured_frames {
490 if let Some(last_frame) = frames.last_mut() {
491 *last_frame = self.frame_transform.transform(replay_meta, frame)?;
492 }
493 }
494 Ok(())
495 }
496
497 fn refresh_replay_meta(&mut self, processor: &ReplayProcessor) -> SubtrActorResult<()> {
498 let player_count = processor.player_count();
499 if self.last_replay_meta_player_count == Some(player_count) {
500 return Ok(());
501 }
502
503 let replay_meta = processor.get_replay_meta()?;
504 self.graph.on_replay_meta(&replay_meta)?;
505 self.replay_meta = Some(replay_meta);
506 self.last_replay_meta_player_count = Some(player_count);
507 Ok(())
508 }
509}
510
511impl<T, F> Collector for StatsCollector<T, F>
512where
513 F: FrameTransform<Output = T>,
514{
515 fn process_frame(
516 &mut self,
517 processor: &ReplayProcessor,
518 _frame: &boxcars::Frame,
519 frame_number: usize,
520 current_time: f32,
521 ) -> SubtrActorResult<TimeAdvance> {
522 self.refresh_replay_meta(processor)?;
523
524 let dt = self
525 .last_sample_time
526 .map(|last_time| (current_time - last_time).max(0.0))
527 .unwrap_or(0.0);
528 let frame_input = match self.sample_mode {
529 SampleMode::Aggregate => FrameInput::aggregate(
530 processor,
531 frame_number,
532 current_time,
533 dt,
534 self.last_demolish_count,
535 self.last_boost_pad_event_count,
536 self.last_touch_event_count,
537 self.last_player_stat_event_count,
538 self.last_goal_event_count,
539 ),
540 SampleMode::Timeline => FrameInput::timeline(processor, frame_number, current_time, dt),
541 };
542 self.graph.evaluate_with_state(&frame_input)?;
543
544 if self.captured_frames.is_some() {
545 let replay_meta = self
546 .replay_meta
547 .as_ref()
548 .expect("replay metadata should be initialized before snapshotting")
549 .clone();
550 if let Some(emitted_dt) = self.frame_persistence.on_frame(frame_number, current_time) {
551 let mut frame = self.modules.snapshot_frame(&self.graph, &replay_meta)?;
552 frame.dt = emitted_dt;
553 self.capture_frame_snapshot(&replay_meta, frame)?;
554 }
555 }
556
557 self.last_sample_time = Some(current_time);
558 if matches!(self.sample_mode, SampleMode::Aggregate) {
559 self.last_demolish_count = processor.demolishes.len();
560 self.last_boost_pad_event_count = processor.boost_pad_events.len();
561 self.last_touch_event_count = processor.touch_events.len();
562 self.last_player_stat_event_count = processor.player_stat_events.len();
563 self.last_goal_event_count = processor.goal_events.len();
564 }
565
566 Ok(TimeAdvance::NextFrame)
567 }
568
569 fn finish_replay(&mut self, _processor: &ReplayProcessor) -> SubtrActorResult<()> {
570 self.graph.finish()?;
571 let Some(replay_meta) = self.replay_meta.as_ref().cloned() else {
572 return Ok(());
573 };
574 let Some(_) = self.graph.state::<FrameInfo>() else {
575 return Ok(());
576 };
577 let mut final_snapshot = self.modules.snapshot_frame(&self.graph, &replay_meta)?;
578 if self.captured_frames.is_some() {
579 match self
580 .frame_persistence
581 .final_frame_action(final_snapshot.frame_number, final_snapshot.time)
582 {
583 Some(FinalStatsFrameAction::Append { dt }) => {
584 final_snapshot.dt = dt;
585 self.capture_frame_snapshot(&replay_meta, final_snapshot)?;
586 }
587 Some(FinalStatsFrameAction::ReplaceLast { dt }) => {
588 final_snapshot.dt = dt;
589 self.replace_last_frame_snapshot(&replay_meta, final_snapshot)?;
590 }
591 None => {}
592 }
593 }
594 Ok(())
595 }
596}