use luau_common::ByteSlice;
use luau_printf::Arg;
use crate::VmResult;
use crate::native::{NativeCallContext, NativeCallResult};
use crate::thread::{LuaStringBuilder, LuaStringBuilderStorage, Thread};
use super::{L_ESC, digit};
const FLAGS: &[u8] = b"-+ #0";
const MAX_ITEM: usize = 512;
const MAX_FORMAT: usize = 32;
struct ScannedFormat {
index: usize,
len: usize,
item_size: usize,
has_precision: bool,
}
fn add_quoted(thread: &Thread, buffer: &mut LuaStringBuilder<'_, '_>, argument: i32) -> VmResult {
let bytes = unsafe { thread.check_string(argument)? };
let _ = unsafe { buffer.reserve(bytes.len() + 2)? };
unsafe { buffer.push_byte(b'"')? };
for &byte in bytes.as_bytes() {
match byte {
b'"' | b'\\' | b'\n' => unsafe {
buffer.push_byte(b'\\')?;
buffer.push_byte(byte)?;
},
b'\r' => unsafe { buffer.push_bytes(b"\\r")? },
b'\0' => unsafe { buffer.push_bytes(b"\\000")? },
_ => unsafe { buffer.push_byte(byte)? },
}
}
unsafe { buffer.push_byte(b'"')? };
Ok(())
}
fn scan_format(
thread: &Thread,
format: &[u8],
mut index: usize,
form: &mut [u8; MAX_FORMAT],
) -> VmResult<ScannedFormat> {
let start = index;
while index < format.len() && FLAGS.contains(&format[index]) {
index += 1;
}
if index - start >= FLAGS.len() {
return unsafe { crate::error!(thread, "invalid format (repeated flags)") }
.map_err(Into::into);
}
if format.get(index).is_some_and(|byte| digit(*byte)) {
index += 1;
}
if format.get(index).is_some_and(|byte| digit(*byte)) {
index += 1;
}
let has_precision = format.get(index) == Some(&b'.');
if format.get(index) == Some(&b'.') {
index += 1;
if format.get(index).is_some_and(|byte| digit(*byte)) {
index += 1;
}
if format.get(index).is_some_and(|byte| digit(*byte)) {
index += 1;
}
}
if format.get(index).is_some_and(|byte| digit(*byte)) {
return unsafe { crate::error!(thread, "invalid format (width or precision too long)") }
.map_err(Into::into);
}
form[0] = b'%';
if index >= format.len() {
return Ok(ScannedFormat {
index,
len: 1,
item_size: 0,
has_precision,
});
}
let item_size = index - start + 1;
let len = item_size + 1;
form[1..len].copy_from_slice(&format[start..=index]);
Ok(ScannedFormat {
index,
len,
item_size,
has_precision,
})
}
fn add_int64_format(form: &mut [u8; MAX_FORMAT], indicator: u8, item_size: usize) -> usize {
debug_assert!(item_size + 3 <= MAX_FORMAT);
debug_assert_eq!(form[0], b'%');
debug_assert_eq!(form[item_size], indicator);
form[item_size] = b'l';
form[item_size + 1] = b'l';
form[item_size + 2] = indicator;
item_size + 3
}
unsafe fn add_formatted_item(
thread: &Thread,
buffer: &mut LuaStringBuilder<'_, '_>,
form: &[u8],
args: &mut [Arg<'_>],
) -> VmResult {
let mut item = [0u8; MAX_ITEM - 1];
let count = match luau_printf::printf_c_locale_to_slice(
&mut item,
luau_printf::BStr::new(form),
args,
) {
Ok(count) => count,
Err(_) => {
return unsafe {
crate::error!(
thread,
"invalid option '%%%c' to 'format'",
form.last().copied().unwrap_or_default() as i32
)
}
.map_err(Into::into);
}
};
let stored = count.min(item.len());
let len = item[..stored]
.iter()
.position(|byte| *byte == 0)
.unwrap_or(stored);
unsafe { buffer.push_bytes(&item[..len])? };
Ok(())
}
unsafe fn add_counted_formatted_item(
thread: &Thread,
buffer: &mut LuaStringBuilder<'_, '_>,
form: &[u8],
args: &mut [Arg<'_>],
) -> VmResult {
let mut item = [0u8; MAX_ITEM - 1];
let count = match luau_printf::printf_c_locale_to_slice(
&mut item,
luau_printf::BStr::new(form),
args,
) {
Ok(count) => count,
Err(_) => {
return unsafe {
crate::error!(
thread,
"invalid option '%%%c' to 'format'",
form.last().copied().unwrap_or_default() as i32
)
}
.map_err(Into::into);
}
};
let stored = count.min(item.len());
unsafe { buffer.push_bytes(&item[..stored])? };
Ok(())
}
pub(super) fn string_format(ctx: NativeCallContext) -> NativeCallResult {
let thread = ctx.raw_thread();
unsafe {
let top = thread.get_top();
let mut argument = 1;
let format = thread.check_string(argument)?;
let mut index = 0usize;
let mut buffer_storage = LuaStringBuilderStorage::uninit();
let mut buffer = LuaStringBuilder::new(thread, &mut buffer_storage);
while index < format.len() {
if format[index] != L_ESC {
buffer.push_byte(format[index])?;
index += 1;
} else if index + 1 < format.len() && format[index + 1] == L_ESC {
buffer.push_byte(format[index + 1])?;
index += 2;
} else if index + 1 < format.len() && format[index + 1] == b'*' {
index += 2;
argument += 1;
if argument > top {
return crate::error!(thread, "missing argument #%d", argument)
.map_err(Into::into);
}
buffer.push_any_value(argument)?;
} else {
let mut form = [0u8; MAX_FORMAT];
argument += 1;
if argument > top {
return crate::error!(thread, "missing argument #%d", argument)
.map_err(Into::into);
}
let scanned = scan_format(thread, format, index + 1, &mut form)?;
index = scanned.index;
let indicator = format.get(index).copied().unwrap_or(0);
index = index.saturating_add(1);
match indicator {
b'c' => {
add_counted_formatted_item(
thread,
&mut buffer,
&form[..scanned.len],
&mut [Arg::int(thread.check_number(argument)? as i32)],
)?;
continue;
}
b'd' | b'i' => {
let value = if let Some(value) = thread.to_integer64(argument) {
value
} else {
thread.check_number(argument)? as i64
};
let form_len = add_int64_format(&mut form, indicator, scanned.item_size);
add_formatted_item(
thread,
&mut buffer,
&form[..form_len],
&mut [Arg::sint(value)],
)?;
continue;
}
b'o' | b'u' | b'x' | b'X' => {
let value = if let Some(value) = thread.to_integer64(argument) {
value as u64
} else {
let number = thread.check_number(argument)?;
if number < 0.0 {
(number as i64) as u64
} else {
number as u64
}
};
let form_len = add_int64_format(&mut form, indicator, scanned.item_size);
add_formatted_item(
thread,
&mut buffer,
&form[..form_len],
&mut [Arg::uint(value)],
)?;
continue;
}
b'e' | b'E' | b'f' | b'g' | b'G' => {
add_formatted_item(
thread,
&mut buffer,
&form[..scanned.len],
&mut [Arg::float(thread.check_number(argument)?)],
)?;
continue;
}
b'q' => {
add_quoted(thread, &mut buffer, argument)?;
continue;
}
b's' => {
let string = thread.check_string(argument)?;
if scanned.len == 2 || (!scanned.has_precision && string.len() >= 100) {
buffer.push_bytes(string)?;
continue;
}
add_formatted_item(
thread,
&mut buffer,
&form[..scanned.len],
&mut [Arg::string(string.as_bstr())],
)?;
continue;
}
b'*' => {
return crate::error!(thread, "'%*' does not take a form")
.map_err(Into::into);
}
_ => {
return crate::error!(
thread,
"invalid option '%%%c' to 'format'",
indicator as i32
)
.map_err(Into::into);
}
}
}
}
buffer.finish()?;
Ok(1)
}
}