Skip to main content

blue_lang_runtime/
alloc.rs

1//! How many allocations a stretch of a blue program made, and how many bytes.
2//!
3//! `blue profile` reports time per definition; this is the other half of what
4//! a definition costs (`theory/BLUE-MEMORY.md` §VI M0: you cannot configure
5//! what you cannot account for). One counter, read at two points: a caller
6//! snapshots [`allocations`] when an activation enters and again when it
7//! leaves, and the difference is what that activation allocated, everything it
8//! called included. Subtracting what its callees allocated gives the
9//! exclusive figure, exactly as the profiler does with time.
10//!
11//! [`Counting`] is a `GlobalAlloc` over `System` that counts each allocation
12//! on the thread that made it. A binary opts in with one line:
13//!
14//! ```ignore
15//! #[global_allocator]
16//! static ALLOC: blue_lang_runtime::alloc::Counting = blue_lang_runtime::alloc::Counting;
17//! ```
18//!
19//! The count is per thread so a profile of one interpreter is not charged for
20//! another thread's work, and it costs two thread-local additions per
21//! allocation, nothing on a free. Until a binary installs [`Counting`],
22//! [`allocations`] answers `None`: a program that was never counted reads as
23//! unmeasured, never as having allocated nothing.
24
25use std::alloc::{GlobalAlloc, Layout, System};
26use std::cell::Cell;
27use std::sync::atomic::{AtomicBool, Ordering};
28
29/// Allocations made and bytes requested on one thread, from its start.
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct Allocations {
32    /// Allocation calls. A `realloc` that grows counts as one.
33    pub count: u64,
34    /// Bytes requested. A growing `realloc` counts the bytes it added.
35    pub bytes: u64,
36}
37
38impl Allocations {
39    /// What was allocated between `earlier` and `self`, on one thread.
40    #[must_use]
41    pub fn since(self, earlier: Self) -> Self {
42        Self {
43            count: self.count.saturating_sub(earlier.count),
44            bytes: self.bytes.saturating_sub(earlier.bytes),
45        }
46    }
47}
48
49impl std::ops::AddAssign for Allocations {
50    fn add_assign(&mut self, other: Self) {
51        self.count += other.count;
52        self.bytes += other.bytes;
53    }
54}
55
56/// Set by the first allocation [`Counting`] serves, so [`allocations`] can tell
57/// "not counted" from "counted zero".
58static INSTALLED: AtomicBool = AtomicBool::new(false);
59thread_local! {
60    // `const` initialisation: no lazy init, so no allocation from inside the
61    // allocator, and no destructor to run during thread teardown.
62    static COUNT: Cell<u64> = const { Cell::new(0) };
63    static BYTES: Cell<u64> = const { Cell::new(0) };
64}
65
66/// This thread's running totals, or `None` when [`Counting`] is not the global
67/// allocator of this binary.
68#[must_use]
69pub fn allocations() -> Option<Allocations> {
70    if !INSTALLED.load(Ordering::Relaxed) {
71        return None;
72    }
73    Some(Allocations {
74        count: COUNT.try_with(Cell::get).unwrap_or(0),
75        bytes: BYTES.try_with(Cell::get).unwrap_or(0),
76    })
77}
78
79fn charge(bytes: usize) {
80    // `try_with`: an allocation made while this thread's locals are being torn
81    // down is served and simply not counted.
82    let _ = COUNT.try_with(|c| c.set(c.get() + 1));
83    let _ = BYTES.try_with(|b| b.set(b.get() + bytes as u64));
84}
85
86/// The system allocator, counting what each thread allocates. See the module
87/// documentation for how a binary installs it.
88pub struct Counting;
89
90// SAFETY: every method forwards to `System` with the caller's arguments
91// unchanged; the counting touches only `Cell`s and an atomic flag, neither of
92// which allocates.
93unsafe impl GlobalAlloc for Counting {
94    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
95        INSTALLED.store(true, Ordering::Relaxed);
96        charge(layout.size());
97        unsafe { System.alloc(layout) }
98    }
99
100    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
101        INSTALLED.store(true, Ordering::Relaxed);
102        charge(layout.size());
103        unsafe { System.alloc_zeroed(layout) }
104    }
105
106    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
107        unsafe { System.dealloc(ptr, layout) }
108    }
109
110    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
111        if new_size > layout.size() {
112            charge(new_size - layout.size());
113        }
114        unsafe { System.realloc(ptr, layout, new_size) }
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    /// This crate's own test binary does not install [`Counting`], so the
123    /// answer must be "unmeasured", not a zero that reads as "allocated
124    /// nothing". `tests/alloc_count.rs` installs it and measures.
125    #[test]
126    fn an_uncounted_binary_reads_as_unmeasured() {
127        let _ = vec![1u8; 64];
128        assert_eq!(allocations(), None);
129    }
130
131    #[test]
132    fn since_is_the_difference_and_never_underflows() {
133        let a = Allocations {
134            count: 3,
135            bytes: 100,
136        };
137        let b = Allocations {
138            count: 5,
139            bytes: 160,
140        };
141        assert_eq!(
142            b.since(a),
143            Allocations {
144                count: 2,
145                bytes: 60
146            }
147        );
148        assert_eq!(a.since(b), Allocations::default());
149    }
150}