Skip to main content

luau_vm/string/
intern.rs

1use core::ptr::{self, NonNull};
2
3use luau_common::{BStr, ByteSlice};
4
5use crate::VmErrorResult;
6use crate::gc::{FIXED_BIT, GcObject, bit_mask};
7use crate::handle::RawHandle;
8use crate::handle::sealed::Sealed;
9use crate::memory::{LuaPage, MemoryRuntime};
10use crate::thread::Thread;
11use crate::types::LUA_TSTRING;
12
13pub const MAX_STRING_SIZE: usize = 1 << 30;
14pub const ATOM_UNDEFINED: i16 = i16::MIN;
15
16/// `LUA_MINSTRTABSIZE`
17pub(crate) const LUA_MIN_STRING_TABLE_SIZE: usize = 32;
18
19/// `luaS_hash`
20pub fn hash(bytes: &[u8]) -> u32 {
21    let mut a = 0u32;
22    let mut b = 0u32;
23    let mut hash = bytes.len() as u32;
24    let mut cursor = bytes;
25    let mut len = bytes.len();
26
27    while len >= 32 {
28        let block0 = u32::from_ne_bytes(cursor[..4].try_into().unwrap());
29        let block1 = u32::from_ne_bytes(cursor[4..8].try_into().unwrap());
30        let block2 = u32::from_ne_bytes(cursor[8..12].try_into().unwrap());
31
32        a = a.wrapping_add(block0);
33        b = b.wrapping_add(block1);
34        hash = hash.wrapping_add(block2);
35
36        a ^= hash;
37        a = a.wrapping_sub(hash.rotate_right(14));
38        b ^= a;
39        b = b.wrapping_sub(a.rotate_right(11));
40        hash ^= b;
41        hash = hash.wrapping_sub(b.rotate_right(25));
42
43        cursor = &cursor[12..];
44        len -= 12;
45    }
46
47    for index in (0..len).rev() {
48        let byte = cursor[index];
49        hash ^= hash
50            .wrapping_shl(5)
51            .wrapping_add(hash.wrapping_shr(2))
52            .wrapping_add(u32::from(byte));
53    }
54
55    hash
56}
57
58#[repr(C)]
59pub struct StringTable {
60    pub hash: *mut *mut RawTString,
61    pub n_use: u32,
62    pub size: i32,
63}
64
65#[repr(C)]
66pub struct RawTString {
67    pub tt: u8,
68    pub marked: u8,
69    pub memcat: u8,
70    pub atom: i16,
71    pub next: *mut RawTString,
72    pub hash: u32,
73    pub len: u32,
74    pub data: [u8; 0],
75}
76
77#[derive(Clone, Copy, PartialEq, Eq)]
78#[repr(transparent)]
79pub struct TString {
80    raw: NonNull<RawTString>,
81}
82
83/// Unstable interned-string allocation capability.
84///
85/// # Safety
86///
87/// The thread and all string/page handles must be live and belong to the same
88/// VM. Sizes and pages must match their allocations, and callers must preserve
89/// interning-table, atom, rooting, and GC invariants.
90#[allow(
91    clippy::missing_safety_doc,
92    reason = "all methods share the capability-level safety contract"
93)]
94pub trait StringRuntime: Sealed {
95    /// `luaS_resize`
96    unsafe fn resize(&self, new_size: i32) -> VmErrorResult;
97
98    /// `luaS_newlstr`
99    unsafe fn intern_string(&self, bytes: &BStr) -> VmErrorResult<TString>;
100
101    /// `luaS_free`
102    unsafe fn free_string(&self, string: TString, page: LuaPage);
103
104    /// `luaS_updateatom`
105    unsafe fn update_atom(&self, string: TString);
106
107    /// `luaS_bufstart`
108    unsafe fn buffer_start(&self, size: usize) -> VmErrorResult<TString>;
109
110    /// `luaS_buffinish`
111    unsafe fn buffer_finish(&self, string: TString) -> VmErrorResult<TString>;
112}
113
114impl TString {
115    /// `luaS_fix`
116    ///
117    /// # Safety
118    ///
119    /// The string must still belong to a live VM, and the caller must ensure
120    /// that no concurrent or reentrant collector operation accesses its mark
121    /// bits while they are updated.
122    pub unsafe fn fix(&self) {
123        unsafe {
124            self.as_ptr().as_mut().unwrap_unchecked().marked |= bit_mask(FIXED_BIT);
125        }
126    }
127
128    /// Constructs a non-owning string handle from a raw VM record.
129    ///
130    /// # Safety
131    ///
132    /// `raw` must address a live `RawTString` owned by the VM in which the
133    /// handle will be used. The handle does not root the string or extend its
134    /// lifetime.
135    pub const unsafe fn from_raw(raw: NonNull<RawTString>) -> Self {
136        Self { raw }
137    }
138
139    /// Constructs a non-owning string handle from a VM string record.
140    ///
141    /// # Safety
142    ///
143    /// `raw` must be a live, VM-owned string record. The returned handle may
144    /// not outlive that record and does not keep it reachable from the GC.
145    pub unsafe fn from_ref(raw: &RawTString) -> Self {
146        Self {
147            raw: NonNull::from(raw),
148        }
149    }
150
151    pub fn atom(&self) -> i16 {
152        unsafe { (*self.as_ptr()).atom }
153    }
154
155    pub fn set_atom(&self, atom: i16) {
156        unsafe {
157            (*self.as_ptr()).atom = atom;
158        }
159    }
160
161    pub const fn size_string(len: usize) -> usize {
162        core::mem::offset_of!(RawTString, data) + len
163    }
164
165    pub fn data_ptr(&self) -> *const u8 {
166        unsafe { (&raw const (*self.as_ptr()).data).cast::<u8>() }
167    }
168
169    /// Returns the address of the string's inline byte storage.
170    ///
171    /// # Safety
172    ///
173    /// The string must be live. Mutating interned string bytes can invalidate
174    /// hashes and table invariants, so writes are only valid during the VM's
175    /// owned string-construction protocol before the string is published.
176    pub unsafe fn data_mut_ptr(&self) -> *mut u8 {
177        unsafe { (&raw mut (*self.as_ptr()).data).cast::<u8>() }
178    }
179
180    /// Borrows the string's bytes for the duration of this handle borrow.
181    ///
182    /// # Safety
183    ///
184    /// The string must remain live and its storage must not be freed for the
185    /// returned borrow. Because this handle does not root the string, the
186    /// caller must also prevent VM transitions that could collect it.
187    pub unsafe fn as_bytes(&self) -> &[u8] {
188        unsafe { core::slice::from_raw_parts(self.data_ptr(), (*self.as_ptr()).len as usize) }
189    }
190
191    /// Borrows the string's bytes as a `BStr`.
192    ///
193    /// # Safety
194    ///
195    /// The string must remain live and its storage must not be freed for the
196    /// returned borrow. Because this handle does not root the string, the
197    /// caller must also prevent VM transitions that could collect it.
198    pub unsafe fn as_bstr(&self) -> &BStr {
199        unsafe { self.as_bytes().as_bstr() }
200    }
201}
202
203impl crate::handle::sealed::Sealed for TString {}
204
205impl RawHandle for TString {
206    type Raw = RawTString;
207
208    fn as_ptr(&self) -> *mut Self::Raw {
209        self.raw.as_ptr()
210    }
211}
212
213impl AsRef<TString> for TString {
214    fn as_ref(&self) -> &TString {
215        self
216    }
217}
218
219fn bucket_index(hash: u32, size: i32) -> usize {
220    debug_assert!(size > 0);
221    debug_assert!(size & (size - 1) == 0);
222    (hash & (size as u32 - 1)) as usize
223}
224
225impl StringRuntime for Thread {
226    /// `luaS_resize`
227    unsafe fn resize(&self, new_size: i32) -> VmErrorResult {
228        unsafe {
229            let new_hash = self.new_array::<*mut RawTString>(new_size as usize, 0)?;
230
231            for index in 0..new_size as usize {
232                new_hash.add(index).write(ptr::null_mut());
233            }
234
235            let global = self.global();
236            let string_table = &global.as_ptr().as_ref().unwrap_unchecked().string_table;
237            let old_hash = string_table.hash;
238            let old_size = string_table.size;
239
240            for index in 0..old_size as usize {
241                let mut string = NonNull::new(old_hash.add(index).read());
242                while let Some(current) = string {
243                    let next = NonNull::new(current.as_ref().next);
244                    let bucket = bucket_index(current.as_ref().hash, new_size);
245
246                    current.as_ptr().as_mut().unwrap_unchecked().next = new_hash.add(bucket).read();
247                    new_hash.add(bucket).write(current.as_ptr());
248
249                    string = next;
250                }
251            }
252
253            if !old_hash.is_null() {
254                self.free_array(old_hash, old_size as usize, 0);
255            }
256
257            let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
258            string_table.size = new_size;
259            string_table.hash = new_hash;
260        }
261        Ok(())
262    }
263
264    /// `luaS_newlstr`
265    unsafe fn intern_string(&self, bytes: &BStr) -> VmErrorResult<TString> {
266        let bytes = bytes.as_bytes();
267        let len = bytes.len();
268        let hash = crate::string::hash(bytes);
269        unsafe {
270            let global = self.global();
271            {
272                let string_table = &global.as_ptr().as_ref().unwrap_unchecked().string_table;
273                debug_assert!(string_table.size > 0);
274
275                let bucket = bucket_index(hash, string_table.size);
276                let mut entry = NonNull::new(string_table.hash.add(bucket).read());
277
278                while let Some(current) = entry {
279                    let current_string = TString::from_raw(current);
280                    if current.as_ref().len as usize == len && current_string.as_bytes() == bytes {
281                        let mut object: GcObject = current_string.into();
282                        if global.is_dead(object) {
283                            object.change_white();
284                        }
285
286                        return Ok(current_string);
287                    }
288
289                    entry = NonNull::new(current.as_ref().next);
290                }
291            }
292
293            if len > MAX_STRING_SIZE {
294                return self.too_big();
295            }
296
297            let active_memcat = self.as_ptr().as_ref().unwrap_unchecked().active_memcat;
298            let string_handle =
299                self.new_gco::<TString>(TString::size_string(len), active_memcat)?;
300            GcObject::from(string_handle).init_header(self, LUA_TSTRING as u8);
301            let string_ref = string_handle.as_ptr().as_mut().unwrap_unchecked();
302            string_ref.atom = ATOM_UNDEFINED;
303            string_ref.hash = hash;
304            string_ref.len = len as u32;
305            ptr::copy_nonoverlapping(bytes.as_ptr(), string_handle.data_mut_ptr(), len);
306            let (should_resize, next_size) = {
307                let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
308                let bucket = bucket_index(hash, string_table.size);
309                string_ref.next = string_table.hash.add(bucket).read();
310                string_table.hash.add(bucket).write(string_handle.as_ptr());
311
312                string_table.n_use += 1;
313                (
314                    string_table.n_use > string_table.size as u32
315                        && string_table.size <= i32::MAX / 2,
316                    string_table.size * 2,
317                )
318            };
319
320            if should_resize {
321                self.resize(next_size)?;
322            }
323
324            Ok(string_handle)
325        }
326    }
327
328    /// `luaS_free`
329    unsafe fn free_string(&self, string: TString, page: LuaPage) {
330        unsafe {
331            let len = string.as_ptr().as_ref().unwrap_unchecked().len as usize;
332            let memcat = string.as_ptr().as_ref().unwrap_unchecked().memcat;
333
334            let global = self.global();
335            let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
336            let bucket = bucket_index(
337                string.as_ptr().as_ref().unwrap_unchecked().hash,
338                string_table.size,
339            );
340            let mut slot = string_table.hash.add(bucket);
341            let mut found = false;
342            let string_raw = string.as_ptr();
343
344            while let Some(current) = NonNull::new(slot.read()) {
345                if current.as_ptr() == string_raw {
346                    slot.write(current.as_ref().next);
347                    found = true;
348                    string_table.n_use -= 1;
349                    break;
350                }
351
352                slot = &raw mut (*current.as_ptr()).next;
353            }
354
355            debug_assert!(found || string.as_ptr().as_ref().unwrap_unchecked().next.is_null());
356
357            self.free_gco(string.into(), TString::size_string(len), memcat, page);
358        }
359    }
360
361    /// `luaS_updateatom`
362    unsafe fn update_atom(&self, string: TString) {
363        unsafe {
364            if string.atom() != ATOM_UNDEFINED {
365                return;
366            }
367
368            let global = self.global();
369            let atom = if let Some(user_atom) = global.user_atom_callback() {
370                user_atom(self, string.as_bstr())
371            } else {
372                -1
373            };
374            string.set_atom(atom);
375        }
376    }
377
378    /// `luaS_bufstart`
379    unsafe fn buffer_start(&self, size: usize) -> VmErrorResult<TString> {
380        if size > MAX_STRING_SIZE {
381            return unsafe { self.too_big() };
382        }
383
384        unsafe {
385            let string = self.new_gco::<TString>(
386                TString::size_string(size),
387                self.as_ptr().as_ref().unwrap_unchecked().active_memcat,
388            )?;
389            GcObject::from(string).init_header(self, LUA_TSTRING as u8);
390            let string_ref = string.as_ptr().as_mut().unwrap_unchecked();
391            string_ref.atom = ATOM_UNDEFINED;
392            string_ref.hash = 0;
393            string_ref.len = size as u32;
394            string_ref.next = ptr::null_mut();
395
396            Ok(string)
397        }
398    }
399
400    /// `luaS_buffinish`
401    unsafe fn buffer_finish(&self, string: TString) -> VmErrorResult<TString> {
402        let bytes = unsafe { string.as_bytes() };
403        let hash = crate::string::hash(bytes);
404        unsafe {
405            let global = self.global();
406            let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
407            let bucket = bucket_index(hash, string_table.size);
408
409            let mut entry = NonNull::new(string_table.hash.add(bucket).read());
410            while let Some(current) = entry {
411                let current_string = TString::from_raw(current);
412                if current_string.as_ptr().as_ref().unwrap_unchecked().len
413                    == string.as_ptr().as_ref().unwrap_unchecked().len
414                    && current_string.as_bytes() == bytes
415                {
416                    let mut object: GcObject = current_string.into();
417                    if global.is_dead(object) {
418                        object.change_white();
419                    }
420
421                    return Ok(current_string);
422                }
423
424                entry = NonNull::new(current.as_ref().next);
425            }
426
427            let string_ref = string.as_ptr().as_mut().unwrap_unchecked();
428            string_ref.hash = hash;
429            string_ref.atom = ATOM_UNDEFINED;
430            string_ref.next = string_table.hash.add(bucket).read();
431            string_table.hash.add(bucket).write(string.as_ptr());
432
433            string_table.n_use += 1;
434            if string_table.n_use > string_table.size as u32 && string_table.size <= i32::MAX / 2 {
435                self.resize(string_table.size * 2)?;
436            }
437
438            Ok(string)
439        }
440    }
441}