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