luau_vm/thread/
string_builder.rs

1use core::mem::MaybeUninit;
2use core::ptr::{self, NonNull};
3
4use luau_common::ByteSlice;
5
6use crate::gc::GcRuntime;
7use crate::state::ThreadState;
8use crate::string::StringRuntime;
9use crate::string::TString;
10use crate::thread::{LUA_BUFFER_SIZE, LUA_TNONE, Thread};
11use crate::types::{LUA_TBOOLEAN, LUA_TINTEGER, LUA_TNIL, LUA_TNUMBER, LUA_TSTRING};
12use crate::{VmErrorResult, VmResult};
13
14pub struct LuaStringBuilderStorage {
15    buffer: [MaybeUninit<u8>; LUA_BUFFER_SIZE],
16}
17
18impl LuaStringBuilderStorage {
19    pub fn uninit() -> Self {
20        let buffer =
21            unsafe { MaybeUninit::<[MaybeUninit<u8>; LUA_BUFFER_SIZE]>::uninit().assume_init() };
22        Self { buffer }
23    }
24
25    fn as_ptr(&self) -> *const u8 {
26        self.buffer.as_ptr().cast()
27    }
28
29    fn as_mut_ptr(&mut self) -> *mut u8 {
30        self.buffer.as_mut_ptr().cast()
31    }
32}
33
34pub struct LuaStringBuilder<'thread, 'storage> {
35    thread: &'thread Thread,
36    storage: Option<TString>,
37    cursor: *mut u8,
38    end: *mut u8,
39    inline: &'storage mut LuaStringBuilderStorage,
40}
41
42impl<'thread, 'storage> LuaStringBuilder<'thread, 'storage> {
43    fn base_ptr(&self) -> *const u8 {
44        match self.storage {
45            Some(storage) => storage.data_ptr(),
46            None => self.inline.as_ptr(),
47        }
48    }
49
50    fn len(&self) -> usize {
51        unsafe { self.cursor.cast_const().offset_from(self.base_ptr()) as usize }
52    }
53
54    fn capacity(&self) -> usize {
55        unsafe { self.end.cast_const().offset_from(self.base_ptr()) as usize }
56    }
57
58    fn available(&self) -> usize {
59        unsafe { self.end.offset_from(self.cursor) as usize }
60    }
61
62    fn bytes(&self) -> &[u8] {
63        unsafe { core::slice::from_raw_parts(self.base_ptr(), self.len()) }
64    }
65
66    /// `getnextbuffersize`
67    fn get_next_buffer_size(&self, desired_size: usize) -> VmErrorResult<usize> {
68        let capacity = self.capacity();
69        let growth = capacity / 2;
70
71        if usize::MAX - growth < capacity {
72            return unsafe { crate::error!(self.thread, "buffer too large") };
73        }
74
75        Ok((capacity + growth).max(desired_size))
76    }
77
78    /// `extendstrbuf`
79    unsafe fn extend_strbuf(
80        &mut self,
81        additional_size: usize,
82        box_loc: i32,
83    ) -> VmErrorResult<*mut u8> {
84        let thread = self.thread;
85        let old_storage = self.storage;
86        let base = self.base_ptr();
87        let len = self.len();
88        let Some(desired_size) = self.capacity().checked_add(additional_size) else {
89            return unsafe { crate::error!(thread, "buffer too large") };
90        };
91        let next_size = self.get_next_buffer_size(desired_size)?;
92        unsafe {
93            let new_storage = thread.buffer_start(next_size)?;
94            ptr::copy_nonoverlapping(base, new_storage.data_mut_ptr(), len);
95
96            if old_storage.is_none() {
97                thread.push_nil()?;
98                thread.insert(box_loc);
99            }
100
101            thread
102                .stack_top()
103                .offset(box_loc as isize)
104                .value_unchecked()
105                .set_string_value(new_storage);
106
107            self.storage = Some(new_storage);
108            let data = new_storage.data_mut_ptr();
109            self.cursor = data.add(len);
110            self.end = data.add(next_size);
111            Ok(self.cursor)
112        }
113    }
114
115    /// `luaL_buffinit`
116    pub unsafe fn new(
117        thread: &'thread Thread,
118        inline: &'storage mut LuaStringBuilderStorage,
119    ) -> Self {
120        let cursor = inline.as_mut_ptr();
121        let end = unsafe { cursor.add(LUA_BUFFER_SIZE) };
122        Self {
123            thread,
124            storage: None,
125            cursor,
126            end,
127            inline,
128        }
129    }
130
131    /// `luaL_prepbuffsize`
132    pub unsafe fn reserve(&mut self, size: usize) -> VmErrorResult<NonNull<u8>> {
133        let result = if self.available() < size {
134            unsafe { self.extend_strbuf(size - self.available(), -1)? }
135        } else {
136            self.cursor
137        };
138
139        Ok(unsafe { NonNull::new_unchecked(result) })
140    }
141
142    /// `luaL_addlstring`
143    pub unsafe fn push_bytes(&mut self, bytes: &[u8]) -> VmErrorResult {
144        let len = bytes.len();
145        if self.available() < len {
146            unsafe {
147                self.extend_strbuf(len - self.available(), -1)?;
148            }
149        }
150
151        unsafe {
152            ptr::copy_nonoverlapping(bytes.as_ptr(), self.cursor, len);
153            self.cursor = self.cursor.add(len);
154        }
155        Ok(())
156    }
157
158    /// `luaL_addchar`
159    #[inline(always)]
160    pub unsafe fn push_byte(&mut self, byte: u8) -> VmErrorResult {
161        if self.cursor == self.end {
162            unsafe {
163                self.extend_strbuf(1, -1)?;
164            }
165        }
166
167        unsafe {
168            self.cursor.write(byte);
169            self.cursor = self.cursor.add(1);
170        }
171        Ok(())
172    }
173
174    /// `luaL_addvalue`
175    pub unsafe fn push_stack_value(&mut self) -> VmErrorResult {
176        let thread = self.thread;
177
178        unsafe {
179            if let Some(bytes) = thread.to_string(-1)? {
180                let bytes = bytes.as_bytes();
181                let len = bytes.len();
182                if self.available() < len {
183                    self.extend_strbuf(len - self.available(), -2)?;
184                }
185
186                ptr::copy_nonoverlapping(bytes.as_ptr(), self.cursor, len);
187                self.cursor = self.cursor.add(len);
188                thread.pop(1);
189            }
190        }
191        Ok(())
192    }
193
194    /// `luaL_addvalueany`
195    pub unsafe fn push_any_value(&mut self, index: i32) -> VmResult {
196        let thread = self.thread;
197
198        unsafe {
199            match thread.type_of(index) {
200                LUA_TNONE => {}
201                LUA_TNIL => self.push_bytes(b"nil")?,
202                LUA_TBOOLEAN => self.push_bytes(if thread.to_boolean(index) != 0 {
203                    b"true"
204                } else {
205                    b"false"
206                })?,
207                LUA_TNUMBER => {
208                    let mut bytes = [0u8; crate::number::LUAI_MAXNUM2STR];
209                    let len = crate::number::num_to_str(
210                        &mut bytes,
211                        thread.to_number(index).unwrap_or(0.0),
212                    );
213                    self.push_bytes(&bytes[..len])?;
214                }
215                LUA_TSTRING => {
216                    let bytes = thread.to_string(index)?.unwrap_unchecked();
217                    self.push_bytes(bytes)?;
218                }
219                LUA_TINTEGER => {
220                    let mut bytes = [0u8; crate::number::LUAI_MAXINT2STR];
221                    let len = crate::number::int_to_str(
222                        &mut bytes,
223                        thread.to_integer64(index).unwrap_or(0),
224                    );
225                    self.push_bytes(&bytes[..len])?;
226                }
227                _ => {
228                    self.push_bytes(thread.lua_to_string(index)?)?;
229                    thread.pop(1);
230                }
231            }
232        }
233        Ok(())
234    }
235
236    /// `luaL_pushresult`
237    pub unsafe fn finish(&mut self) -> VmErrorResult {
238        let thread = self.thread;
239
240        if let Some(storage) = self.storage {
241            unsafe {
242                thread.check_gc()?;
243
244                let result = if self.cursor == self.end {
245                    thread.buffer_finish(storage)?
246                } else {
247                    thread.intern_string(self.bytes().as_bstr())?
248                };
249                thread
250                    .stack_top()
251                    .sub(1)
252                    .value_unchecked()
253                    .set_string_value(result);
254            }
255        } else {
256            unsafe { thread.push_string(self.bytes())? };
257        }
258        Ok(())
259    }
260
261    /// `luaL_pushresultsize`
262    pub unsafe fn finish_with_reserved(&mut self, size: usize) -> VmErrorResult {
263        unsafe {
264            self.cursor = self.cursor.add(size);
265            self.finish()
266        }
267    }
268}
luau_vm/thread/string_builder.rs

luau_vm/thread/
string_builder.rs
