use std::rc::Rc;
use super::StdLib;
use crate::error::{LuaError, Result};
use crate::vm::{BinOp, FuncDef, Oper, UnOp, Value, ValueType, VM};
use crate::Error;
pub(super) fn module(stdlib: &mut StdLib) -> Result<()> {
stdlib
.module("table")
.func("concat", concat)?
.func("insert", insert)?
.func("move", r#move)?
.func("pack", pack)?
.func("remove", remove)?
.func("sort", sort)?
.func("unpack", unpack)?;
Ok(())
}
fn check_table(vm: &mut VM, value: &Value, field: &str) -> bool {
if value.to_table().is_ok() {
true
} else if let Some(meta) = vm.get_metatable(value) {
meta.borrow().get(&Value::str(field)) != Value::Nil
} else {
false
}
}
fn concat(vm: &mut VM) -> Result<Value> {
let list = vm.arg(0)?;
if !check_table(vm, &list, "__index") || !check_table(vm, &list, "__len") {
return err!(LuaError::ExpectedType(ValueType::Table, list.value_type()));
}
let sep = match vm.arg_or_nil(1) {
Value::Nil => Vec::new(),
v => v.to_string_coerce()?,
};
let i = match vm.arg_or_nil(2) {
Value::Nil => 1,
v => v.to_number_coerce()?.coerce_int()?,
};
let j = match vm.arg_or_nil(3) {
Value::Nil => vm
.un_op(Oper::Raw(list.clone()), UnOp::Len)?
.to_int_coerce()?,
v => v.to_int_coerce()?,
};
let mut out = Vec::new();
for ind in i..=j {
if ind > i {
out.extend(&sep);
}
let elem = vm.get_table_field(list.clone(), Value::int(ind))?;
if elem == Value::Nil {
return err!(LuaError::TableConcat(ind));
}
out.extend(elem.to_string_coerce()?);
}
Ok(Value::str_bytes(out))
}
fn insert(vm: &mut VM) -> Result<Value> {
let list = vm.arg(0)?;
if !check_table(vm, &list, "__index")
|| !check_table(vm, &list, "__newindex")
|| !check_table(vm, &list, "__len")
{
return err!(LuaError::ExpectedType(ValueType::Table, list.value_type()));
} else if vm.arg_len() > 3 {
return err!(LuaError::ArgumentCount);
}
let border = vm
.un_op(Oper::Raw(list.clone()), UnOp::Len)?
.to_number_coerce()
.and_then(|n| n.coerce_int())
.map_err(|e| e.map_lua_error(|_| LuaError::TableLength))?;
let (pos, value) = match vm.arg_or_nil(2) {
Value::Nil => (border.wrapping_add(1), vm.arg(1)?),
value => {
let pos = vm.arg_int_coerce(1)?;
if pos <= 0 || pos > border + 1 {
return err!(LuaError::TableIndex);
}
for i in (pos..=border).rev() {
let elem = vm.get_table_field(list.clone(), Value::int(i))?;
vm.set_table_field(list.clone(), Value::int(i + 1), elem)?;
}
(pos, value)
}
};
vm.set_table_field(list, Value::int(pos), value)?;
Ok(Value::empty())
}
fn r#move(vm: &mut VM) -> Result<Value> {
let a1 = vm.arg(0)?;
if !check_table(vm, &a1, "__index") || !check_table(vm, &a1, "__newindex") {
return err!(LuaError::ExpectedType(ValueType::Table, a1.value_type()));
}
let a2 = match vm.arg_or_nil(4) {
Value::Nil => a1.clone(),
v => v,
};
if !check_table(vm, &a2, "__index") || !check_table(vm, &a2, "__newindex") {
return err!(LuaError::ExpectedType(ValueType::Table, a2.value_type()));
}
let f = vm.arg_int_coerce(1)?;
let e = vm.arg_int_coerce(2)?;
let t = vm.arg_int_coerce(3)?;
let (n, over) = e.overflowing_sub(f);
if f > e {
return Ok(a2);
} else if n >= i64::MAX || (over && n < 0) {
return err!(LuaError::TableMoveCount);
}
let n = e - f + 1;
if t.checked_add(n - 1).is_none() {
return err!(LuaError::TableMoveWrap);
}
if t > e || t <= f {
for i in 0..n {
let elem = vm.get_table_field(a1.clone(), Value::int(f + i))?;
vm.set_table_field(a2.clone(), Value::int(t + i as i64), elem)?;
}
} else {
for i in (0..n).rev() {
let elem = vm.get_table_field(a1.clone(), Value::int(f + i))?;
vm.set_table_field(a2.clone(), Value::int(t + i as i64), elem)?;
}
}
Ok(a2)
}
fn pack(vm: &mut VM) -> Result<Value> {
let values = vm.arg_split(0);
let table = vm.alloc_table();
{
let mut table = table.borrow_mut();
table.set(Value::str("n"), Value::int(values.len() as i64))?;
for (i, val) in values.into_iter().enumerate() {
table.set(Value::int(i as i64 + 1), val)?;
}
}
Ok(Value::Table(table))
}
fn remove(vm: &mut VM) -> Result<Value> {
let list = vm.arg(0)?;
if !check_table(vm, &list, "__index")
|| !check_table(vm, &list, "__newindex")
|| !check_table(vm, &list, "__len")
{
return err!(LuaError::ExpectedType(ValueType::Table, list.value_type()));
}
let border = vm
.un_op(Oper::Raw(list.clone()), UnOp::Len)?
.to_number_coerce()?
.coerce_int()?;
let pos = match vm.arg_or_nil(1) {
Value::Nil => Value::int(border),
v => v,
};
Ok(match pos.to_number_coerce().and_then(|n| n.coerce_int()) {
Ok(mut pos) => {
if pos != border && (pos <= 0 || pos > border + 1) {
return err!(LuaError::TableIndex);
}
let removed = vm.get_table_field(list.clone(), Value::int(pos))?;
while pos < border {
let value = vm.get_table_field(list.clone(), Value::int(pos + 1))?;
vm.set_table_field(list.clone(), Value::int(pos), value)?;
pos += 1;
}
vm.set_table_field(list.clone(), Value::int(pos), Value::Nil)?;
removed
}
_ => {
let value = vm.get_table_field(list.clone(), pos.clone())?;
vm.set_table_field(list, pos, Value::Nil)?;
value
}
})
}
fn sort(vm: &mut VM) -> Result<Value> {
let list = vm.arg(0)?;
if !check_table(vm, &list, "__index")
|| !check_table(vm, &list, "__newindex")
|| !check_table(vm, &list, "__len")
{
return Err(vm.arg_error(
0,
Error::from_lua(LuaError::ExpectedType(ValueType::Table, list.value_type())),
));
}
let comp = match vm.arg_or_nil(1) {
Value::Nil => None,
v => Some(v.to_func()?),
};
let border = vm
.un_op(Oper::Raw(list.clone()), UnOp::Len)?
.to_number_coerce()?
.coerce_int()?;
if border >= i64::MAX {
return err!(LuaError::TableSortSize);
}
_quicksort(vm, list, 1, border, &comp)?;
Ok(Value::empty())
}
fn _quicksort(
vm: &mut VM,
list: Value,
low: i64,
high: i64,
comp: &Option<Rc<FuncDef>>,
) -> Result<()> {
if low >= 1 && high >= 1 && low < high {
let pivot = _quicksort_partition(vm, list.clone(), low, high, comp)?;
_quicksort(vm, list.clone(), low, pivot, comp)?;
_quicksort(vm, list, pivot + 1, high, comp)?;
}
Ok(())
}
fn _quicksort_partition(
vm: &mut VM,
list: Value,
low: i64,
high: i64,
comp: &Option<Rc<FuncDef>>,
) -> Result<i64> {
let pivot = vm.get_table_field(list.clone(), Value::int((high + low) / 2))?;
let mut i = low - 1;
let mut j = high + 1;
loop {
let a = loop {
i += 1;
let a = vm.get_table_field(list.clone(), Value::int(i))?;
if !_quicksort_comp(vm, a.clone(), pivot.clone(), comp)? {
break a;
}
if i == high {
return err!(LuaError::TableSortFunc);
}
};
let b = loop {
j -= 1;
let b = vm.get_table_field(list.clone(), Value::int(j))?;
if !_quicksort_comp(vm, pivot.clone(), b.clone(), comp)? {
break b;
}
if j < i {
return err!(LuaError::TableSortFunc);
}
};
if i >= j {
return Ok(j);
}
vm.set_table_field(list.clone(), Value::int(i), b)?;
vm.set_table_field(list.clone(), Value::int(j), a)?;
}
}
fn _quicksort_comp(vm: &mut VM, a: Value, b: Value, comp: &Option<Rc<FuncDef>>) -> Result<bool> {
match comp {
Some(func) => vm
.call_recursive(func.clone(), vec![a, b])
.map(|v| v.is_truthy()),
None => vm
.bin_op(Oper::Raw(a), BinOp::Lt, Oper::Raw(b))
.map(|v| v.is_truthy()),
}
}
const TABLE_UNPACK_MAX: i64 = 999_999;
fn unpack(vm: &mut VM) -> Result<Value> {
let list = vm.arg(0)?;
if !check_table(vm, &list, "__index") || !check_table(vm, &list, "__len") {
return err!(LuaError::ExpectedType(ValueType::Table, list.value_type()));
}
let i = match vm.arg_or_nil(1) {
Value::Nil => 1,
v => v.to_int_coerce()?,
};
let j = match vm.arg_or_nil(2) {
Value::Nil => vm
.un_op(Oper::Raw(list.clone()), UnOp::Len)?
.to_int_coerce()
.map_err(|e| e.map_lua_error(|_| LuaError::TableLength))?,
v => v.to_number_coerce()?.coerce_int()?,
};
let (count, over) = j.overflowing_sub(i);
if count >= TABLE_UNPACK_MAX || (over && count < 0) {
return err!(LuaError::TableUnpack);
} else if over {
return Ok(Value::empty());
}
let mut out = Vec::new();
for i in i..=j {
out.push(vm.get_table_field(list.clone(), Value::int(i))?);
}
Ok(Value::Mult(out))
}