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

mlua/
memory.rs
