henad_core/export/
stats_csv.rs

1//! Stat series written as comma-separated values (CSV), one row per sampled tick.
2//!
3//! The column layout is fixed from the *first* sample and reused for every later row, so the
4//! header and every row always agree. A model whose `stats()` shape changes mid-run is a
5//! programming error, and [`StatsWriter::push`] reports it as one.
6
7use std::fmt;
8use std::io::{self, Write};
9
10use crate::export::csv::{escape_field, fmt_f64};
11use crate::view::{StatEntry, StatValue};
12
13/// A stat series could not be written.
14#[derive(Debug)]
15pub enum StatsWriteError {
16    /// Error of the underlying writer.
17    Io(io::Error),
18    /// A series changed shape after the first sample fixed the column layout.
19    Shape(String),
20}
21
22impl fmt::Display for StatsWriteError {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        match self {
25            Self::Io(err) => write!(f, "{err}"),
26            Self::Shape(message) => f.write_str(message),
27        }
28    }
29}
30
31impl std::error::Error for StatsWriteError {
32    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
33        match self {
34            Self::Io(err) => Some(err),
35            Self::Shape(_) => None,
36        }
37    }
38}
39
40impl From<io::Error> for StatsWriteError {
41    fn from(err: io::Error) -> Self {
42        Self::Io(err)
43    }
44}
45
46/// Separator between a series label and a component suffix, e.g. `Average Velocity.x`.
47const SUFFIX_SEP: char = '.';
48
49/// One output column, and the part of the stat series feeding it.
50#[derive(Debug, Clone, PartialEq, Eq)]
51struct Column {
52    /// Column name before CSV escaping.
53    name: String,
54    /// Header text, already escaped for CSV.
55    header: String,
56    /// Label of the series feeding the column.
57    label: &'static str,
58    /// Index into the `Vec<StatEntry>` returned by `stats()`.
59    series: usize,
60    part: Part,
61}
62
63impl Column {
64    /// Returns this column's value in `stats`, the sample taken at `tick`.
65    fn value(&self, tick: u64, stats: &[StatEntry]) -> Result<f64, StatsWriteError> {
66        let Some(entry) = stats.get(self.series) else {
67            return Err(StatsWriteError::Shape(format!(
68                "stat series count changed mid-run: column '{}' needs series {} but tick {tick} has {}",
69                self.header,
70                self.series,
71                stats.len()
72            )));
73        };
74        part_value(&entry.value, self.part).ok_or_else(|| {
75            StatsWriteError::Shape(format!(
76                "stat series '{}' changed shape mid-run at tick {tick}: column '{}' no longer applies",
77                entry.label, self.header
78            ))
79        })
80    }
81}
82
83/// The scalar pulled out of a [`StatValue`] for one column.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85enum Part {
86    Scalar,
87    VecX,
88    VecY,
89    /// Magnitude of a `Vector2D`, so a vector series is still usable without recombining lanes.
90    VecMagnitude,
91    /// Count in one histogram bucket.
92    Bucket(usize),
93    /// Total across all histogram buckets.
94    BucketTotal,
95}
96
97/// Column layout of a stat series, planned from one sample.
98///
99/// A scalar series is one column, a vector is three (`.x`, `.y` and `.magnitude`), and a histogram
100/// is one per bucket plus `.total`.
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct StatColumns {
103    columns: Vec<Column>,
104}
105
106impl StatColumns {
107    /// Derives the column layout from the sample `stats`.
108    pub fn plan(stats: &[StatEntry]) -> Self {
109        let mut columns = Vec::new();
110        for (series, entry) in stats.iter().enumerate() {
111            let mut push = |suffix: Option<&str>, part: Part| {
112                let name = match suffix {
113                    Some(suffix) => format!("{}{SUFFIX_SEP}{suffix}", entry.label),
114                    None => entry.label.to_owned(),
115                };
116                let header = escape_field(&name);
117                columns.push(Column {
118                    name,
119                    header,
120                    label: entry.label,
121                    series,
122                    part,
123                });
124            };
125            match &entry.value {
126                StatValue::Scalar(_) => push(None, Part::Scalar),
127                StatValue::Vector2D { .. } => {
128                    push(Some("x"), Part::VecX);
129                    push(Some("y"), Part::VecY);
130                    push(Some("magnitude"), Part::VecMagnitude);
131                }
132                StatValue::Histogram { edges, counts } => {
133                    // Each bucket is named by its own range, and the columns can be read without the edge list.
134                    // `edges` holds the boundaries, a lower and an upper edge per bucket. A bucket whose
135                    // edges are missing is named by its index.
136                    for bucket in 0..counts.len() {
137                        let range = match (edges.get(bucket), edges.get(bucket + 1)) {
138                            (Some(lo), Some(hi)) => format!("[{}, {})", fmt_f64(*lo), fmt_f64(*hi)),
139                            _ => format!("bucket {bucket}"),
140                        };
141                        push(Some(&range), Part::Bucket(bucket));
142                    }
143                    push(Some("total"), Part::BucketTotal);
144                }
145            }
146        }
147        Self { columns }
148    }
149
150    /// Number of columns.
151    pub fn len(&self) -> usize {
152        self.columns.len()
153    }
154
155    /// Returns whether the layout has no columns.
156    pub fn is_empty(&self) -> bool {
157        self.columns.is_empty()
158    }
159
160    /// Returns the name of column `i`, before CSV escaping.
161    ///
162    /// # Panics
163    ///
164    /// Panics when `i` is not below [`Self::len`].
165    pub fn name(&self, i: usize) -> &str {
166        &self.columns[i].name
167    }
168
169    /// Returns the header of column `i`, escaped for CSV.
170    ///
171    /// # Panics
172    ///
173    /// Panics when `i` is not below [`Self::len`].
174    pub fn header(&self, i: usize) -> &str {
175        &self.columns[i].header
176    }
177
178    /// Returns whether column `i` counts a single histogram bucket.
179    ///
180    /// # Panics
181    ///
182    /// Panics when `i` is not below [`Self::len`].
183    pub fn is_bucket(&self, i: usize) -> bool {
184        matches!(self.columns[i].part, Part::Bucket(_))
185    }
186
187    /// Returns the index of the column called `name`.
188    ///
189    /// A bare vector or histogram label resolves to its magnitude or total column, the value
190    /// [`StatValue::scalar`] returns. An exact column name wins over a bare label.
191    pub fn resolve(&self, name: &str) -> Option<usize> {
192        self.columns.iter().position(|column| column.name == name).or_else(|| {
193            self.columns.iter().position(|column| {
194                column.label == name && matches!(column.part, Part::VecMagnitude | Part::BucketTotal)
195            })
196        })
197    }
198
199    /// Replaces the contents of `out` with one value per column, read from the sample `stats` taken at `tick`.
200    ///
201    /// # Errors
202    ///
203    /// Returns [`StatsWriteError::Shape`] when `stats` no longer fits the planned layout.
204    pub fn extract(&self, tick: u64, stats: &[StatEntry], out: &mut Vec<f64>) -> Result<(), StatsWriteError> {
205        out.clear();
206        for column in &self.columns {
207            out.push(column.value(tick, stats)?);
208        }
209        Ok(())
210    }
211}
212
213/// Streams a stat time series to a writer as CSV.
214///
215/// Call [`push`](Self::push) once per sampled tick, then [`finish`](Self::finish). The first `push`
216/// writes the header, with the columns of its sample. A [`StatDescriptor`](crate::view::StatDescriptor)
217/// does not say whether its value is a scalar, a vector or a histogram.
218#[derive(Debug)]
219pub struct StatsWriter<W: Write> {
220    out: W,
221    /// Column layout, `None` until the first `push` fixes it.
222    columns: Option<StatColumns>,
223    rows: u64,
224}
225
226impl<W: Write> StatsWriter<W> {
227    /// Returns a writer over `out` that has written nothing.
228    pub fn new(out: W) -> Self {
229        Self {
230            out,
231            columns: None,
232            rows: 0,
233        }
234    }
235
236    /// Writes the row of the sample `stats`, taken at `tick`.
237    ///
238    /// The first call fixes the column layout and writes the header.
239    ///
240    /// # Errors
241    ///
242    /// Returns [`StatsWriteError::Io`] when a write fails, and [`StatsWriteError::Shape`] when `stats` does not fit
243    /// the layout the first sample fixed.
244    pub fn push(&mut self, tick: u64, stats: &[StatEntry]) -> Result<(), StatsWriteError> {
245        if self.columns.is_none() {
246            let columns = StatColumns::plan(stats);
247            write!(self.out, "tick")?;
248            for column in &columns.columns {
249                write!(self.out, ",{}", column.header)?;
250            }
251            writeln!(self.out)?;
252            self.columns = Some(columns);
253        }
254        let columns = self.columns.as_ref().map_or(&[][..], |planned| &planned.columns[..]);
255
256        write!(self.out, "{tick}")?;
257        for column in columns {
258            let value = column.value(tick, stats)?;
259            write!(self.out, ",{}", fmt_f64(value))?;
260        }
261        writeln!(self.out)?;
262        self.rows += 1;
263        Ok(())
264    }
265
266    /// Flushes the writer and returns the number of rows written.
267    ///
268    /// Note that a `BufWriter` dropped without a flush discards its write errors.
269    ///
270    /// # Errors
271    ///
272    /// Returns [`StatsWriteError::Io`] when the flush fails.
273    pub fn finish(self) -> Result<u64, StatsWriteError> {
274        Ok(self.into_inner()?.1)
275    }
276
277    /// Flushes the writer and returns it with the number of rows written.
278    ///
279    /// # Errors
280    ///
281    /// Returns [`StatsWriteError::Io`] when the flush fails.
282    pub fn into_inner(mut self) -> Result<(W, u64), StatsWriteError> {
283        self.out.flush()?;
284        Ok((self.out, self.rows))
285    }
286
287    /// Number of rows written so far, not counting the header.
288    pub fn rows(&self) -> u64 {
289        self.rows
290    }
291}
292
293/// Returns the scalar that `part` extracts from `value`, or `None` when `value` has no such part.
294///
295/// A `None` means the series changed shape after the layout was fixed.
296fn part_value(value: &StatValue, part: Part) -> Option<f64> {
297    match (value, part) {
298        (StatValue::Scalar(v), Part::Scalar) => Some(*v),
299        (StatValue::Vector2D { x, .. }, Part::VecX) => Some(*x),
300        (StatValue::Vector2D { y, .. }, Part::VecY) => Some(*y),
301        (StatValue::Vector2D { x, y }, Part::VecMagnitude) => Some(x.hypot(*y)),
302        (StatValue::Histogram { counts, .. }, Part::Bucket(bucket)) => counts.get(bucket).map(|c| *c as f64),
303        (StatValue::Histogram { counts, .. }, Part::BucketTotal) => Some(counts.iter().sum::<u64>() as f64),
304        _ => None,
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    const C: [u8; 4] = [0, 0, 0, 255];
313
314    fn scalar(label: &'static str, v: f64) -> StatEntry {
315        StatEntry {
316            label,
317            value: StatValue::Scalar(v),
318            color: C,
319        }
320    }
321
322    fn vec2(label: &'static str, x: f64, y: f64) -> StatEntry {
323        StatEntry {
324            label,
325            value: StatValue::Vector2D { x, y },
326            color: C,
327        }
328    }
329
330    fn hist(label: &'static str, edges: Vec<f64>, counts: Vec<u64>) -> StatEntry {
331        StatEntry {
332            label,
333            value: StatValue::Histogram { edges, counts },
334            color: C,
335        }
336    }
337
338    /// Runs samples through a writer and returns the CSV text.
339    fn render(samples: &[(u64, Vec<StatEntry>)]) -> String {
340        let mut buf = Vec::new();
341        let mut writer = StatsWriter::new(&mut buf);
342        for (tick, stats) in samples {
343            writer.push(*tick, stats).expect("push should succeed");
344        }
345        writer.finish().expect("finish should succeed");
346        String::from_utf8(buf).expect("output should be utf8")
347    }
348
349    #[test]
350    fn scalars_write_one_column_each() {
351        let csv = render(&[
352            (0, vec![scalar("Alive", 10.0), scalar("Dead", 5.0)]),
353            (1, vec![scalar("Alive", 12.0), scalar("Dead", 3.0)]),
354        ]);
355        assert_eq!(csv, "tick,Alive,Dead\n0,10,5\n1,12,3\n");
356    }
357
358    #[test]
359    fn vectors_split_into_x_y_and_magnitude() {
360        let csv = render(&[(7, vec![vec2("Velocity", 3.0, 4.0)])]);
361        assert_eq!(csv, "tick,Velocity.x,Velocity.y,Velocity.magnitude\n7,3,4,5\n");
362    }
363
364    #[test]
365    fn histograms_label_buckets_by_range_and_add_a_total() {
366        let csv = render(&[(0, vec![hist("Speed", vec![0.0, 1.0, 2.0], vec![4, 6])])]);
367        assert_eq!(csv, "tick,\"Speed.[0, 1)\",\"Speed.[1, 2)\",Speed.total\n0,4,6,10\n");
368    }
369
370    /// A histogram with fewer edges than buckets still produces a column per bucket.
371    #[test]
372    fn histogram_falls_back_to_bucket_index_without_edges() {
373        let csv = render(&[(0, vec![hist("H", vec![], vec![1, 2])])]);
374        assert_eq!(csv, "tick,H.bucket 0,H.bucket 1,H.total\n0,1,2,3\n");
375    }
376
377    #[test]
378    fn header_is_written_once_for_many_rows() {
379        let csv = render(&[
380            (0, vec![scalar("A", 1.0)]),
381            (1, vec![scalar("A", 2.0)]),
382            (2, vec![scalar("A", 3.0)]),
383        ]);
384        assert_eq!(csv.lines().filter(|l| l.starts_with("tick")).count(), 1);
385        assert_eq!(csv.lines().count(), 4);
386    }
387
388    #[test]
389    fn no_samples_writes_nothing() {
390        let csv = render(&[]);
391        assert!(csv.is_empty(), "expected empty output, got {csv:?}");
392    }
393
394    #[test]
395    fn fractional_values_keep_precision() {
396        let csv = render(&[(0, vec![scalar("A", 0.1 + 0.2)])]);
397        // The sum keeps its round-trip digits past 0.3.
398        assert!(csv.contains("0.30000000000000004"), "got {csv}");
399    }
400
401    #[test]
402    fn non_finite_values_become_empty_cells() {
403        let csv = render(&[(0, vec![scalar("A", f64::NAN), scalar("B", f64::INFINITY)])]);
404        assert_eq!(csv, "tick,A,B\n0,,\n");
405    }
406
407    #[test]
408    fn labels_containing_commas_are_quoted() {
409        let csv = render(&[(0, vec![scalar("Susceptible, count", 1.0), scalar("B", 2.0)])]);
410        assert_eq!(csv, "tick,\"Susceptible, count\",B\n0,1,2\n");
411        // Three fields per row, so the comma did not shift the columns.
412        assert_eq!(csv.lines().count(), 2);
413    }
414
415    #[test]
416    fn finish_returns_the_row_count_excluding_the_header() {
417        let mut buf = Vec::new();
418        let mut writer = StatsWriter::new(&mut buf);
419        writer.push(0, &[scalar("A", 1.0)]).expect("push");
420        writer.push(1, &[scalar("A", 2.0)]).expect("push");
421        assert_eq!(writer.finish().expect("finish"), 2);
422    }
423
424    #[test]
425    fn a_series_disappearing_mid_run_is_an_error() {
426        let mut buf = Vec::new();
427        let mut writer = StatsWriter::new(&mut buf);
428        writer.push(0, &[scalar("A", 1.0), scalar("B", 2.0)]).expect("push");
429        let err = writer.push(1, &[scalar("A", 1.0)]).expect_err("should reject");
430        assert!(err.to_string().contains("changed mid-run"), "got {err}");
431    }
432
433    #[test]
434    fn a_series_changing_kind_mid_run_is_an_error() {
435        let mut buf = Vec::new();
436        let mut writer = StatsWriter::new(&mut buf);
437        writer.push(0, &[scalar("A", 1.0)]).expect("push");
438        let err = writer.push(1, &[vec2("A", 1.0, 2.0)]).expect_err("should reject");
439        assert!(err.to_string().contains("changed shape"), "got {err}");
440    }
441
442    #[test]
443    fn a_histogram_losing_buckets_mid_run_is_an_error() {
444        let mut buf = Vec::new();
445        let mut writer = StatsWriter::new(&mut buf);
446        writer
447            .push(0, &[hist("H", vec![0.0, 1.0, 2.0], vec![1, 2])])
448            .expect("push");
449        let err = writer
450            .push(1, &[hist("H", vec![0.0, 1.0], vec![1])])
451            .expect_err("should reject");
452        assert!(err.to_string().contains("changed shape"), "got {err}");
453    }
454
455    #[test]
456    fn mixed_kinds_keep_series_order() {
457        let csv = render(&[(
458            0,
459            vec![
460                scalar("S", 1.0),
461                vec2("V", 0.0, 2.0),
462                hist("H", vec![0.0, 1.0], vec![3]),
463            ],
464        )]);
465        assert_eq!(csv, "tick,S,V.x,V.y,V.magnitude,\"H.[0, 1)\",H.total\n0,1,0,2,2,3,3\n");
466    }
467
468    #[test]
469    fn a_model_with_no_stats_still_writes_ticks() {
470        let csv = render(&[(0, vec![]), (5, vec![])]);
471        assert_eq!(csv, "tick\n0\n5\n");
472    }
473
474    #[test]
475    fn stat_columns_name_what_the_writer_writes() {
476        let stats = vec![
477            scalar("Susceptible, count", 7.0),
478            vec2("V", 3.0, 4.0),
479            hist("H", vec![0.0, 1.0, 2.0], vec![4, 6]),
480        ];
481        let columns = StatColumns::plan(&stats);
482        let mut values = Vec::new();
483        columns
484            .extract(9, &stats, &mut values)
485            .expect("a sample fits its own plan");
486
487        let csv = render(&[(9, stats)]);
488        let headers: Vec<&str> = (0..columns.len()).map(|i| columns.header(i)).collect();
489        assert_eq!(csv.lines().next(), Some(format!("tick,{}", headers.join(",")).as_str()));
490
491        let records = crate::export::csv::parse_records(&csv).expect("the writer writes valid CSV");
492        let names: Vec<&str> = (0..columns.len()).map(|i| columns.name(i)).collect();
493        assert_eq!(records[0][1..], names[..], "names are the headers unescaped");
494        let row: Vec<String> = values.iter().map(|&value| fmt_f64(value)).collect();
495        assert_eq!(records[1][1..], row[..]);
496        assert_eq!(names[0], "Susceptible, count");
497
498        let buckets: Vec<&str> = (0..columns.len())
499            .filter(|&i| columns.is_bucket(i))
500            .map(|i| columns.name(i))
501            .collect();
502        assert_eq!(buckets, ["H.[0, 1)", "H.[1, 2)"]);
503    }
504
505    #[test]
506    fn a_bare_label_resolves_to_the_column_of_its_scalar() {
507        let stats = vec![
508            scalar("S", 1.0),
509            vec2("V", 3.0, 4.0),
510            hist("H", vec![0.0, 1.0, 2.0], vec![4, 6]),
511        ];
512        let columns = StatColumns::plan(&stats);
513        let mut values = Vec::new();
514        columns
515            .extract(0, &stats, &mut values)
516            .expect("a sample fits its own plan");
517        for entry in &stats {
518            let i = columns.resolve(entry.label).expect("every label resolves");
519            assert_eq!(values[i], entry.value.scalar(), "{}", entry.label);
520        }
521        assert_eq!(columns.resolve("V"), columns.resolve("V.magnitude"));
522        assert_eq!(columns.resolve("H"), columns.resolve("H.total"));
523        assert_eq!(columns.resolve("V.y").map(|i| columns.name(i)), Some("V.y"));
524        assert_eq!(columns.resolve("Missing"), None);
525
526        // An exact name wins over a bare label.
527        let shadowed = StatColumns::plan(&[vec2("V", 3.0, 4.0), scalar("V", 1.0)]);
528        assert_eq!(shadowed.resolve("V"), Some(3));
529    }
530
531    /// The two paths by which a stat series reaches a file must produce the same file.
532    ///
533    /// The app writes through a [`StatsWriter`] while recording, and replays
534    /// [`crate::view::StatsHistory`] otherwise. Both claim the column layout the headless runner
535    /// writes, and only one of them holds structured values.
536    mod parity {
537        use super::super::StatsWriter;
538        use crate::view::{StatDescriptor, StatEntry, StatValue, StatsHistory};
539
540        const C: [u8; 4] = [9, 9, 9, 255];
541
542        /// A scalar beside a vector, so the replay has to reach both stores.
543        fn sample(i: u64) -> Vec<StatEntry> {
544            vec![
545                StatEntry {
546                    label: "Alive",
547                    value: StatValue::Scalar(i as f64 * 1.5),
548                    color: C,
549                },
550                StatEntry {
551                    label: "Average Velocity",
552                    value: StatValue::Vector2D {
553                        x: i as f64,
554                        y: -(i as f64) * 0.25,
555                    },
556                    color: C,
557                },
558            ]
559        }
560
561        fn render(samples: &[(u64, Vec<StatEntry>)]) -> String {
562            let mut writer = StatsWriter::new(Vec::new());
563            for (tick, stats) in samples {
564                writer.push(*tick, stats).expect("push should succeed");
565            }
566            let (csv, _) = writer.into_inner().expect("finish should succeed");
567            String::from_utf8(csv).expect("output should be utf8")
568        }
569
570        #[test]
571        fn a_replayed_history_writes_what_a_recording_writes() {
572            let descriptors = vec![
573                StatDescriptor::new("Alive", C),
574                StatDescriptor::new("Average Velocity", C),
575            ];
576            let mut history = StatsHistory::new(descriptors, Some(64));
577            let mut recorded = Vec::new();
578            for tick in 0..20 {
579                let stats = sample(tick);
580                history.push_entries(&stats, tick);
581                recorded.push((tick, stats));
582            }
583
584            let replayed: Vec<(u64, Vec<StatEntry>)> = (0..history.len())
585                .map(|j| (history.tick(j).expect("tick"), history.entries(j).expect("entries")))
586                .collect();
587
588            assert_eq!(render(&replayed), render(&recorded));
589        }
590
591        /// A history that has wrapped writes the tail of what a recording holds, column for column.
592        #[test]
593        fn a_wrapped_history_writes_the_tail_of_the_recording() {
594            let descriptors = vec![
595                StatDescriptor::new("Alive", C),
596                StatDescriptor::new("Average Velocity", C),
597            ];
598            let mut history = StatsHistory::new(descriptors, Some(5));
599            let mut recorded = Vec::new();
600            for tick in 0..20 {
601                let stats = sample(tick);
602                history.push_entries(&stats, tick);
603                recorded.push((tick, stats));
604            }
605
606            let replayed: Vec<(u64, Vec<StatEntry>)> = (0..history.len())
607                .map(|j| (history.tick(j).expect("tick"), history.entries(j).expect("entries")))
608                .collect();
609
610            let full = render(&recorded);
611            let tail = render(&replayed);
612            let header = full.lines().next().expect("a header");
613            assert_eq!(
614                tail.lines().next(),
615                Some(header),
616                "the columns must not depend on the path"
617            );
618            for line in tail.lines().skip(1) {
619                assert!(full.contains(line), "replayed row {line:?} is not in the recording");
620            }
621            assert_eq!(tail.lines().count(), 6, "five samples plus the header");
622        }
623    }
624}
henad_core/export/stats_csv.rs

henad_core/export/
stats_csv.rs
