Merge pull request #1815 from o1-labs/dw/move-alloc-test

Vendor: move alloc-test into mina-rust monorepo

Author: dannywillems

.github/workflows/tests.yaml

@@ -283,6 +283,7 @@ jobs:
               --exclude transaction_fuzzer \
               --exclude mina-archive-breadcrumb-compare \
               --exclude webrtc-sniffer
+              # --exclude alloc-test
               # --exclude mina-tree
               # --exclude redux-rs
               # --exclude vrf

CHANGELOG.md

@@ -31,6 +31,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changes
 
+- **Dependencies**: vendor alloc-test into `vendor/alloc-test` to centralize all
+  code in the monorepo ([#1815](https://github.com/o1-labs/mina-rust/pull/1815))
 - **Dependencies**: vendor redux-rs into `vendor/redux` to centralize all code
   in the monorepo ([#1814](https://github.com/o1-labs/mina-rust/pull/1814))
 - **Build System**: enforce GNU sed on macOS, require `brew install gnu-sed`

Cargo.lock

@@ -27,6 +27,7 @@ members = [
   "tools/ledger-tool",
   "tools/transport",
   "tools/webrtc-sniffer",
+  "vendor/alloc-test",
   "vendor/redux",
   "vendor/salsa-simple",
   "vrf",
@@ -42,7 +43,7 @@ indexing_slicing = "warn"
 [workspace.dependencies]
 aes = "0.8"
 aes-gcm = "0.10.3"
-alloc-test = { git = "https://github.com/openmina/alloc-test.git" }
+alloc-test = { path = "vendor/alloc-test" }
 anyhow = "1.0.81"
 argon2 = { version = "0.5.3", features = ["std"] }
 ark-ec = { version = "0.5.0", features = ["std"] }

Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "alloc-test"
+version = "0.1.1"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+clap = { version = "4.0.18", features = ["derive", "env"], optional = true }
+derive_builder = "0.11.2"
+derive_more = { version = "0.99.17", features = [
+  "display",
+], default-features = false }
+num = "0.4.0"
+serde = { version = "1.0.146", features = ["derive"] }
+serde_json = { version = "1.0", optional = true }
+thiserror = "1.0.37"
+toml = { version = "0.5.9", optional = true }
+wasm-bindgen = "0.2.83"
+wasm-bindgen-test = { version = "0.3.33", optional = true }
+
+[features]
+default = ["benchmark"]
+benchmark = ["dep:clap", "dep:toml", "dep:serde_json", "dep:wasm-bindgen-test"]

vendor/alloc-test/Cargo.toml

@@ -0,0 +1,78 @@
+use std::alloc::{GlobalAlloc, Layout};
+
+#[derive(Debug, Default)]
+pub struct TracingAllocator<H: 'static, A>(A, H)
+where
+    A: GlobalAlloc;
+
+impl<H, A> TracingAllocator<H, A>
+where
+    A: GlobalAlloc,
+{
+    pub const fn new(hooks: H, allocator: A) -> Self {
+        TracingAllocator(allocator, hooks)
+    }
+}
+
+// pub const fn default_tracing_allocator() -> TracingAllocator<MemoryTracingHooks, System> {
+//     TracingAllocator(System, MemoryTracingHooks)
+// }
+
+unsafe impl<H, A> GlobalAlloc for TracingAllocator<H, A>
+where
+    A: GlobalAlloc,
+    H: AllocHooks,
+{
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let size = layout.size();
+        let align = layout.align();
+        let pointer = self.0.alloc(layout);
+        self.1.on_alloc(pointer, size, align);
+        pointer
+    }
+
+    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
+        let size = layout.size();
+        let align = layout.align();
+        self.0.dealloc(pointer, layout);
+        self.1.on_dealloc(pointer, size, align);
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let size = layout.size();
+        let align = layout.align();
+        let pointer = self.0.alloc_zeroed(layout);
+        self.1.on_alloc_zeroed(pointer, size, align);
+        pointer
+    }
+
+    unsafe fn realloc(&self, old_pointer: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let old_size = layout.size();
+        let align = layout.align();
+        let new_pointer = self.0.realloc(old_pointer, layout, new_size);
+        self.1
+            .on_realloc(old_pointer, new_pointer, old_size, new_size, align);
+        new_pointer
+    }
+}
+
+/// Trait for implementing allocation hooks in a tracing allocator.
+///
+/// # Safety
+///
+/// Implementors must ensure that hook methods do not allocate memory
+/// (to avoid infinite recursion) and are safe to call from any thread
+/// at any time during allocation operations.
+pub unsafe trait AllocHooks {
+    fn on_alloc(&self, pointer: *mut u8, size: usize, align: usize);
+    fn on_dealloc(&self, pointer: *mut u8, size: usize, align: usize);
+    fn on_alloc_zeroed(&self, pointer: *mut u8, size: usize, align: usize);
+    fn on_realloc(
+        &self,
+        old_pointer: *mut u8,
+        new_pointer: *mut u8,
+        old_size: usize,
+        new_size: usize,
+        align: usize,
+    );
+}

vendor/alloc-test/src/alloc/allocator.rs

@@ -0,0 +1,46 @@
+use super::measure::MemoryStats;
+
+pub fn alloc_benchmark<F: FnOnce() -> O, O>(id: &str, f: F) -> MemoryStats {
+    let (_, stats) = crate::alloc::measure::trace_allocs(f);
+    log!("\nmemory allocation stats for `{id}`:\n{stats}");
+    stats
+}
+
+pub fn alloc_log_toml<F: Fn() -> O, O>(id: &str, f: F) -> MemoryStats {
+    let (_, stats) = super::measure::trace_allocs(f);
+    log!(
+        "\nperformance stats for `{id}`:\n{stats}",
+        stats = toml::to_string(&stats).unwrap()
+    );
+    stats
+}
+
+#[macro_export]
+macro_rules! alloc_bench {
+    ($test:ident, $thresh:expr) => {
+        $crate::threshold::check_threshold_with_args(
+            || $crate::alloc::benchmark::alloc_benchmark(stringify!($test), $test),
+            "alloc_bench",
+            stringify!($test),
+            $thresh,
+        )
+    };
+}
+
+#[macro_export]
+macro_rules! alloc_bench_cmp_with_toml {
+    ($test:ident $(,)?) => {{
+        let value = $crate::alloc::benchmark::alloc_log_toml(stringify!($test), $test);
+        Result::<
+            $crate::alloc::measure::MemoryStats,
+            $crate::threshold::CheckThresholdError<$crate::alloc::compare::AllocThresholdsError>,
+        >::Ok(value)
+    }};
+    ($test:ident, $toml:expr, $limits:expr $(,)?) => {{
+        $crate::threshold::check_threshold_with_str(
+            || $crate::alloc::benchmark::alloc_benchmark(stringify!($test), $test),
+            $toml,
+            $limits,
+        )
+    }};
+}

vendor/alloc-test/src/alloc/benchmark.rs

@@ -0,0 +1,106 @@
+use derive_builder::Builder;
+use thiserror::Error;
+
+use super::measure::MemoryStats;
+use crate::threshold::{Threshold, ThresholdError, ThresholdFor};
+
+/// Limits for each allocation statistics parameter.
+#[derive(Debug, Builder)]
+pub struct AllocThresholds {
+    #[builder(default)]
+    pub current: Threshold<usize>,
+    #[builder(default)]
+    pub peak: Threshold<usize>,
+    #[builder(default)]
+    pub total_size: Threshold<usize>,
+    #[builder(default)]
+    pub total_num: Threshold<usize>,
+    #[builder(default)]
+    pub reallocs: Threshold<usize>,
+}
+
+#[derive(Debug, Error)]
+#[error("Allocation parameter `{param}`: {error}")]
+pub struct AllocThresholdsError {
+    error: ThresholdError<usize>,
+    param: &'static str,
+}
+
+macro_rules! check {
+    ($f:ident, $l:expr, $v:expr, $r:expr) => {
+        $l.$f
+            .check(&$v.$f, &$r.$f)
+            .map_err(|e| AllocThresholdsError {
+                error: e,
+                param: stringify!($f),
+            })
+    };
+}
+
+impl ThresholdFor<MemoryStats> for AllocThresholds {
+    type Error = AllocThresholdsError;
+
+    fn check_threshold(
+        &self,
+        value: &MemoryStats,
+        ref_value: &MemoryStats,
+    ) -> Result<(), Self::Error> {
+        self.check(value, ref_value)
+    }
+}
+
+impl AllocThresholds {
+    pub fn check(
+        &self,
+        value: &MemoryStats,
+        ref_value: &MemoryStats,
+    ) -> Result<(), AllocThresholdsError> {
+        check!(current, self, value, ref_value)?;
+        check!(peak, self, value, ref_value)?;
+        check!(total_size, self, value, ref_value)?;
+        check!(total_num, self, value, ref_value)?;
+        check!(reallocs, self, value, ref_value)?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn limits() {
+        let rs = MemoryStats {
+            current: 100,
+            peak: 1000,
+            total_size: 2000,
+            total_num: 100,
+            reallocs: 0,
+        };
+        let vs = MemoryStats {
+            current: 110,
+            peak: 1100,
+            total_size: 2200,
+            total_num: 110,
+            reallocs: 1,
+        };
+
+        let ls = AllocThresholdsBuilder::default().build().unwrap();
+        assert!(ls.check_threshold(&rs, &rs).is_ok());
+        assert!(ls.check_threshold(&vs, &rs).is_ok());
+
+        let ls = AllocThresholdsBuilder::default()
+            .current(Threshold::Cap(1))
+            .build()
+            .unwrap();
+        let r = ls.check(&vs, &rs);
+        assert!(r.unwrap_err().param == "current");
+
+        let ls = AllocThresholdsBuilder::default()
+            .reallocs(Threshold::Cap(0))
+            .build()
+            .unwrap();
+        let r = ls.check(&vs, &rs);
+        assert!(r.unwrap_err().param == "reallocs");
+    }
+}