twizzler/alloc/
arena.rs

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

twizzler/alloc/
arena.rs
