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},
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_inplace<T>(
148        &mut self,
149        f: impl FnOnce(RefMut<MaybeUninit<T>>) -> Result<RefMut<T>>,
150    ) -> Result<OwnedGlobalPtr<T, ArenaAllocator>> {
151        let reserve = self
152            .base_mut()
153            .reserve(Layout::new::<T>())
154            .map_err(|_| AllocError)?;
155        let gp = GlobalPtr::<u8>::new(self.id(), reserve);
156        let res = gp.cast::<MaybeUninit<T>>();
157        let res = unsafe { res.resolve_mut() };
158        let gp = f(res)?.global();
159        Ok(unsafe { OwnedGlobalPtr::from_global(gp.cast(), self.allocator()) })
160    }
161
162    pub fn allocator(&self) -> ArenaAllocator {
163        ArenaAllocator {
164            ptr: GlobalPtr::new(self.id(), NULLPAGE_SIZE as u64),
165        }
166    }
167}
168
169impl AsRef<ObjectHandle> for ArenaObject {
170    fn as_ref(&self) -> &ObjectHandle {
171        self.obj.handle()
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use crate::object::ObjectBuilder;
179
180    #[test]
181    fn test_arena_object_new() {
182        let builder = ObjectBuilder::default();
183        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
184
185        // Verify the object was created successfully
186        assert!(arena.object().handle().id() != ObjID::new(0));
187    }
188
189    #[test]
190    fn test_arena_allocator() {
191        let builder = ObjectBuilder::default();
192        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
193        let allocator = arena.allocator();
194
195        // Test basic allocation
196        let layout = Layout::new::<u64>();
197        let ptr = allocator.alloc(layout).expect("Failed to allocate");
198
199        // Verify the pointer is valid
200        assert!(ptr.offset() >= NULLPAGE_SIZE as u64 * 2);
201    }
202
203    #[test]
204    fn test_arena_alloc_value() {
205        let builder = ObjectBuilder::default();
206        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
207
208        let value = 42u64;
209        let owned_ptr = arena.alloc(value).expect("Failed to allocate value");
210
211        // Verify the allocated value
212        let resolved = { owned_ptr.resolve() };
213        assert_eq!(*resolved, 42u64);
214    }
215
216    #[test]
217    fn test_arena_alloc_inplace() {
218        let builder = ObjectBuilder::default();
219        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
220
221        let owned_ptr = arena
222            .alloc_inplace(|uninit| Ok(uninit.write(100u32)))
223            .expect("Failed to allocate in place");
224
225        // Verify the allocated value
226        let resolved = { owned_ptr.resolve() };
227        assert_eq!(*resolved, 100u32);
228    }
229
230    #[test]
231    fn test_arena_multiple_allocations() {
232        let builder = ObjectBuilder::default();
233        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
234
235        // Allocate multiple values
236        let ptr1 = arena.alloc(1u64).expect("Failed to allocate first value");
237        let ptr2 = arena.alloc(2u64).expect("Failed to allocate second value");
238        let ptr3 = arena.alloc(3u64).expect("Failed to allocate third value");
239
240        // Verify all values are correct and pointers are different
241        let val1 = { ptr1.resolve() };
242        let val2 = { ptr2.resolve() };
243        let val3 = { ptr3.resolve() };
244
245        assert_eq!(*val1, 1u64);
246        assert_eq!(*val2, 2u64);
247        assert_eq!(*val3, 3u64);
248
249        // Verify pointers are at different offsets
250        assert_ne!(ptr1.offset(), ptr2.offset());
251        assert_ne!(ptr2.offset(), ptr3.offset());
252        assert_ne!(ptr1.offset(), ptr3.offset());
253    }
254
255    #[test]
256    fn test_arena_tx_object() {
257        let builder = ObjectBuilder::default();
258        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
259
260        let mut tx_obj = arena.as_tx().expect("Failed to create tx object");
261        let owned_ptr = tx_obj.alloc(999u64).expect("Failed to allocate in tx");
262
263        // Verify the allocated value
264        let resolved = { owned_ptr.resolve() };
265        assert_eq!(*resolved, 999u64);
266    }
267
268    #[test]
269    fn test_arena_alignment() {
270        let builder = ObjectBuilder::default();
271        let arena = ArenaObject::new(builder).expect("Failed to create ArenaObject");
272
273        // Allocate values with different alignments
274        let ptr1 = arena.alloc(1u8).expect("Failed to allocate u8");
275        let ptr2 = arena.alloc(2u64).expect("Failed to allocate u64");
276
277        // Verify alignment requirements are met
278        assert_eq!(ptr1.offset() % ArenaBase::MIN_ALIGN as u64, 0);
279        assert_eq!(ptr2.offset() % ArenaBase::MIN_ALIGN as u64, 0);
280    }
281
282    #[test]
283    fn test_arena_from_objid() {
284        let builder = ObjectBuilder::default();
285        let arena1 = ArenaObject::new(builder).expect("Failed to create ArenaObject");
286        let obj_id = arena1.object().id();
287
288        // Create a new ArenaObject from the same object ID
289        let arena2 = ArenaObject::from_objid(obj_id).expect("Failed to create from objid");
290
291        // Verify they reference the same object
292        assert_eq!(arena1.object().id(), arena2.object().id());
293    }
294}
twizzler/alloc/arena.rs

twizzler/alloc/
arena.rs
