virt_arena/
lib.rs

1#[cfg(unix)]
2mod unix;
3#[cfg(windows)]
4mod windows;
5
6use std::{alloc::Layout, mem::MaybeUninit, ptr::NonNull};
7
8#[cfg(unix)]
9type RawArena = unix::VirtArena;
10#[cfg(windows)]
11type RawArena = windows::VirtArena;
12
13/// A memory arena which leverages the virtual memory system
14/// for allocating structures in a single contiguous memory region.
15#[derive(Default)]
16pub struct VirtArena(RawArena);
17
18impl VirtArena {
19    /// Allocates a memory for the given `layout`.
20    pub fn alloc(&self, layout: Layout) -> NonNull<u8> {
21        self.0.alloc(layout)
22    }
23
24    /// Allocates a struct `T` inside the arena and clears its memory to 0.
25    ///
26    /// # Safety
27    /// Look into [std::mem::zeroed] for safety concerns.
28    #[allow(clippy::mut_from_ref)]
29    pub unsafe fn alloc_zeroed<T: Sized>(&self) -> &mut T {
30        let layout = Layout::new::<T>();
31        self.0.alloc_zeroed(layout).cast().as_mut()
32    }
33
34    /// Allocates a slice `[T]` inside the arena and clears its memory to 0.
35    ///
36    /// # Safety
37    /// Look into [std::mem::zeroed] for safety concerns.
38    #[allow(clippy::mut_from_ref)]
39    pub unsafe fn alloc_slice_zeroed<T: Sized>(&self, count: usize) -> &mut [T] {
40        let layout = Layout::array::<T>(count).expect("Failed to read the array layout");
41        let ptr = self.0.alloc_zeroed(layout).cast();
42        std::slice::from_raw_parts_mut(ptr.as_ptr(), count)
43    }
44
45    /// Allocates memory for struct `T`.
46    #[allow(clippy::mut_from_ref)]
47    pub fn alloc_uninit<T: Sized>(&self) -> &mut MaybeUninit<T> {
48        let layout = Layout::new::<T>();
49        unsafe { self.0.alloc(layout).cast().as_mut() }
50    }
51
52    /// Allocates memory for slice `[T]`.
53    #[allow(clippy::mut_from_ref)]
54    pub fn alloc_slice_uninit<T: Sized>(&self, count: usize) -> &mut [MaybeUninit<T>] {
55        let layout = Layout::array::<T>(count).expect("Failed to read the array layout");
56        let ptr = self.0.alloc(layout).cast();
57        unsafe { std::slice::from_raw_parts_mut(ptr.as_ptr(), count) }
58    }
59
60    /// Allocates a struct `T` inside the arena and sets its
61    /// content to the output of `fun`.   
62    #[allow(clippy::mut_from_ref)]
63    pub fn alloc_with<T: Sized>(&self, fun: impl FnOnce() -> T) -> &mut T {
64        let uninit = self.alloc_uninit();
65        uninit.write(fun());
66        unsafe { uninit.assume_init_mut() }
67    }
68
69    /// Allocates a struct `T` inside the arena and moves `val` into the allocation.    
70    #[allow(clippy::mut_from_ref)]
71    pub fn alloc_value<T: Sized>(&self, val: T) -> &mut T {
72        self.alloc_with(move || val)
73    }
74
75    /// Returns the number of bytes currently allocated from the arena.
76    pub fn bytes_used(&self) -> usize {
77        self.0.bytes_used()
78    }
79
80    /// Restes the arena storage, Invalidating all the references allocated.
81    /// This method does not run the destructors! Those need to be run manually.
82    pub fn reset(&mut self) {
83        self.0.reset()
84    }
85}
86
87// We don't use any thread local storage so this should be fine
88unsafe impl Send for VirtArena {}
89
90const VIRT_ALLOC_SIZE: usize = 128 * (1 << 30); // 128 GiB is assumed to be enoght for every use case of this arena
91
92trait VirtArenaRaw {
93    fn bytes_used(&self) -> usize;
94    fn reset(&mut self);
95
96    fn alloc(&self, layout: Layout) -> NonNull<u8>;
97
98    unsafe fn alloc_zeroed(&self, layout: Layout) -> NonNull<u8> {
99        let ptr = self.alloc(layout);
100        ptr.write_bytes(0, layout.size());
101        ptr
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use std::num::NonZeroU32;
108
109    use super::*;
110
111    #[test]
112    fn arena() {
113        let mut arena = VirtArena::default();
114
115        struct Test {
116            thing: Option<NonZeroU32>,
117        }
118
119        let test1 = unsafe { arena.alloc_zeroed::<Test>() };
120        assert!(test1.thing.is_none());
121
122        let test2 = unsafe { arena.alloc_zeroed::<Test>() };
123        assert!(test2.thing.is_none());
124
125        test1.thing = Some(NonZeroU32::new(345).unwrap());
126
127        arena.reset();
128
129        let test1 = unsafe { arena.alloc_uninit::<Test>().assume_init_mut() };
130        assert_eq!(test1.thing.map(|v| v.get()), Some(345));
131
132        let test2 = unsafe { arena.alloc_uninit::<Test>().assume_init_mut() };
133        assert!(test2.thing.is_none());
134    }
135}
virt_arena/lib.rs

virt_arena/
lib.rs
