Skip to main content

luau_vm/memory/
allocator.rs

1use core::alloc::Layout;
2use core::mem::align_of;
3use core::ptr::NonNull;
4
5use crate::gc::RawGcObject;
6use crate::memory::RawLuaPage;
7use crate::state::{RawCallInfo, RawLuaState, RawMainState};
8use crate::value::RawTValue;
9
10pub const VM_ALLOC_ALIGN: usize = {
11    let mut align = align_of::<usize>();
12    if align_of::<RawMainState>() > align {
13        align = align_of::<RawMainState>();
14    }
15    if align_of::<RawLuaPage>() > align {
16        align = align_of::<RawLuaPage>();
17    }
18    if align_of::<RawLuaState>() > align {
19        align = align_of::<RawLuaState>();
20    }
21    if align_of::<RawTValue>() > align {
22        align = align_of::<RawTValue>();
23    }
24    if align_of::<RawCallInfo>() > align {
25        align = align_of::<RawCallInfo>();
26    }
27    if align_of::<RawGcObject>() > align {
28        align = align_of::<RawGcObject>();
29    }
30    align
31};
32
33/// Allocates memory for a Luau state.
34///
35/// The VM passes exact Rust layouts for every allocation. Implementations must
36/// return pointers that satisfy the requested layout and must deallocate or
37/// reallocate pointers using the corresponding previous layout. Reallocation to
38/// an equal or smaller size must not fail.
39///
40/// # Safety
41///
42/// Returned pointers must be valid for reads and writes of `layout.size()`
43/// bytes, aligned to `layout.align()`, and remain allocated until passed back
44/// to `reallocate` or `deallocate`. `reallocate` must preserve the first
45/// `old_layout.size().min(new_layout.size())` bytes on success, and must
46/// return `Some` when `new_layout.size() <= old_layout.size()`.
47#[allow(
48    clippy::missing_safety_doc,
49    reason = "allocator methods share the unsafe trait contract above"
50)]
51pub unsafe trait LuaAllocator {
52    unsafe fn allocate(&self, layout: Layout) -> Option<NonNull<u8>>;
53
54    unsafe fn reallocate(
55        &self,
56        ptr: NonNull<u8>,
57        old_layout: Layout,
58        new_layout: Layout,
59    ) -> Option<NonNull<u8>>;
60
61    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout);
62}
63
64#[derive(Debug, Default, Clone, Copy)]
65pub struct SystemLuaAllocator;
66
67unsafe impl LuaAllocator for SystemLuaAllocator {
68    unsafe fn allocate(&self, layout: Layout) -> Option<NonNull<u8>> {
69        debug_assert!(layout.size() > 0);
70        unsafe { NonNull::new(std::alloc::alloc(layout)) }
71    }
72
73    unsafe fn reallocate(
74        &self,
75        ptr: NonNull<u8>,
76        old_layout: Layout,
77        new_layout: Layout,
78    ) -> Option<NonNull<u8>> {
79        debug_assert!(old_layout.size() > 0);
80        debug_assert!(new_layout.size() > 0);
81
82        if old_layout.align() == new_layout.align() {
83            unsafe {
84                NonNull::new(std::alloc::realloc(
85                    ptr.as_ptr(),
86                    old_layout,
87                    new_layout.size(),
88                ))
89            }
90        } else {
91            let new_ptr = unsafe { self.allocate(new_layout)? };
92            unsafe {
93                core::ptr::copy_nonoverlapping(
94                    ptr.as_ptr(),
95                    new_ptr.as_ptr(),
96                    old_layout.size().min(new_layout.size()),
97                );
98                self.deallocate(ptr, old_layout);
99            }
100            Some(new_ptr)
101        }
102    }
103
104    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
105        debug_assert!(layout.size() > 0);
106        unsafe { std::alloc::dealloc(ptr.as_ptr(), layout) };
107    }
108}
109
110pub(crate) enum VmAllocator {
111    System(SystemLuaAllocator),
112    Custom(Box<dyn LuaAllocator>),
113}
114
115impl VmAllocator {
116    pub(crate) fn system() -> Self {
117        Self::System(SystemLuaAllocator)
118    }
119
120    pub(crate) fn custom<A: LuaAllocator + 'static>(allocator: A) -> Self {
121        Self::Custom(Box::new(allocator))
122    }
123
124    pub(crate) unsafe fn allocate(&self, layout: Layout) -> Option<NonNull<u8>> {
125        match &self {
126            Self::System(allocator) => unsafe { allocator.allocate(layout) },
127            Self::Custom(allocator) => unsafe { allocator.allocate(layout) },
128        }
129    }
130
131    pub(crate) unsafe fn reallocate(
132        &self,
133        ptr: NonNull<u8>,
134        old_layout: Layout,
135        new_layout: Layout,
136    ) -> Option<NonNull<u8>> {
137        match &self {
138            Self::System(allocator) => unsafe { allocator.reallocate(ptr, old_layout, new_layout) },
139            Self::Custom(allocator) => unsafe { allocator.reallocate(ptr, old_layout, new_layout) },
140        }
141    }
142
143    pub(crate) unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
144        match &self {
145            Self::System(allocator) => unsafe { allocator.deallocate(ptr, layout) },
146            Self::Custom(allocator) => unsafe { allocator.deallocate(ptr, layout) },
147        }
148    }
149}