Measuring allocator and span statistics

Cargo.toml

@@ -41,13 +41,8 @@ opt-level = 2
 opt-level = 3
 
 [workspace.dependencies]
-anyhow = "1.0.93"
-argh = "0.1.12"
-clap = { version = "4.4", features = ["derive"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
-tracing = "0.1.41"
-tracing-subscriber = { version = "0.3.18", features = ["env-filter", "ansi"] }
 base64 = "0.22.1"
 ruint = { version = "1.12.3", features = ["num-traits", "rand"] }
 rand = "0.8.5"
@@ -82,5 +77,13 @@ ark-poly = "0.5"
 ark-std = { version = "0.5", features = ["std"] }
 ark-serialize = "0.5"
 
+# Bin
+anyhow = "1.0.93"
+argh = "0.1.12"
+tracing = "0.1.41"
+tracing-subscriber = { version = "0.3.18", features = ["env-filter", "ansi"] }
+nu-ansi-term = "0.46.0"
+
+
 [patch.crates-io]
 ruint = { git = "https://github.com/recmo/uint", branch = "main" }

noir-r1cs/Cargo.toml

@@ -6,15 +6,14 @@ edition = "2021"
 [lints]
 workspace = true
 
+[[bin]]
+name = "noir-r1cs"
+path = "src/bin/main.rs"
+
 [dependencies]
-argh.workspace = true
 serde.workspace = true
 serde_json.workspace = true
-anyhow.workspace = true
-tracing.workspace = true
-tracing-subscriber.workspace = true
 base64.workspace = true
-clap.workspace = true
 zeroize.workspace = true
 itertools.workspace = true
 ruint.workspace = true
@@ -42,10 +41,10 @@ ark-poly.workspace = true
 ark-std.workspace = true
 ark-serialize.workspace = true
 
-[[bin]]
-name = "circuit_stats"
-path = "src/bin/circuit_stats.rs"
-
-[[bin]]
-name = "prover"
-path = "src/bin/prover.rs"
+# Binary 
+# See <https://github.com/rust-lang/cargo/issues/1982>
+argh.workspace = true
+anyhow.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true
+nu-ansi-term.workspace = true

noir-r1cs/src/bin/main.rs

@@ -0,0 +1,22 @@
+#![allow(missing_docs)]
+mod cmd;
+mod measuring_alloc;
+mod span_stats;
+
+use {
+    self::{cmd::Command, measuring_alloc::MeasuringAllocator, span_stats::SpanStats},
+    anyhow::Result,
+    tracing_subscriber::{self, fmt, layer::SubscriberExt as _, Registry},
+};
+
+#[global_allocator]
+static ALLOC: MeasuringAllocator = MeasuringAllocator::new();
+
+fn main() -> Result<()> {
+    let subscriber = Registry::default().with(SpanStats);
+    tracing::subscriber::set_global_default(subscriber)?;
+
+    // Run CLI command
+    let args = argh::from_env::<cmd::Args>();
+    args.run()
+}

noir-r1cs/src/bin/measuring_alloc.rs

@@ -0,0 +1,56 @@
+use std::{
+    alloc::{GlobalAlloc, Layout, System as SystemAlloc},
+    sync::atomic::{AtomicUsize, Ordering},
+};
+
+/// Custom allocator that keeps track of statistics to see program memory
+/// consumption.
+pub struct MeasuringAllocator {
+    current: AtomicUsize,
+    max:     AtomicUsize,
+    count:   AtomicUsize,
+}
+
+impl MeasuringAllocator {
+    pub const fn new() -> Self {
+        Self {
+            current: AtomicUsize::new(0),
+            max:     AtomicUsize::new(0),
+            count:   AtomicUsize::new(0),
+        }
+    }
+
+    pub fn current(&self) -> usize {
+        self.current.load(Ordering::SeqCst)
+    }
+
+    pub fn max(&self) -> usize {
+        self.max.load(Ordering::SeqCst)
+    }
+
+    pub fn reset_max(&self) -> usize {
+        let current = self.current();
+        self.max.store(current, Ordering::SeqCst);
+        current
+    }
+
+    pub fn count(&self) -> usize {
+        self.count.load(Ordering::SeqCst)
+    }
+}
+
+#[allow(unsafe_code)]
+unsafe impl GlobalAlloc for MeasuringAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // We just ignore the race conditions here...
+        let prev = self.current.fetch_add(layout.size(), Ordering::SeqCst);
+        self.max.fetch_max(prev + layout.size(), Ordering::SeqCst);
+        self.count.fetch_add(1, Ordering::SeqCst);
+        SystemAlloc.alloc(layout)
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        self.current.fetch_sub(layout.size(), Ordering::SeqCst);
+        SystemAlloc.dealloc(ptr, layout)
+    }
+}

noir-r1cs/src/bin/span_stats.rs

@@ -0,0 +1,253 @@
+//! Using `tracing` spans to print performance statistics for the program.
+//!
+//! NOTE: This module is only included in the bin, not in the lib.
+use {
+    crate::ALLOC,
+    nu_ansi_term::{Color, Style},
+    std::{
+        fmt::{self, Display, Formatter, Write as _},
+        sync::{Arc, Mutex},
+        time::Instant,
+    },
+    tracing::{
+        field::{Field, Visit},
+        span::{self, Attributes, Id},
+        Level, Subscriber,
+    },
+    tracing_subscriber::{self, layer::Context, registry::LookupSpan, Layer},
+};
+
+const DIM: &'static str = "\x1b[2m";
+const UNDIM: &'static str = "\x1b[22m";
+
+// Span extension data
+pub(crate) struct Data {
+    depth:       usize,
+    time:        Instant,
+    memory:      usize,
+    allocations: usize,
+
+    /// `peak_memory` will be updated as it is not monotonic
+    peak_memory: usize,
+    children:    bool,
+    kvs:         Vec<(&'static str, String)>,
+}
+
+impl Data {
+    pub fn new(attrs: &Attributes<'_>, depth: usize) -> Self {
+        let mut span = Self {
+            depth,
+            time: Instant::now(),
+            memory: ALLOC.current(),
+            allocations: ALLOC.count(),
+            peak_memory: ALLOC.current(),
+            children: false,
+            kvs: Vec::new(),
+        };
+        attrs.record(&mut span);
+        span
+    }
+}
+
+impl Visit for Data {
+    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
+        self.kvs.push((field.name(), format!("{:?}", value)))
+    }
+}
+
+pub struct FmtEvent<'a>(&'a mut String);
+
+impl<'a> Visit for FmtEvent<'a> {
+    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
+        match field.name() {
+            "message" => {
+                write!(self.0, " {:?}", value).unwrap();
+            }
+            name => {
+                write!(self.0, " {}={:?}", name, value).unwrap();
+            }
+        }
+    }
+}
+
+/// Logging layer that keeps track of time and memory consumption of spans.
+pub struct SpanStats;
+
+impl<S> Layer<S> for SpanStats
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
+    fn on_new_span(&self, attrs: &Attributes, id: &Id, ctx: Context<S>) {
+        let span = ctx.span(id).expect("invalid span in on_new_span");
+
+        // Add Data if it hasn't already
+        if span.extensions().get::<Data>().is_none() {
+            let depth = span
+                .parent()
+                .map(|s| {
+                    s.extensions()
+                        .get::<Data>()
+                        .expect("parent span has no data")
+                        .depth
+                        + 1
+                })
+                .unwrap_or(0);
+            let data = Data::new(attrs, depth);
+            span.extensions_mut().insert(data);
+        }
+
+        // Flag child on parent
+        if let Some(parent) = span.parent() {
+            if let Some(data) = parent.extensions_mut().get_mut::<Data>() {
+                data.children = true;
+            }
+        }
+
+        // Fetch data
+        let ext = span.extensions();
+        let data = ext.get::<Data>().expect("span does not have data");
+
+        let mut buffer = String::with_capacity(100);
+
+        // Box draw tree indentation
+        if data.depth >= 1 {
+            for _ in 0..(data.depth - 1) {
+                let _ = write!(&mut buffer, "│ ");
+            }
+            let _ = write!(&mut buffer, "├─");
+        }
+        let _ = write!(&mut buffer, "╮ ");
+
+        // Span name
+        let _ = write!(
+            &mut buffer,
+            "{DIM}{}::{UNDIM}{}",
+            span.metadata().target(),
+            span.metadata().name()
+        );
+
+        // KV args
+        for (key, val) in &data.kvs {
+            let _ = write!(&mut buffer, " {}={}", key, val);
+        }
+
+        eprintln!("{}", buffer);
+    }
+
+    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
+        let span = ctx.current_span().id().and_then(|id| ctx.span(id));
+
+        let mut buffer = String::with_capacity(100);
+
+        // Span indentation + time in span
+        if let Some(span) = &span {
+            // Flag child on parent
+            if let Some(parent) = span.parent() {
+                if let Some(data) = parent.extensions_mut().get_mut::<Data>() {
+                    data.children = true;
+                }
+            }
+
+            if let Some(data) = span.extensions().get::<Data>() {
+                // Box draw tree indentation
+                for _ in 0..(data.depth + 1) {
+                    let _ = write!(&mut buffer, "│ ");
+                }
+
+                // Time
+                let elapsed = data.time.elapsed();
+                let _ = write!(
+                    &mut buffer,
+                    "{DIM}{:6}s {UNDIM}",
+                    human(elapsed.as_secs_f64())
+                );
+            }
+        }
+
+        // Log level
+        match *event.metadata().level() {
+            Level::TRACE => write!(&mut buffer, "TRACE"),
+            Level::DEBUG => write!(&mut buffer, "DEBUG"),
+            Level::INFO => write!(&mut buffer, "\x1b[1;32mINFO\x1b[0m"),
+            Level::WARN => write!(&mut buffer, "\x1b[1;38;5;208mWARN\x1b[0m"),
+            Level::ERROR => write!(&mut buffer, "\x1b[1;31mERROR\x1b[0m"),
+        }
+        .unwrap();
+
+        let mut visitor = FmtEvent(&mut buffer);
+        event.record(&mut visitor);
+
+        eprintln!("{}", buffer);
+    }
+
+    fn on_close(&self, id: Id, ctx: Context<S>) {
+        let span = ctx.span(&id).expect("invalid span in on_close");
+        let ext = span.extensions();
+        let data = ext.get::<Data>().expect("span does not have data");
+
+        let duration = data.time.elapsed();
+        let peak_memory: usize = std::cmp::max(ALLOC.max(), data.peak_memory);
+        let allocations = ALLOC.count() - data.allocations;
+        let own = peak_memory - data.memory;
+
+        let mut buffer = String::with_capacity(100);
+
+        // Box draw tree indentation
+        if data.depth >= 1 {
+            for _ in 0..(data.depth - 1) {
+                let _ = write!(&mut buffer, "│ ");
+            }
+            let _ = write!(&mut buffer, "├─");
+        }
+        let _ = write!(&mut buffer, "╯ ");
+
+        // Short span name if not childless
+        if data.children {
+            let _ = write!(&mut buffer, "{DIM}{}: {UNDIM}", span.metadata().name());
+        }
+
+        // Print stats
+        let _ = write!(
+            &mut buffer,
+            "{}s{DIM} duration, {UNDIM}{}B{DIM} peak memory, {UNDIM}{}B{DIM} local, \
+             {UNDIM}{:#}{DIM} allocations{UNDIM}",
+            human(duration.as_secs_f64()),
+            human(peak_memory as f64),
+            human(own as f64),
+            human(allocations as f64)
+        );
+
+        eprintln!("{}", buffer);
+    }
+}
+
+/// Pretty print a float using SI-prefixes.
+pub fn human(value: f64) -> impl Display {
+    pub struct Human(f64);
+    impl Display for Human {
+        fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+            let log10 = if self.0.is_normal() {
+                self.0.abs().log10()
+            } else {
+                0.0
+            };
+            let si_power = ((log10 / 3.0).floor() as isize).clamp(-10, 10);
+            let value = self.0 * 10_f64.powi((-si_power * 3) as i32);
+            let digits = f.precision().unwrap_or(3) - 1 - (log10 - 3.0 * si_power as f64) as usize;
+            let separator = if f.alternate() { "" } else { "\u{202F}" };
+            if f.width() == Some(6) && digits == 0 {
+                write!(f, " ")?;
+            }
+            write!(f, "{value:.digits$}{separator}")?;
+            let suffix = "qryzafpnμm kMGTPEZYRQ"
+                .chars()
+                .nth((si_power + 10) as usize)
+                .unwrap();
+            if suffix != ' ' {
+                write!(f, "{suffix}")?;
+            }
+            Ok(())
+        }
+    }
+    Human(value)
+}