nordisk/facett

facett-map/tests/common/heap.rs

//! **The measuring allocator** — one writer for every "what did this cost the heap"
//! guard in this crate.
//!
//! `#[global_allocator]` is a per-BINARY item and integration tests are separate
//! binaries, so the *installation* cannot be shared; the body can, and this is it.
//! A test file that wants to measure declares
//!
//! ```ignore
//! #[path = "common/heap.rs"]
//! mod heap;
//! #[global_allocator]
//! static A: heap::Counting = heap::Counting;
//! ```
//!
//! and reaches for [`measure`]. It goes through `#[path]` rather than `mod common;`
//! because `common/mod.rs` is the wgpu device gate and pulls `wgpu` in — a heap
//! measurement must stay runnable on a host with no adapter and on a default
//! (feature-less) build.
//!
//! # Why the counters are THREAD-LOCAL
//!
//! The allocator is process-wide but `#[test]` bodies run on parallel threads, so a
//! global tally measures "whatever the whole process allocated during my window",
//! which is not what any assertion means. MEASURED 2026-08-03 in
//! `zero_copy_gpu_lane.rs` at `--test-threads=2`: a body that allocates nothing read
//! **777 378 B**, and the number moved run to run. Scoping to the measuring thread
//! fixes it by construction: another thread's allocation simply is not this thread's
//! count. `Cell` with a `const` initialiser so the TLS access inside `alloc` cannot
//! itself allocate and recurse, and `try_with` so an allocation during TLS teardown
//! is ignored rather than panicking inside the allocator.
//!
//! # Why COUNT and not only BYTES
//!
//! Bytes alone cannot see a per-item temporary: a throwaway `Vec` that is built and
//! immediately memcpy'd into a shared buffer allocates roughly the same bytes as the
//! shared buffer's own growth, but it does it once **per item** instead of
//! `O(log n)` times in total. The allocation COUNT is the observable that separates
//! the two shapes, so both are reported.

#![allow(dead_code)]

use std::alloc::{GlobalAlloc, Layout, System};

thread_local! {
    static BYTES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
    static ALLOCS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
    static ARMED: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}

/// Add one allocation of `n` bytes to this thread's tally, but only while armed.
#[inline]
fn tally(n: usize) {
    let armed = ARMED.try_with(std::cell::Cell::get).unwrap_or(false);
    if armed {
        let _ = BYTES.try_with(|b| b.set(b.get() + n));
        let _ = ALLOCS.try_with(|c| c.set(c.get() + 1));
    }
}

/// The `System` allocator with a thread-local tally in front of it.
pub struct Counting;

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        tally(l.size());
        unsafe { System.alloc(l) }
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) }
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, n: usize) -> *mut u8 {
        // Only the GROWTH is new bytes — a `Vec` doubling its way to `S` bytes has
        // then allocated `S`, not `2S`. The realloc still counts as one allocation,
        // which is what makes a doubling buffer read as `O(log n)` events.
        tally(n.saturating_sub(l.size()));
        unsafe { System.realloc(p, l, n) }
    }
}

/// What one call cost this thread's heap: `(value, allocations, bytes)`.
///
/// Build the fixture BEFORE calling this — only `f` is inside the armed window.
pub fn measure<T>(f: impl FnOnce() -> T) -> (T, usize, usize) {
    BYTES.with(|b| b.set(0));
    ALLOCS.with(|c| c.set(0));
    ARMED.with(|a| a.set(true));
    let r = f();
    ARMED.with(|a| a.set(false));
    (r, ALLOCS.with(std::cell::Cell::get), BYTES.with(std::cell::Cell::get))
}