use core::mem::MaybeUninit;
use core::ptr::{self, NonNull};
use luau_common::ByteSlice;
use crate::gc::GcRuntime;
use crate::state::ThreadState;
use crate::string::StringRuntime;
use crate::string::TString;
use crate::thread::{LUA_BUFFER_SIZE, LUA_TNONE, Thread};
use crate::types::{LUA_TBOOLEAN, LUA_TINTEGER, LUA_TNIL, LUA_TNUMBER, LUA_TSTRING};
use crate::{VmErrorResult, VmResult};
pub struct LuaStringBuilderStorage {
buffer: [MaybeUninit<u8>; LUA_BUFFER_SIZE],
}
impl LuaStringBuilderStorage {
pub fn uninit() -> Self {
let buffer =
unsafe { MaybeUninit::<[MaybeUninit<u8>; LUA_BUFFER_SIZE]>::uninit().assume_init() };
Self { buffer }
}
fn as_ptr(&self) -> *const u8 {
self.buffer.as_ptr().cast()
}
fn as_mut_ptr(&mut self) -> *mut u8 {
self.buffer.as_mut_ptr().cast()
}
}
pub struct LuaStringBuilder<'thread, 'storage> {
thread: &'thread Thread,
storage: Option<TString>,
cursor: *mut u8,
end: *mut u8,
inline: &'storage mut LuaStringBuilderStorage,
}
impl<'thread, 'storage> LuaStringBuilder<'thread, 'storage> {
fn base_ptr(&self) -> *const u8 {
match self.storage {
Some(storage) => storage.data_ptr(),
None => self.inline.as_ptr(),
}
}
fn len(&self) -> usize {
unsafe { self.cursor.cast_const().offset_from(self.base_ptr()) as usize }
}
fn capacity(&self) -> usize {
unsafe { self.end.cast_const().offset_from(self.base_ptr()) as usize }
}
fn available(&self) -> usize {
unsafe { self.end.offset_from(self.cursor) as usize }
}
fn bytes(&self) -> &[u8] {
unsafe { core::slice::from_raw_parts(self.base_ptr(), self.len()) }
}
fn get_next_buffer_size(&self, desired_size: usize) -> VmErrorResult<usize> {
let capacity = self.capacity();
let growth = capacity / 2;
if usize::MAX - growth < capacity {
return unsafe { crate::error!(self.thread, "buffer too large") };
}
Ok((capacity + growth).max(desired_size))
}
unsafe fn extend_strbuf(
&mut self,
additional_size: usize,
box_loc: i32,
) -> VmErrorResult<*mut u8> {
let thread = self.thread;
let old_storage = self.storage;
let base = self.base_ptr();
let len = self.len();
let Some(desired_size) = self.capacity().checked_add(additional_size) else {
return unsafe { crate::error!(thread, "buffer too large") };
};
let next_size = self.get_next_buffer_size(desired_size)?;
unsafe {
let new_storage = thread.buffer_start(next_size)?;
ptr::copy_nonoverlapping(base, new_storage.data_mut_ptr(), len);
if old_storage.is_none() {
thread.push_nil()?;
thread.insert(box_loc);
}
thread
.stack_top()
.offset(box_loc as isize)
.value_unchecked()
.set_string_value(new_storage);
self.storage = Some(new_storage);
let data = new_storage.data_mut_ptr();
self.cursor = data.add(len);
self.end = data.add(next_size);
Ok(self.cursor)
}
}
pub unsafe fn new(
thread: &'thread Thread,
inline: &'storage mut LuaStringBuilderStorage,
) -> Self {
let cursor = inline.as_mut_ptr();
let end = unsafe { cursor.add(LUA_BUFFER_SIZE) };
Self {
thread,
storage: None,
cursor,
end,
inline,
}
}
pub unsafe fn reserve(&mut self, size: usize) -> VmErrorResult<NonNull<u8>> {
let result = if self.available() < size {
unsafe { self.extend_strbuf(size - self.available(), -1)? }
} else {
self.cursor
};
Ok(unsafe { NonNull::new_unchecked(result) })
}
pub unsafe fn push_bytes(&mut self, bytes: &[u8]) -> VmErrorResult {
let len = bytes.len();
if self.available() < len {
unsafe {
self.extend_strbuf(len - self.available(), -1)?;
}
}
unsafe {
ptr::copy_nonoverlapping(bytes.as_ptr(), self.cursor, len);
self.cursor = self.cursor.add(len);
}
Ok(())
}
#[inline(always)]
pub unsafe fn push_byte(&mut self, byte: u8) -> VmErrorResult {
if self.cursor == self.end {
unsafe {
self.extend_strbuf(1, -1)?;
}
}
unsafe {
self.cursor.write(byte);
self.cursor = self.cursor.add(1);
}
Ok(())
}
pub unsafe fn push_stack_value(&mut self) -> VmErrorResult {
let thread = self.thread;
unsafe {
if let Some(bytes) = thread.to_string(-1)? {
let bytes = bytes.as_bytes();
let len = bytes.len();
if self.available() < len {
self.extend_strbuf(len - self.available(), -2)?;
}
ptr::copy_nonoverlapping(bytes.as_ptr(), self.cursor, len);
self.cursor = self.cursor.add(len);
thread.pop(1);
}
}
Ok(())
}
pub unsafe fn push_any_value(&mut self, index: i32) -> VmResult {
let thread = self.thread;
unsafe {
match thread.type_of(index) {
LUA_TNONE => {}
LUA_TNIL => self.push_bytes(b"nil")?,
LUA_TBOOLEAN => self.push_bytes(if thread.to_boolean(index) != 0 {
b"true"
} else {
b"false"
})?,
LUA_TNUMBER => {
let mut bytes = [0u8; crate::number::LUAI_MAXNUM2STR];
let len = crate::number::num_to_str(
&mut bytes,
thread.to_number(index).unwrap_or(0.0),
);
self.push_bytes(&bytes[..len])?;
}
LUA_TSTRING => {
let bytes = thread.to_string(index)?.unwrap_unchecked();
self.push_bytes(bytes)?;
}
LUA_TINTEGER => {
let mut bytes = [0u8; crate::number::LUAI_MAXINT2STR];
let len = crate::number::int_to_str(
&mut bytes,
thread.to_integer64(index).unwrap_or(0),
);
self.push_bytes(&bytes[..len])?;
}
_ => {
self.push_bytes(thread.lua_to_string(index)?)?;
thread.pop(1);
}
}
}
Ok(())
}
pub unsafe fn finish(&mut self) -> VmErrorResult {
let thread = self.thread;
if let Some(storage) = self.storage {
unsafe {
thread.check_gc()?;
let result = if self.cursor == self.end {
thread.buffer_finish(storage)?
} else {
thread.intern_string(self.bytes().as_bstr())?
};
thread
.stack_top()
.sub(1)
.value_unchecked()
.set_string_value(result);
}
} else {
unsafe { thread.push_string(self.bytes())? };
}
Ok(())
}
pub unsafe fn finish_with_reserved(&mut self, size: usize) -> VmErrorResult {
unsafe {
self.cursor = self.cursor.add(size);
self.finish()
}
}
}