Read/write records (#1117)

raviqqe · Apr 28, 2024 · 101419a · 101419a
1 parent b32989d
commit 101419a
Show file tree

Hide file tree

Showing 8 changed files with 252 additions and 126 deletions.
diff --git a/cmd/profile/src/main.rs b/cmd/profile/src/main.rs
@@ -12,13 +12,13 @@ use stak_configuration::DEFAULT_HEAP_SIZE;
 use stak_device::StdioDevice;
 use stak_primitive::SmallPrimitiveSet;
 use stak_profiler::{
-    calculate_durations, calculate_flamegraph, collapse_stacks, DurationRecord, ProcedureRecord,
-    StackProfiler, StackedRecord,
+    calculate_durations, calculate_flamegraph, collapse_stacks, read_records, reverse_stacks,
+    write_records, DurationRecord, ProcedureRecord, StackProfiler,
 };
 use stak_vm::Vm;
 use std::{
     fs::{read, OpenOptions},
-    io::{stdin, stdout, BufRead, BufWriter},
+    io::{stdin, stdout, BufWriter},
     path::PathBuf,
 };
 
@@ -62,6 +62,8 @@ enum Analysis {
     Duration,
     /// Calculates collapsed stacks.
     StackCollapse,
+    /// Calculates reversed stacks.
+    StackReverse,
     /// Calculates a flamegraph.
     Flamegraph,
 }
@@ -88,21 +90,19 @@ fn main() -> Result<(), MainError> {
             let writer = BufWriter::new(stdout().lock());
 
             match arguments.command {
-                Analysis::Duration => calculate_durations(
-                    reader.lines().map(|line| {
-                        let mut record = line?.parse::<ProcedureRecord>()?;
-                        record.stack_mut().reverse_frames();
-                        Ok(record)
-                    }),
+                Analysis::Duration => write_records(
+                    calculate_durations(read_records::<ProcedureRecord>(reader)),
                     writer,
                 )?,
-                Analysis::StackCollapse => collapse_stacks(
-                    reader.lines().map(|line| line?.parse::<DurationRecord>()),
+                Analysis::StackCollapse => write_records(
+                    collapse_stacks(read_records::<DurationRecord>(reader)),
                     writer,
                 )?,
-                Analysis::Flamegraph => {
-                    calculate_flamegraph(reader.lines().map(|line| line?.parse()), writer)?
-                }
+                Analysis::StackReverse => write_records(
+                    reverse_stacks(read_records::<DurationRecord>(reader)),
+                    writer,
+                )?,
+                Analysis::Flamegraph => calculate_flamegraph(read_records(reader), writer)?,
             }
         }
     }

diff --git a/profiler/src/collapse.rs b/profiler/src/collapse.rs
@@ -1,35 +1,28 @@
 use crate::{Error, StackedRecord};
-use std::io::Write;
 
 /// Collapses stacks.
-pub fn collapse_stacks(
-    records: impl IntoIterator<Item = Result<impl StackedRecord, Error>>,
-    mut writer: impl Write,
-) -> Result<(), Error> {
-    for record in records {
+pub fn collapse_stacks<R: StackedRecord>(
+    records: impl IntoIterator<Item = Result<R, Error>>,
+) -> impl Iterator<Item = Result<R, Error>> {
+    records.into_iter().map(|record| {
         let mut record = record?;
 
         record.stack_mut().collapse_frames();
 
-        writeln!(writer, "{record}")?;
-    }
-
-    Ok(())
+        Ok(record)
+    })
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::{DurationRecord, Stack};
-    use indoc::indoc;
     use pretty_assertions::assert_eq;
 
     #[test]
     fn collapse() {
-        let mut buffer = vec![];
-
-        collapse_stacks(
-            [
+        assert_eq!(
+            collapse_stacks([
                 Ok(DurationRecord::new(
                     Stack::new(vec![Some("bar".into()), Some("foo".into())]),
                     123,
@@ -43,19 +36,22 @@ mod tests {
                     ]),
                     42,
                 )),
-            ],
-            &mut buffer,
-        )
-        .unwrap();
-
-        assert_eq!(
-            String::from_utf8(buffer).unwrap(),
-            indoc!(
-                "
-                bar;foo\t123
-                baz;bar;foo\t42
-                "
-            )
+            ])
+            .collect::<Vec<_>>(),
+            vec![
+                Ok(DurationRecord::new(
+                    Stack::new(vec![Some("bar".into()), Some("foo".into())]),
+                    123,
+                )),
+                Ok(DurationRecord::new(
+                    Stack::new(vec![
+                        Some("baz".into()),
+                        Some("bar".into()),
+                        Some("foo".into()),
+                    ]),
+                    42,
+                )),
+            ]
         );
     }
 }
diff --git a/profiler/src/duration.rs b/profiler/src/duration.rs
@@ -1,69 +1,62 @@
 use crate::{DurationRecord, Error, ProcedureOperation, ProcedureRecord, Stack, StackedRecord};
-use std::io::Write;
 
 /// Calculates durations.
 pub fn calculate_durations(
     records: impl IntoIterator<Item = Result<ProcedureRecord, Error>>,
-    mut writer: impl Write,
-) -> Result<(), Error> {
-    let mut first = true;
+) -> impl Iterator<Item = Result<DurationRecord, Error>> {
     let mut stack = vec![];
 
-    for record in records {
-        let record = record?;
+    records
+        .into_iter()
+        .enumerate()
+        .map(move |(index, record)| {
+            let record = record?;
 
-        if first {
-            stack.push(ProcedureRecord::new(
-                ProcedureOperation::Call,
-                Stack::new(vec![None]),
-                record.time(),
-            ));
-            first = false;
-        }
-
-        match record.operation() {
-            ProcedureOperation::Call => {
-                stack.push(record);
-            }
-            ProcedureOperation::Return => calculate_duration(&mut stack, &record, &mut writer)?,
-            ProcedureOperation::ReturnCall => {
-                calculate_duration(&mut stack, &record, &mut writer)?;
-                stack.push(record);
+            if index == 0 {
+                stack.push(ProcedureRecord::new(
+                    ProcedureOperation::Call,
+                    Stack::new(vec![None]),
+                    record.time(),
+                ));
             }
-        }
-    }
 
-    Ok(())
+            Ok(match record.operation() {
+                ProcedureOperation::Call => {
+                    stack.push(record);
+                    None
+                }
+                ProcedureOperation::Return => Some(calculate_duration(&mut stack, &record)?),
+                ProcedureOperation::ReturnCall => {
+                    let duration = calculate_duration(&mut stack, &record)?;
+                    stack.push(record);
+                    Some(duration)
+                }
+            })
+        })
+        .flat_map(Result::transpose)
 }
 
 fn calculate_duration(
     stack: &mut Vec<ProcedureRecord>,
     record: &ProcedureRecord,
-    mut writer: impl Write,
-) -> Result<(), Error> {
+) -> Result<DurationRecord, Error> {
     let previous = stack.pop().ok_or(Error::MissingCallRecord)?;
 
-    writeln!(
-        writer,
-        "{}",
-        DurationRecord::new(previous.stack().clone(), record.time() - previous.time()),
-    )?;
-
-    Ok(())
+    Ok(DurationRecord::new(
+        previous.stack().clone(),
+        record.time() - previous.time(),
+    ))
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use indoc::indoc;
     use pretty_assertions::assert_eq;
 
     #[test]
     fn analyze_call() {
-        let mut buffer = vec![];
-
-        calculate_durations(
-            [
+        assert_eq!(
+            calculate_durations([
                 Ok(ProcedureRecord::new(
                     ProcedureOperation::Call,
                     Stack::new(vec![
@@ -82,20 +75,23 @@ mod tests {
                     ]),
                     42,
                 )),
-            ],
-            &mut buffer,
-        )
-        .unwrap();
-
-        assert_eq!(String::from_utf8(buffer).unwrap(), "baz;bar;foo\t42\n");
+            ])
+            .collect::<Vec<_>>(),
+            vec![Ok(DurationRecord::new(
+                Stack::new(vec![
+                    Some("baz".into()),
+                    Some("bar".into()),
+                    Some("foo".into()),
+                ]),
+                42,
+            ))]
+        );
     }
 
     #[test]
     fn analyze_nested_calls() {
-        let mut buffer = vec![];
-
-        calculate_durations(
-            [
+        assert_eq!(
+            calculate_durations([
                 Ok(ProcedureRecord::new(
                     ProcedureOperation::Call,
                     Stack::new(vec![Some("baz".into())]),
@@ -134,20 +130,26 @@ mod tests {
                     Stack::new(vec![Some("baz".into())]),
                     126,
                 )),
-            ],
-            &mut buffer,
-        )
-        .unwrap();
-
-        assert_eq!(
-            String::from_utf8(buffer).unwrap(),
-            indoc!(
-                "
-                baz;bar;foo\t40
-                baz;bar\t83
-                baz\t126
-                "
-            )
+            ])
+            .collect::<Vec<_>>(),
+            [
+                Ok(DurationRecord::new(
+                    Stack::new(vec![
+                        Some("baz".into()),
+                        Some("bar".into()),
+                        Some("foo".into()),
+                    ]),
+                    40,
+                )),
+                Ok(DurationRecord::new(
+                    Stack::new(vec![Some("baz".into()), Some("bar".into())]),
+                    83,
+                )),
+                Ok(DurationRecord::new(
+                    Stack::new(vec![Some("baz".into())]),
+                    126,
+                )),
+            ]
         );
     }
 }