twizzler/alloc/
arena.rs

1use std::{
2    alloc::{AllocError, Layout},
3    mem::MaybeUninit,
4};
5
6use twizzler_abi::object::{ObjID, MAX_SIZE, NULLPAGE_SIZE};
7use twizzler_rt_abi::{
8    error::ResourceError,
9    object::{MapFlags, ObjectHandle},
10};
11
12use super::{Allocator, OwnedGlobalPtr, SingleObjectAllocator};
13use crate::{
14    marker::BaseType,
15    object::{Object, ObjectBuilder, RawObject, TxObject},
16    ptr::{GlobalPtr, RefMut, RefSliceMut},
17    Result,
18};
19
20pub struct ArenaObject {
21    obj: Object<ArenaBase>,
22}
23
24impl ArenaObject {
25    pub fn object(&self) -> &Object<ArenaBase> {
26        &self.obj
27    }
28
29    pub fn from_allocator(alloc: ArenaAllocator) -> Result<Self> {
30        Self::from_objid(alloc.ptr.id())
31    }
32
33    pub fn from_objid(id: ObjID) -> Result<Self> {
34        Ok(Self {
35            obj: Object::map(id, MapFlags::READ | MapFlags::WRITE | MapFlags::PERSIST)?,
36        })
37    }
38
39    pub fn new(builder: ObjectBuilder<ArenaBase>) -> Result<Self> {
40        let obj = builder.build(ArenaBase {
41            next: (NULLPAGE_SIZE * 2) as u64,
42        })?;
43        Ok(Self { obj })
44    }
45
46    pub fn into_tx(self) -> Result<TxObject<ArenaBase>> {
47        self.obj.into_tx()
48    }
49
50    pub fn as_tx(&self) -> Result<TxObject<ArenaBase>> {
51        self.obj.as_tx()
52    }
53
54    pub fn allocator(&self) -> ArenaAllocator {
55        ArenaAllocator {
56            ptr: GlobalPtr::new(self.obj.id(), NULLPAGE_SIZE as u64),
57        }
58    }
59
60    pub fn alloc<T>(&self, value: T) -> Result<OwnedGlobalPtr<T, ArenaAllocator>> {
61        self.alloc_inplace(|p| Ok(p.write(value)))
62    }
63
64    pub fn alloc_inplace<T>(
65        &self,
66        f: impl FnOnce(RefMut<MaybeUninit<T>>) -> Result<RefMut<T>>,
67    ) -> Result<OwnedGlobalPtr<T, ArenaAllocator>> {
68        let gp = self
69            .allocator()
70            .alloc_with(|x| f(x).map_err(|_| AllocError))?;
71        Ok(unsafe { OwnedGlobalPtr::from_global(gp.cast(), self.allocator()) })
72    }
73}
74
75#[derive(Clone, Copy)]
76pub struct ArenaAllocator {
77    ptr: GlobalPtr<ArenaBase>,
78}
79
80impl ArenaAllocator {
81    pub fn new(ptr: GlobalPtr<ArenaBase>) -> Self {
82        Self { ptr }
83    }
84}
85
86impl SingleObjectAllocator for ArenaAllocator {}
87
88#[repr(C)]
89pub struct ArenaBase {
90    next: u64,
91}
92
93impl BaseType for ArenaBase {}
94
95impl ArenaBase {
96    const MIN_ALIGN: usize = 16;
97    fn reserve(&mut self, layout: Layout) -> Result<u64> {
98        let align = std::cmp::max(layout.align(), Self::MIN_ALIGN);
99        let len = std::cmp::max(layout.size(), Self::MIN_ALIGN) as u64;
100        let next_cell = self.next;
101        let next = next_cell.next_multiple_of(align as u64);
102        if next + len > MAX_SIZE as u64 {
103            return Err(ResourceError::OutOfMemory.into());
104        }
105
106        self.next = next + len;
107        Ok(next)
108    }
109}
110
111impl Allocator for ArenaAllocator {
112    fn alloc(
113        &self,
114        layout: std::alloc::Layout,
115    ) -> core::result::Result<GlobalPtr<u8>, std::alloc::AllocError> {
116        // TODO: use try_resolve
117        let mut allocator = unsafe { self.ptr.resolve().into_tx() }.map_err(|_| AllocError)?;
118        let reserve = allocator.reserve(layout).map_err(|_| AllocError)?;
119        let gp = GlobalPtr::new(allocator.handle().id(), reserve);
120        Ok(gp)
121    }
122
123    fn alloc_with<T>(
124        &self,
125        f: impl FnOnce(
126            RefMut<MaybeUninit<T>>,
127        ) -> core::result::Result<RefMut<T>, std::alloc::AllocError>,
128    ) -> core::result::Result<GlobalPtr<u8>, AllocError> {
129        let mut allocator = unsafe { self.ptr.resolve().into_tx() }.map_err(|_| AllocError)?;
130        let reserve = allocator
131            .reserve(Layout::new::<T>())
132            .map_err(|_| AllocError)?;
133        let gp = GlobalPtr::<u8>::new(allocator.handle().id(), reserve);
134        let res = gp.cast::<MaybeUninit<T>>();
135        let res = unsafe { res.resolve_mut() };
136        Ok(f(res)?.global().cast())
137    }
138
139    unsafe fn dealloc(&self, _ptr: GlobalPtr<u8>, _layout: std::alloc::Layout) {}
140}
141
142impl TxObject<ArenaBase> {
143    pub fn alloc<T>(&mut self, value: T) -> Result<OwnedGlobalPtr<T, ArenaAllocator>> {
144        self.alloc_inplace(|p| Ok(p.write(value)))
145    }
146
147    pub fn alloc_with_slice<T: Copy>(
148        &mut self,
149        slice: &[T],
150    ) -> Result<(usize, OwnedGlobalPtr<T, ArenaAllocator>)> {
151        let layout = Layout::array::<T>(slice.len()).unwrap();
152        let reserve = self.base_mut().reserve(layout)?;
153        let gp = GlobalPtr::<T>::new(self.id(), reserve);
154        let res = unsafe { gp.resolve_mut() };
155        let mut slice_alloc = unsafe { RefSliceMut::from_ref(res, slice.len()) };
156        slice_alloc.copy_from_slice(slice);
157        Ok(unsafe {
158            (
159                slice.len(),
160                OwnedGlobalPtr::from_global(gp.cast(), self.allocator()),
161            )
162        })
163    }
164
165    pub fn alloc_inplace<T>(
166        &mut self,
167        f: impl FnOnce(RefMut<MaybeUninit<T>>) -> Result<RefMut<T>>,
168    ) -> Result<OwnedGlobalPtr<T, ArenaAllocator>> {
169        let reserve = self
170            .base_mut()
171            .reserve(Layout::new::<T>())
172            .map_err(|_| AllocError)?;
173        let gp = GlobalPtr::<u8>::new(self.id(), reserve);
174        let res = gp.cast::<MaybeUninit<T>>();
175        let res = unsafe { res.resolve_mut() };
176        let gp = f(res)?.global();
177        Ok(unsafe { OwnedGlobalPtr::from_global(gp.cast(), self.allocator()) })
178    }
179
180    pub fn allocator(&self) -> ArenaAllocator {
181        ArenaAllocator {
182            ptr: GlobalPtr::new(self.id(), NULLPAGE_SIZE as u64),
183        }
184    }
185}
186
187impl AsRef<ObjectHandle> for ArenaObject {
188    fn as_ref(&self) -> &ObjectHandle {
189        self.obj.handle()
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196    use crate::object::ObjectBuilder;
197
198    #[test]
199    fn test_arena_object_new() {
200        let builder = ObjectBuilder::default();
201        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
202
203        // Verify the object was created successfully
204        assert!(arena.object().handle().id() != ObjID::new(0));
205    }
206
207    #[test]
208    fn test_arena_allocator() {
209        let builder = ObjectBuilder::default();
210        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
211        let allocator = arena.allocator();
212
213        // Test basic allocation
214        let layout = Layout::new::<u64>();
215        let ptr = allocator.alloc(layout).expect("Failed to allocate");
216
217        // Verify the pointer is valid
218        assert!(ptr.offset() >= NULLPAGE_SIZE as u64 * 2);
219    }
220
221    #[test]
222    fn test_arena_alloc_value() {
223        let builder = ObjectBuilder::default();
224        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
225
226        let value = 42u64;
227        let owned_ptr = arena.alloc(value).expect("Failed to allocate value");
228
229        // Verify the allocated value
230        let resolved = { owned_ptr.resolve() };
231        assert_eq!(*resolved, 42u64);
232    }
233
234    #[test]
235    fn test_arena_alloc_inplace() {
236        let builder = ObjectBuilder::default();
237        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
238
239        let owned_ptr = arena
240            .alloc_inplace(|uninit| Ok(uninit.write(100u32)))
241            .expect("Failed to allocate in place");
242
243        // Verify the allocated value
244        let resolved = { owned_ptr.resolve() };
245        assert_eq!(*resolved, 100u32);
246    }
247
248    #[test]
249    fn test_arena_multiple_allocations() {
250        let builder = ObjectBuilder::default();
251        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
252
253        // Allocate multiple values
254        let ptr1 = arena.alloc(1u64).expect("Failed to allocate first value");
255        let ptr2 = arena.alloc(2u64).expect("Failed to allocate second value");
256        let ptr3 = arena.alloc(3u64).expect("Failed to allocate third value");
257
258        // Verify all values are correct and pointers are different
259        let val1 = { ptr1.resolve() };
260        let val2 = { ptr2.resolve() };
261        let val3 = { ptr3.resolve() };
262
263        assert_eq!(*val1, 1u64);
264        assert_eq!(*val2, 2u64);
265        assert_eq!(*val3, 3u64);
266
267        // Verify pointers are at different offsets
268        assert_ne!(ptr1.offset(), ptr2.offset());
269        assert_ne!(ptr2.offset(), ptr3.offset());
270        assert_ne!(ptr1.offset(), ptr3.offset());
271    }
272
273    #[test]
274    fn test_arena_tx_object() {
275        let builder = ObjectBuilder::default();
276        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
277
278        let mut tx_obj = arena.as_tx().expect("Failed to create tx object");
279        let owned_ptr = tx_obj.alloc(999u64).expect("Failed to allocate in tx");
280
281        // Verify the allocated value
282        let resolved = { owned_ptr.resolve() };
283        assert_eq!(*resolved, 999u64);
284    }
285
286    #[test]
287    fn test_arena_alignment() {
288        let builder = ObjectBuilder::default();
289        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
290
291        // Allocate values with different alignments
292        let ptr1 = arena.alloc(1u8).expect("Failed to allocate u8");
293        let ptr2 = arena.alloc(2u64).expect("Failed to allocate u64");
294
295        // Verify alignment requirements are met
296        assert_eq!(ptr1.offset() % ArenaBase::MIN_ALIGN as u64, 0);
297        assert_eq!(ptr2.offset() % ArenaBase::MIN_ALIGN as u64, 0);
298    }
299
300    #[test]
301    fn test_arena_from_objid() {
302        let builder = ObjectBuilder::default();
303        let arena1 = ArenaObject::new(builder).expect("Failed to create ArenaObject");
304        let obj_id = arena1.object().id();
305
306        // Create a new ArenaObject from the same object ID
307        let arena2 = ArenaObject::from_objid(obj_id).expect("Failed to create from objid");
308
309        // Verify they reference the same object
310        assert_eq!(arena1.object().id(), arena2.object().id());
311    }
312}
twizzler/alloc/arena.rs

twizzler/alloc/
arena.rs
