tenant.openagents/omega

1use std::{
2    alloc::{self, handle_alloc_error},
3    cell::Cell,
4    num::NonZeroUsize,
5    ops::{Deref, DerefMut},
6    ptr::{self, NonNull},
7    rc::Rc,
8};
9
10struct ArenaElement {
11    value: *mut u8,
12    drop: unsafe fn(*mut u8),
13}
14
15impl Drop for ArenaElement {
16    #[inline(always)]
17    fn drop(&mut self) {
18        unsafe { (self.drop)(self.value) };
19    }
20}
21
22struct Chunk {
23    start: *mut u8,
24    end: *mut u8,
25    offset: *mut u8,
26}
27
28impl Drop for Chunk {
29    fn drop(&mut self) {
30        unsafe {
31            let chunk_size = self.end.offset_from_unsigned(self.start);
32            // SAFETY: This succeeded during allocation.
33            let layout = alloc::Layout::from_size_align_unchecked(chunk_size, 1);
34            alloc::dealloc(self.start, layout);
35        }
36    }
37}
38
39impl Chunk {
40    fn new(chunk_size: NonZeroUsize) -> Self {
41        // this only fails if chunk_size is unreasonably huge
42        let layout = alloc::Layout::from_size_align(chunk_size.get(), 1).unwrap();
43        let start = unsafe { alloc::alloc(layout) };
44        if start.is_null() {
45            handle_alloc_error(layout);
46        }
47        let end = unsafe { start.add(chunk_size.get()) };
48        Self {
49            start,
50            end,
51            offset: start,
52        }
53    }
54
55    fn allocate(&mut self, layout: alloc::Layout) -> Option<NonNull<u8>> {
56        // Compute the allocation bounds in integer address space so that the
57        // bounds check happens before any pointer offsetting. Offsetting a
58        // pointer past the end of its allocation is undefined behavior even if
59        // the result is never dereferenced (as happens on the chunk-spill
60        // path), so `ptr::add` cannot be used until we know the result stays
61        // in bounds. `checked_add` also handles the documented case where
62        // `align_offset` returns `usize::MAX`.
63        let base = self.offset.addr();
64        let aligned_addr = base.checked_add(self.offset.align_offset(layout.align()))?;
65        let next_addr = aligned_addr.checked_add(layout.size())?;
66
67        if next_addr <= self.end.addr() {
68            let aligned = self.offset.with_addr(aligned_addr);
69            self.offset = self.offset.with_addr(next_addr);
70            NonNull::new(aligned)
71        } else {
72            None
73        }
74    }
75
76    fn reset(&mut self) {
77        self.offset = self.start;
78    }
79}
80
81pub struct Arena {
82    chunks: Vec<Chunk>,
83    elements: Vec<ArenaElement>,
84    valid: Rc<Cell<bool>>,
85    current_chunk_index: usize,
86    chunk_size: NonZeroUsize,
87    scope_depth: usize,
88}
89
90impl Drop for Arena {
91    fn drop(&mut self) {
92        self.force_clear();
93    }
94}
95
96impl Arena {
97    pub fn new(chunk_size: usize) -> Self {
98        let chunk_size = NonZeroUsize::try_from(chunk_size).unwrap();
99        Self {
100            chunks: vec![Chunk::new(chunk_size)],
101            elements: Vec::new(),
102            valid: Rc::new(Cell::new(true)),
103            current_chunk_index: 0,
104            chunk_size,
105            scope_depth: 0,
106        }
107    }
108
109    pub fn capacity(&self) -> usize {
110        self.chunks.len() * self.chunk_size.get()
111    }
112
113    /// Marks the start of a scope (e.g. a window draw) whose allocations must stay
114    /// live until the scope ends, even if `clear` is called by a nested scope in
115    /// the meantime.
116    pub fn begin_scope(&mut self) {
117        self.scope_depth += 1;
118    }
119
120    /// Ends the innermost scope started with `begin_scope`.
121    ///
122    /// Panics if no scope is active: an unbalanced `end_scope` would let `clear`
123    /// run while an enclosing scope still references arena memory, which is
124    /// exactly the use-after-free this bookkeeping exists to prevent, so failing
125    /// loudly here is preferable.
126    pub fn end_scope(&mut self) {
127        self.scope_depth = self
128            .scope_depth
129            .checked_sub(1)
130            .expect("Arena::end_scope called without a matching begin_scope");
131    }
132
133    /// Drops all allocations and resets the arena, unless a scope is still active.
134    ///
135    /// When a draw triggers a nested draw (e.g. re-entrant window procedure
136    /// invocations on Windows, or opening a window from within a draw), the nested
137    /// draw's clear must not free memory the outer draw still references, so it is
138    /// deferred: the outer draw's own clear will drop both draws' allocations.
139    pub fn clear(&mut self) {
140        if self.scope_depth == 0 {
141            self.force_clear();
142        } else {
143            log::debug!(
144                "deferring arena clear; {} enclosing scope(s) still active",
145                self.scope_depth
146            );
147        }
148    }
149
150    fn force_clear(&mut self) {
151        self.valid.set(false);
152        self.valid = Rc::new(Cell::new(true));
153        self.elements.clear();
154        for chunk_index in 0..=self.current_chunk_index {
155            self.chunks[chunk_index].reset();
156        }
157        self.current_chunk_index = 0;
158    }
159
160    #[inline(always)]
161    pub fn alloc<T>(&mut self, f: impl FnOnce() -> T) -> ArenaBox<T> {
162        #[inline(always)]
163        unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
164        where
165            F: FnOnce() -> T,
166        {
167            unsafe { ptr::write(ptr, f()) };
168        }
169
170        unsafe fn drop<T>(ptr: *mut u8) {
171            unsafe { std::ptr::drop_in_place(ptr.cast::<T>()) };
172        }
173
174        let layout = alloc::Layout::new::<T>();
175        let mut current_chunk = &mut self.chunks[self.current_chunk_index];
176        let ptr = if let Some(ptr) = current_chunk.allocate(layout) {
177            ptr.as_ptr()
178        } else {
179            self.current_chunk_index += 1;
180            if self.current_chunk_index >= self.chunks.len() {
181                self.chunks.push(Chunk::new(self.chunk_size));
182                assert_eq!(self.current_chunk_index, self.chunks.len() - 1);
183                log::trace!(
184                    "increased element arena capacity to {}kb",
185                    self.capacity() / 1024,
186                );
187            }
188            current_chunk = &mut self.chunks[self.current_chunk_index];
189            if let Some(ptr) = current_chunk.allocate(layout) {
190                ptr.as_ptr()
191            } else {
192                panic!(
193                    "Arena chunk_size of {} is too small to allocate {} bytes",
194                    self.chunk_size,
195                    layout.size()
196                );
197            }
198        };
199
200        unsafe { inner_writer(ptr.cast(), f) };
201        self.elements.push(ArenaElement {
202            value: ptr,
203            drop: drop::<T>,
204        });
205
206        ArenaBox {
207            ptr: ptr.cast(),
208            valid: self.valid.clone(),
209        }
210    }
211}
212
213pub struct ArenaBox<T: ?Sized> {
214    ptr: *mut T,
215    valid: Rc<Cell<bool>>,
216}
217
218impl<T: ?Sized> ArenaBox<T> {
219    #[inline(always)]
220    pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaBox<U> {
221        ArenaBox {
222            ptr: f(&mut self),
223            valid: self.valid,
224        }
225    }
226
227    #[track_caller]
228    fn validate(&self) {
229        assert!(
230            self.valid.get(),
231            "attempted to dereference an ArenaRef after its Arena was cleared"
232        );
233    }
234}
235
236impl<T: ?Sized> Deref for ArenaBox<T> {
237    type Target = T;
238
239    #[inline(always)]
240    fn deref(&self) -> &Self::Target {
241        self.validate();
242        unsafe { &*self.ptr }
243    }
244}
245
246impl<T: ?Sized> DerefMut for ArenaBox<T> {
247    #[inline(always)]
248    fn deref_mut(&mut self) -> &mut Self::Target {
249        self.validate();
250        unsafe { &mut *self.ptr }
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use std::{cell::Cell, rc::Rc};
257
258    use super::*;
259
260    #[test]
261    fn test_arena() {
262        let mut arena = Arena::new(1024);
263        let a = arena.alloc(|| 1u64);
264        let b = arena.alloc(|| 2u32);
265        let c = arena.alloc(|| 3u16);
266        let d = arena.alloc(|| 4u8);
267        assert_eq!(*a, 1);
268        assert_eq!(*b, 2);
269        assert_eq!(*c, 3);
270        assert_eq!(*d, 4);
271
272        arena.clear();
273        let a = arena.alloc(|| 5u64);
274        let b = arena.alloc(|| 6u32);
275        let c = arena.alloc(|| 7u16);
276        let d = arena.alloc(|| 8u8);
277        assert_eq!(*a, 5);
278        assert_eq!(*b, 6);
279        assert_eq!(*c, 7);
280        assert_eq!(*d, 8);
281
282        // Ensure drop gets called.
283        let dropped = Rc::new(Cell::new(false));
284        struct DropGuard(Rc<Cell<bool>>);
285        impl Drop for DropGuard {
286            fn drop(&mut self) {
287                self.0.set(true);
288            }
289        }
290        arena.alloc(|| DropGuard(dropped.clone()));
291        arena.clear();
292        assert!(dropped.get());
293    }
294
295    #[test]
296    fn test_arena_grow() {
297        let mut arena = Arena::new(8);
298        arena.alloc(|| 1u64);
299        arena.alloc(|| 2u64);
300
301        assert_eq!(arena.capacity(), 16);
302
303        arena.alloc(|| 3u32);
304        arena.alloc(|| 4u32);
305
306        assert_eq!(arena.capacity(), 24);
307    }
308
309    #[test]
310    fn test_arena_alignment() {
311        let mut arena = Arena::new(256);
312        let x1 = arena.alloc(|| 1u8);
313        let x2 = arena.alloc(|| 2u16);
314        let x3 = arena.alloc(|| 3u32);
315        let x4 = arena.alloc(|| 4u64);
316        let x5 = arena.alloc(|| 5u64);
317
318        assert_eq!(*x1, 1);
319        assert_eq!(*x2, 2);
320        assert_eq!(*x3, 3);
321        assert_eq!(*x4, 4);
322        assert_eq!(*x5, 5);
323
324        assert_eq!(x1.ptr.align_offset(std::mem::align_of_val(&*x1)), 0);
325        assert_eq!(x2.ptr.align_offset(std::mem::align_of_val(&*x2)), 0);
326    }
327
328    #[test]
329    #[should_panic(expected = "attempted to dereference an ArenaRef after its Arena was cleared")]
330    fn test_arena_use_after_clear() {
331        let mut arena = Arena::new(16);
332        let value = arena.alloc(|| 1u64);
333
334        arena.clear();
335        let _read_value = *value;
336    }
337
338    #[test]
339    fn test_clear_deferred_while_scope_active() {
340        struct DropCounter(Rc<Cell<usize>>);
341        impl Drop for DropCounter {
342            fn drop(&mut self) {
343                self.0.set(self.0.get() + 1);
344            }
345        }
346
347        let drops = Rc::new(Cell::new(0));
348        let mut arena = Arena::new(1024);
349
350        // Outer draw starts and allocates.
351        arena.begin_scope();
352        let outer = arena.alloc(|| 42u64);
353        arena.alloc({
354            let drops = drops.clone();
355            || DropCounter(drops)
356        });
357
358        // Nested draw runs to completion and requests a clear.
359        arena.begin_scope();
360        let inner = arena.alloc(|| 7u64);
361        arena.alloc({
362            let drops = drops.clone();
363            || DropCounter(drops)
364        });
365        arena.end_scope();
366        arena.clear();
367
368        // The clear must be deferred: the outer draw's allocations are still live.
369        assert_eq!(*outer, 42);
370        assert_eq!(*inner, 7);
371        assert_eq!(drops.get(), 0);
372
373        // Once the outer draw finishes, its clear drops both draws' allocations.
374        arena.end_scope();
375        arena.clear();
376        assert_eq!(drops.get(), 2);
377    }
378
379    #[test]
380    fn test_clear_without_scope_is_immediate() {
381        let mut arena = Arena::new(1024);
382        let value = arena.alloc(|| 1u64);
383        assert_eq!(*value, 1);
384        arena.clear();
385        assert!(!value.valid.get());
386    }
387
388    #[test]
389    #[should_panic(expected = "Arena::end_scope called without a matching begin_scope")]
390    fn test_unbalanced_end_scope_panics() {
391        let mut arena = Arena::new(1024);
392        arena.begin_scope();
393        arena.end_scope();
394        arena.end_scope();
395    }
396}
397