Skip to main content

moonpool_assertions/
region.rs

1//! Storage for the assertion + each-bucket regions.
2//!
3//! By default the regions are heap-allocated by [`init`] — portable everywhere
4//! (wasm, macOS, Linux), single process, no sharing. An exploration backend that
5//! needs cross-`fork` sharing calls [`install_region`] with `MAP_SHARED` pointers
6//! it owns; this crate then accounts into that memory and [`clear`]s its view when
7//! the backend frees it.
8//!
9//! Pointers are thread-locals set before forking, so forked children inherit them
10//! (the `MAP_SHARED` region is the same physical memory in parent and child).
11
12use std::alloc::{Layout, alloc_zeroed, dealloc};
13use std::cell::Cell;
14use std::sync::atomic::{AtomicU32, Ordering};
15
16use crate::buckets::{EACH_BUCKET_MEM_SIZE, EachBucket, MAX_EACH_BUCKETS};
17use crate::slots::{ASSERTION_TABLE_MEM_SIZE, AssertionSlot, MAX_ASSERTION_SLOTS};
18
19/// Alignment for both regions. `AssertionSlot`/`EachBucket` contain `u64`/`i64`
20/// fields and the layout places the slot array at offset 8, so 8-byte alignment
21/// satisfies every field access.
22const REGION_ALIGN: usize = 8;
23
24thread_local! {
25    static ASSERTION_TABLE: Cell<*mut u8> = const { Cell::new(std::ptr::null_mut()) };
26    static EACH_BUCKET_PTR: Cell<*mut u8> = const { Cell::new(std::ptr::null_mut()) };
27    /// True when the current pointers are heap regions this crate allocated (and
28    /// must free); false when an external backend installed them (it frees).
29    static HEAP_OWNED: Cell<bool> = const { Cell::new(false) };
30}
31
32fn table_layout() -> Layout {
33    // Infallible: size is a compile-time constant and align is a power of two.
34    Layout::from_size_align(ASSERTION_TABLE_MEM_SIZE, REGION_ALIGN)
35        .expect("assertion table layout: const size, power-of-two align")
36}
37
38fn bucket_layout() -> Layout {
39    // Infallible: size is a compile-time constant and align is a power of two.
40    Layout::from_size_align(EACH_BUCKET_MEM_SIZE, REGION_ALIGN)
41        .expect("each-bucket layout: const size, power-of-two align")
42}
43
44/// Get the raw pointer to the assertion table region (null if uninitialized).
45#[must_use]
46pub fn assertion_table_ptr() -> *mut u8 {
47    ASSERTION_TABLE.with(Cell::get)
48}
49
50/// Get the raw pointer to the each-bucket region (null if uninitialized).
51#[must_use]
52pub fn each_bucket_ptr() -> *mut u8 {
53    EACH_BUCKET_PTR.with(Cell::get)
54}
55
56/// Allocate the regions on the heap (zeroed). Idempotent: a no-op if a region is
57/// already present (whether heap-owned or installed by a backend).
58pub fn init() {
59    if !assertion_table_ptr().is_null() {
60        return;
61    }
62    // Safety: layouts have non-zero size; alloc_zeroed returns zeroed memory of the
63    // requested size/alignment, matching the `[count: u32, _pad, slots..]` layout.
64    let table = unsafe { alloc_zeroed(table_layout()) };
65    if table.is_null() {
66        std::alloc::handle_alloc_error(table_layout());
67    }
68    let buckets = unsafe { alloc_zeroed(bucket_layout()) };
69    if buckets.is_null() {
70        std::alloc::handle_alloc_error(bucket_layout());
71    }
72    ASSERTION_TABLE.with(|c| c.set(table));
73    EACH_BUCKET_PTR.with(|c| c.set(buckets));
74    HEAP_OWNED.with(|c| c.set(true));
75}
76
77/// Point accounting at caller-owned regions (e.g. `MAP_SHARED` memory from an
78/// exploration backend). The caller owns the memory and is responsible for
79/// freeing it after calling [`clear`]. Frees any heap regions this crate
80/// previously allocated.
81///
82/// Both pointers must reference at least [`crate::ASSERTION_TABLE_MEM_SIZE`] /
83/// [`crate::EACH_BUCKET_MEM_SIZE`] bytes respectively, and stay valid until
84/// [`clear`] is called.
85pub fn install_region(table: *mut u8, buckets: *mut u8) {
86    free_heap_regions();
87    ASSERTION_TABLE.with(|c| c.set(table));
88    EACH_BUCKET_PTR.with(|c| c.set(buckets));
89    HEAP_OWNED.with(|c| c.set(false));
90}
91
92/// Drop this crate's view of the regions. Frees heap regions it owns; for
93/// installed (external) regions it only nulls the pointers — the backend frees
94/// its own memory.
95pub fn clear() {
96    free_heap_regions();
97    ASSERTION_TABLE.with(|c| c.set(std::ptr::null_mut()));
98    EACH_BUCKET_PTR.with(|c| c.set(std::ptr::null_mut()));
99}
100
101fn free_heap_regions() {
102    if !HEAP_OWNED.with(Cell::get) {
103        return;
104    }
105    let table = assertion_table_ptr();
106    if !table.is_null() {
107        // Safety: heap-owned table was allocated by init() with table_layout().
108        unsafe { dealloc(table, table_layout()) };
109    }
110    let buckets = each_bucket_ptr();
111    if !buckets.is_null() {
112        // Safety: heap-owned buckets were allocated by init() with bucket_layout().
113        unsafe { dealloc(buckets, bucket_layout()) };
114    }
115    HEAP_OWNED.with(|c| c.set(false));
116}
117
118/// Zero both regions for a between-run reset. No-op if not initialized.
119pub fn reset() {
120    let table = assertion_table_ptr();
121    if !table.is_null() {
122        // Safety: region is ASSERTION_TABLE_MEM_SIZE bytes (heap or installed).
123        unsafe { std::ptr::write_bytes(table, 0, ASSERTION_TABLE_MEM_SIZE) };
124    }
125    let buckets = each_bucket_ptr();
126    if !buckets.is_null() {
127        // Safety: region is EACH_BUCKET_MEM_SIZE bytes (heap or installed).
128        unsafe { std::ptr::write_bytes(buckets, 0, EACH_BUCKET_MEM_SIZE) };
129    }
130}
131
132/// Reset only the per-seed split triggers, preserving pass/fail counts and
133/// watermarks/frontiers across seeds (coverage-preserving multi-seed reset).
134///
135/// `pass_count`/`fail_count` accumulate across seeds so contract validation never
136/// reports a false "was never reached" when a seed doesn't reach a guarded
137/// assertion. Forking uses `split_triggered`, not counts. No-op if uninitialized.
138pub fn prepare_next_seed_reset() {
139    let table = assertion_table_ptr();
140    if !table.is_null() {
141        // Safety: region is the assertion table laid out as
142        // `[count: u32, _pad, slots: [AssertionSlot; MAX_ASSERTION_SLOTS]]`.
143        unsafe {
144            let count_ptr = table.cast::<()>().cast::<AtomicU32>();
145            // MAX_ASSERTION_SLOTS is a small const (128), so the cast is lossless.
146            let max_slots = u32::try_from(MAX_ASSERTION_SLOTS).unwrap_or(u32::MAX);
147            let count = (*count_ptr).load(Ordering::Relaxed).min(max_slots) as usize;
148            let base = table.add(8).cast::<()>().cast::<AssertionSlot>();
149            for i in 0..count {
150                let slot = &mut *base.add(i);
151                // Skip tombstones (msg_hash == 0) left by the duplicate-slot race fix.
152                if slot.msg_hash == 0 {
153                    continue;
154                }
155                slot.split_triggered = 0;
156            }
157        }
158    }
159
160    let buckets = each_bucket_ptr();
161    if !buckets.is_null() {
162        // Safety: region is the each-bucket table laid out as
163        // `[count: u32, _pad, buckets: [EachBucket; MAX_EACH_BUCKETS]]`.
164        unsafe {
165            let count_ptr = buckets.cast::<()>().cast::<AtomicU32>();
166            // MAX_EACH_BUCKETS is a small const (256), so the cast is lossless.
167            let max_buckets = u32::try_from(MAX_EACH_BUCKETS).unwrap_or(u32::MAX);
168            let count = (*count_ptr).load(Ordering::Relaxed).min(max_buckets) as usize;
169            let base = buckets.add(8).cast::<()>().cast::<EachBucket>();
170            for i in 0..count {
171                (*base.add(i)).split_triggered = 0;
172            }
173        }
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::slots::{AssertCmp, AssertKind, assertion_bool, assertion_read_all};
181
182    #[test]
183    fn heap_init_enables_accounting_and_reset() {
184        // Each nextest test runs in its own process, so this thread starts with
185        // null regions.
186        assert!(assertion_table_ptr().is_null());
187        init();
188        assert!(!assertion_table_ptr().is_null());
189
190        // Accounting works on the heap region.
191        assertion_bool(AssertKind::Sometimes, true, true, "site_a");
192        let slots = assertion_read_all();
193        assert_eq!(slots.len(), 1);
194        assert_eq!(slots[0].pass_count, 1);
195
196        // reset() zeroes counts.
197        reset();
198        assert!(assertion_read_all().is_empty());
199
200        clear();
201        assert!(assertion_table_ptr().is_null());
202    }
203
204    #[test]
205    fn prepare_next_seed_preserves_counts_resets_triggers() {
206        init();
207        crate::slots::assertion_numeric(
208            AssertKind::NumericSometimes,
209            AssertCmp::Gt,
210            true,
211            5,
212            0,
213            "n",
214        );
215        assertion_bool(AssertKind::Sometimes, true, true, "s");
216        let before = assertion_read_all();
217        let pass_total: u64 = before.iter().map(|s| s.pass_count).sum();
218        assert!(pass_total >= 2);
219
220        prepare_next_seed_reset();
221        let after = assertion_read_all();
222        // Counts preserved across the seed boundary.
223        let pass_after: u64 = after.iter().map(|s| s.pass_count).sum();
224        assert_eq!(pass_total, pass_after);
225        clear();
226    }
227}