mlua_codemp_patch/
memory.rs

1use std::alloc::{self, Layout};
2use std::os::raw::c_void;
3use std::ptr;
4
5pub(crate) static ALLOCATOR: ffi::lua_Alloc = allocator;
6
7#[repr(C)]
8#[derive(Default)]
9pub(crate) struct MemoryState {
10    used_memory: isize,
11    memory_limit: isize,
12    // Can be set to temporary ignore the memory limit.
13    // This is used when calling `lua_pushcfunction` for lua5.1/jit/luau.
14    ignore_limit: bool,
15    // Indicates that the memory limit was reached on the last allocation.
16    #[cfg(feature = "luau")]
17    limit_reached: bool,
18}
19
20impl MemoryState {
21    #[inline]
22    pub(crate) unsafe fn get(state: *mut ffi::lua_State) -> *mut Self {
23        let mut mem_state = ptr::null_mut();
24        #[cfg(feature = "luau")]
25        {
26            ffi::lua_getallocf(state, &mut mem_state);
27            mlua_assert!(!mem_state.is_null(), "Luau state has no allocator userdata");
28        }
29        #[cfg(not(feature = "luau"))]
30        if ffi::lua_getallocf(state, &mut mem_state) != ALLOCATOR {
31            mem_state = ptr::null_mut();
32        }
33        mem_state as *mut MemoryState
34    }
35
36    #[inline]
37    pub(crate) fn used_memory(&self) -> usize {
38        self.used_memory as usize
39    }
40
41    #[inline]
42    pub(crate) fn memory_limit(&self) -> usize {
43        self.memory_limit as usize
44    }
45
46    #[inline]
47    pub(crate) fn set_memory_limit(&mut self, limit: usize) -> usize {
48        let prev_limit = self.memory_limit;
49        self.memory_limit = limit as isize;
50        prev_limit as usize
51    }
52
53    // This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit/luau
54    // to bypass the memory limit (if set).
55    #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
56    #[inline]
57    pub(crate) unsafe fn relax_limit_with(state: *mut ffi::lua_State, f: impl FnOnce()) {
58        let mem_state = Self::get(state);
59        if !mem_state.is_null() {
60            (*mem_state).ignore_limit = true;
61            f();
62            (*mem_state).ignore_limit = false;
63        } else {
64            f();
65        }
66    }
67
68    // Does nothing apart from calling `f()`, we don't need to bypass any limits
69    #[cfg(any(feature = "lua52", feature = "lua53", feature = "lua54"))]
70    #[inline]
71    pub(crate) unsafe fn relax_limit_with(_state: *mut ffi::lua_State, f: impl FnOnce()) {
72        f();
73    }
74
75    // Returns `true` if the memory limit was reached on the last memory operation
76    #[cfg(feature = "luau")]
77    #[inline]
78    pub(crate) unsafe fn limit_reached(state: *mut ffi::lua_State) -> bool {
79        (*Self::get(state)).limit_reached
80    }
81}
82
83unsafe extern "C-unwind" fn allocator(
84    extra: *mut c_void,
85    ptr: *mut c_void,
86    osize: usize,
87    nsize: usize,
88) -> *mut c_void {
89    let mem_state = &mut *(extra as *mut MemoryState);
90    #[cfg(feature = "luau")]
91    {
92        // Reset the flag
93        mem_state.limit_reached = false;
94    }
95
96    if nsize == 0 {
97        // Free memory
98        if !ptr.is_null() {
99            let layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN);
100            alloc::dealloc(ptr as *mut u8, layout);
101            mem_state.used_memory -= osize as isize;
102        }
103        return ptr::null_mut();
104    }
105
106    // Do not allocate more than isize::MAX
107    if nsize > isize::MAX as usize {
108        return ptr::null_mut();
109    }
110
111    // Are we fit to the memory limits?
112    let mut mem_diff = nsize as isize;
113    if !ptr.is_null() {
114        mem_diff -= osize as isize;
115    }
116    let mem_limit = mem_state.memory_limit;
117    let new_used_memory = mem_state.used_memory + mem_diff;
118    if mem_limit > 0 && new_used_memory > mem_limit && !mem_state.ignore_limit {
119        #[cfg(feature = "luau")]
120        {
121            mem_state.limit_reached = true;
122        }
123        return ptr::null_mut();
124    }
125    mem_state.used_memory += mem_diff;
126
127    if ptr.is_null() {
128        // Allocate new memory
129        let new_layout = match Layout::from_size_align(nsize, ffi::SYS_MIN_ALIGN) {
130            Ok(layout) => layout,
131            Err(_) => return ptr::null_mut(),
132        };
133        let new_ptr = alloc::alloc(new_layout) as *mut c_void;
134        if new_ptr.is_null() {
135            alloc::handle_alloc_error(new_layout);
136        }
137        return new_ptr;
138    }
139
140    // Reallocate memory
141    let old_layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN);
142    let new_ptr = alloc::realloc(ptr as *mut u8, old_layout, nsize) as *mut c_void;
143    if new_ptr.is_null() {
144        alloc::handle_alloc_error(old_layout);
145    }
146    new_ptr
147}
mlua_codemp_patch/memory.rs

mlua_codemp_patch/
memory.rs
