use std::ffi::{c_double, c_float, c_int, c_uchar, c_uint, c_ulong, c_ushort};
use std::iter::{repeat, Peekable};
use std::mem::size_of;
use crate::value::{LuaFloat, LuaInt};
use crate::vm::stdlib::string::index_start;
use crate::{LuaError, Result, Value, VM};
struct PackHead {
endian: PackEndian,
align: usize,
total: usize,
}
impl Default for PackHead {
fn default() -> Self {
Self {
endian: PackEndian::Native,
align: 1,
total: 0,
}
}
}
struct PackOpt {
opt: char,
endian: PackEndian,
size: usize,
padding: usize,
}
#[derive(Clone, Copy)]
enum PackEndian {
Big,
Little,
Native,
}
fn parse_opt(
head: &mut PackHead,
fmt: &mut Peekable<impl Iterator<Item = u8>>,
) -> Result<Option<PackOpt>> {
#[inline]
fn num(fmt: &mut Peekable<impl Iterator<Item = u8>>) -> Option<usize> {
let mut num = match fmt.peek().and_then(|c| char::from(*c).to_digit(10)) {
Some(digit) => {
fmt.next();
digit as usize
}
None => return None,
};
while let Some(digit) = fmt.peek().and_then(|c| char::from(*c).to_digit(10)) {
let (mul, o1) = num.overflowing_mul(10);
let (add, o2) = mul.overflowing_add(digit as usize);
if o1 || o2 {
break;
}
fmt.next();
num = add;
}
Some(num)
}
#[inline]
fn num_checked(fmt: &mut Peekable<impl Iterator<Item = u8>>) -> Result<Option<usize>> {
match num(fmt) {
Some(num) => {
if !(1..=16).contains(&num) {
err!(LuaError::PackSizeLimit(num))
} else {
Ok(Some(num))
}
}
None => Ok(None),
}
}
#[inline]
fn opt(
fmt: &mut Peekable<impl Iterator<Item = u8>>,
total: &mut usize,
c: Option<u8>,
endian: PackEndian,
max_align: usize,
empty: bool,
) -> Result<PackOpt> {
let c = c.or_else(|| fmt.next()).unwrap_or(b'-');
let size = match c {
b'b' | b'B' => size_of::<c_uchar>(),
b'h' | b'H' => size_of::<c_ushort>(),
b'l' | b'L' => size_of::<c_ulong>(),
b'j' | b'J' => size_of::<LuaInt>(),
b'T' => size_of::<usize>(),
b'i' | b'I' => num_checked(fmt)?.unwrap_or(size_of::<c_uint>()),
b'f' => size_of::<c_float>(),
b'd' => size_of::<c_double>(),
b'n' => size_of::<LuaFloat>(),
b'c' => match num(fmt) {
Some(s) => s,
None => return err!(LuaError::PackMissingSize),
},
b'z' => 0,
b's' => num_checked(fmt)?.unwrap_or(size_of::<usize>()),
b'x' => 1,
_ => {
return if empty {
err!(LuaError::PackInvalidX)
} else {
err!(LuaError::PackFormatOption(char::from(c)))
}
}
};
if empty && matches!(c, b'c' | b'z') {
return err!(LuaError::PackInvalidX);
}
let padding = if size <= 1 || matches!(c, b'c' | b'z') {
0
} else {
let align = size.min(max_align);
if !align.is_power_of_two() {
return err!(LuaError::PackAlign);
}
(align - (*total & (align - 1))) & (align - 1)
};
let size = if empty { 0 } else { size };
*total += padding + size;
Ok(PackOpt {
opt: char::from(c),
endian,
size,
padding,
})
}
while let Some(c) = fmt.next() {
match c {
b' ' => {}
b'<' => head.endian = PackEndian::Little,
b'>' => head.endian = PackEndian::Big,
b'=' => head.endian = PackEndian::Native,
b'!' => head.align = num_checked(fmt)?.unwrap_or(std::mem::size_of::<usize>()),
b'X' => {
return Ok(Some(opt(
fmt,
&mut head.total,
None,
head.endian,
head.align,
true,
)?))
}
c => {
return Ok(Some(opt(
fmt,
&mut head.total,
Some(c),
head.endian,
head.align,
false,
)?))
}
}
}
Ok(None)
}
pub(super) fn pack(vm: &mut VM) -> Result<Value> {
#[inline]
fn pack_int(packed: &mut Vec<u8>, num: u128, size: usize, endian: PackEndian) {
let bytes = &num.to_le_bytes();
let bytes = bytes.iter().take(size);
match endian {
PackEndian::Big => packed.extend(bytes.rev()),
PackEndian::Little => packed.extend(bytes),
PackEndian::Native => {
if cfg!(target_endian = "big") {
packed.extend(bytes.rev())
} else {
packed.extend(bytes)
}
}
}
}
let fmt = vm.arg_string_coerce(0)?;
let mut fmt = fmt.into_iter().peekable();
let mut head = PackHead::default();
let mut args = vm.arg_split(1).into_iter();
let mut packed = Vec::new();
while let Some(opt) = parse_opt(&mut head, &mut fmt)? {
packed.extend(repeat(b'\0').take(opt.padding));
if opt.size == 0 && !matches!(opt.opt, 'c' | 'z') {
continue; }
let arg = args.next().unwrap_or(Value::Nil);
match opt.opt {
'b' | 'h' | 'l' | 'j' | 'i' => {
let num = arg.to_number_coerce()?.to_int()?;
if opt.size < size_of::<LuaInt>() {
let limit = (1u64 << (opt.size * 8 - 1)) as i64;
if num < -limit || limit <= num {
return err!(LuaError::PackIntOverflow);
}
}
pack_int(&mut packed, num as i128 as u128, opt.size, opt.endian);
}
'B' | 'H' | 'L' | 'J' | 'I' | 'T' => {
let num = arg.to_number_coerce()?.to_int()? as u64; if opt.size < size_of::<LuaInt>() && num >= (1u64 << (opt.size * 8)) {
return err!(LuaError::PackIntOverflow);
}
pack_int(&mut packed, num as u128, opt.size, opt.endian);
}
'f' => {
let num = arg.to_number_coerce()?.to_float() as f32;
packed.extend_from_slice(&match opt.endian {
PackEndian::Big => num.to_be_bytes(),
PackEndian::Little => num.to_le_bytes(),
PackEndian::Native => num.to_ne_bytes(),
});
}
'd' | 'n' => {
let num = arg.to_number_coerce()?.to_float();
packed.extend_from_slice(&match opt.endian {
PackEndian::Big => num.to_be_bytes(),
PackEndian::Little => num.to_le_bytes(),
PackEndian::Native => num.to_ne_bytes(),
});
}
'c' => {
let str = arg.to_string_coerce()?;
let len = str.len();
if len > opt.size {
return err!(LuaError::PackStringC);
}
packed.extend(str);
if len < opt.size {
packed.extend(repeat(b'\0').take(opt.size - len));
}
}
'z' => {
let str = arg.to_string_coerce()?;
if str.contains(&b'\0') {
return err!(LuaError::StringZeros);
}
head.total += str.len() + 1;
packed.extend(str);
packed.push(b'\0');
}
's' => {
let str = arg.to_string_coerce()?;
if opt.size < size_of::<usize>() && str.len() >= (1 << (opt.size * 8)) {
return err!(LuaError::PackStringLength);
}
pack_int(&mut packed, str.len() as u128, opt.size, opt.endian);
head.total += str.len();
packed.extend(str);
}
'x' => packed.push(b'\0'),
_ => unreachable!(),
}
}
Ok(Value::str_bytes(packed))
}
pub(super) fn packsize(vm: &mut VM) -> Result<Value> {
const MAX: usize = if size_of::<isize>() < size_of::<c_int>() {
isize::MAX as usize
} else {
c_int::MAX as usize
};
let fmt = vm.arg_string_coerce(0)?;
let mut fmt = fmt.into_iter().peekable();
let mut head = PackHead::default();
while let Some(opt) = parse_opt(&mut head, &mut fmt)? {
if matches!(opt.opt, 's' | 'z') {
return err!(LuaError::PackSizeVariableLength);
}
let size = opt.size + opt.padding;
if head.total.wrapping_sub(size) > MAX.saturating_sub(size) {
return err!(LuaError::PackSizeTooLarge);
}
}
Ok(Value::int(head.total as i64))
}
pub(super) fn unpack(vm: &mut VM) -> Result<Value> {
#[inline]
fn unpack_int(bytes: &[u8], endian: PackEndian) -> u128 {
#[inline]
fn parse<T: Iterator<Item = u8>>(iter: T) -> u128 {
u128::from_le_bytes(
iter.chain(repeat(b'\0'))
.take(size_of::<u128>())
.collect::<Vec<_>>()
.try_into()
.unwrap(),
)
}
if matches!(endian, PackEndian::Big)
|| (matches!(endian, PackEndian::Native) && cfg!(target_endian = "big"))
{
parse(bytes.iter().copied().rev())
} else {
parse(bytes.iter().copied())
}
}
let s = vm.arg_string_coerce(1)?;
let mut s = s.as_slice();
let mut head = PackHead::default();
match vm.arg_or_nil(2) {
Value::Nil => {}
v => {
let pos = index_start(v, s.len())?;
if pos > s.len() {
return err!(LuaError::UnpackPosition);
}
head.total += pos;
s = &s[pos..];
}
};
let fmt = vm.arg_string_coerce(0)?;
let mut fmt = fmt.into_iter().peekable();
let mut res = Vec::new();
while let Some(opt) = parse_opt(&mut head, &mut fmt)? {
if s.len() < opt.padding + opt.size {
return err!(LuaError::UnpackTooShort);
}
let (bytes, rest) = s[opt.padding..].split_at(opt.size);
s = rest;
if opt.size == 0 && !matches!(opt.opt, 'c' | 'z') {
continue; }
res.push(match opt.opt {
'b' | 'h' | 'l' | 'j' | 'i' => {
let num = unpack_int(bytes, opt.endian);
let num = if opt.size > size_of::<LuaInt>() {
let mask = ((1 << (opt.size * 8 - 1)) - 1) - (-1i64 as u64 as u128);
let neg = (num as i64) < 0;
if (!neg && num & mask > 0) || (neg && num & mask != mask) {
return err!(LuaError::UnpackIntOverflow(opt.size));
}
num
} else {
let mask = 0x1 << (opt.size * 8 - 1);
(num ^ mask).wrapping_sub(mask)
};
Value::int(num as i64)
}
'B' | 'H' | 'L' | 'J' | 'I' | 'T' => {
let num = unpack_int(bytes, opt.endian);
if opt.size > size_of::<LuaInt>() {
let mask = ((1 << (opt.size * 8 - 1)) - 1) - (-1i64 as u64 as u128);
if num & mask > 0 {
return err!(LuaError::UnpackIntOverflow(opt.size));
}
}
Value::int(num as i64)
}
'f' => {
let bytes = <[u8; 4]>::try_from(bytes).unwrap();
Value::float(match opt.endian {
PackEndian::Big => f32::from_be_bytes(bytes),
PackEndian::Little => f32::from_le_bytes(bytes),
PackEndian::Native => f32::from_ne_bytes(bytes),
} as f64)
}
'd' | 'n' => {
let bytes = <[u8; 8]>::try_from(bytes).unwrap();
Value::float(match opt.endian {
PackEndian::Big => f64::from_be_bytes(bytes),
PackEndian::Little => f64::from_le_bytes(bytes),
PackEndian::Native => f64::from_ne_bytes(bytes),
})
}
'c' => Value::str_bytes(bytes.to_vec()),
'z' => match s.iter().position(|c| c == &b'\0') {
Some(idx) => {
let (bytes, rest) = s.split_at(idx);
s = &rest[1..];
let bytes = bytes.to_vec();
head.total += bytes.len() + 1;
Value::str_bytes(bytes)
}
None => return err!(LuaError::UnpackStringZ),
},
's' => {
let num = unpack_int(bytes, opt.endian) as usize;
if num <= s.len() {
let (bytes, rest) = s.split_at(num);
s = rest;
head.total += num;
Value::str_bytes(bytes.to_vec())
} else {
return err!(LuaError::UnpackStringS);
}
}
'x' => continue,
_ => unreachable!(),
});
}
res.push(Value::int((head.total + 1) as i64));
Ok(Value::Mult(res))
}