moonpool_assertions/
region.rs1use 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
19const 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 static HEAP_OWNED: Cell<bool> = const { Cell::new(false) };
30}
31
32fn table_layout() -> Layout {
33 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 Layout::from_size_align(EACH_BUCKET_MEM_SIZE, REGION_ALIGN)
41 .expect("each-bucket layout: const size, power-of-two align")
42}
43
44#[must_use]
46pub fn assertion_table_ptr() -> *mut u8 {
47 ASSERTION_TABLE.with(Cell::get)
48}
49
50#[must_use]
52pub fn each_bucket_ptr() -> *mut u8 {
53 EACH_BUCKET_PTR.with(Cell::get)
54}
55
56pub fn init() {
59 if !assertion_table_ptr().is_null() {
60 return;
61 }
62 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
77pub 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
92pub 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 unsafe { dealloc(table, table_layout()) };
109 }
110 let buckets = each_bucket_ptr();
111 if !buckets.is_null() {
112 unsafe { dealloc(buckets, bucket_layout()) };
114 }
115 HEAP_OWNED.with(|c| c.set(false));
116}
117
118pub fn reset() {
120 let table = assertion_table_ptr();
121 if !table.is_null() {
122 unsafe { std::ptr::write_bytes(table, 0, ASSERTION_TABLE_MEM_SIZE) };
124 }
125 let buckets = each_bucket_ptr();
126 if !buckets.is_null() {
127 unsafe { std::ptr::write_bytes(buckets, 0, EACH_BUCKET_MEM_SIZE) };
129 }
130}
131
132pub fn prepare_next_seed_reset() {
139 let table = assertion_table_ptr();
140 if !table.is_null() {
141 unsafe {
144 let count_ptr = table.cast::<()>().cast::<AtomicU32>();
145 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 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 unsafe {
165 let count_ptr = buckets.cast::<()>().cast::<AtomicU32>();
166 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 assert!(assertion_table_ptr().is_null());
187 init();
188 assert!(!assertion_table_ptr().is_null());
189
190 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();
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 let pass_after: u64 = after.iter().map(|s| s.pass_count).sum();
224 assert_eq!(pass_total, pass_after);
225 clear();
226 }
227}