use std::cmp::Ordering;
use std::io::{stderr, Write};
use std::rc::Rc;
use super::StdLib;
use crate::error::{LuaError, LuaErrorSrc, Result};
use crate::lexer::Pos;
use crate::util::read_lua_file;
use crate::value::{CallResult, FuncDef, Numeric, ProtectedHandler, ValueType};
use crate::vm::{calls::ReturnType, gc::GcMode, FuncCont};
use crate::{Error, Value, VM};
pub(super) fn module(stdlib: &mut StdLib) -> Result<()> {
stdlib
.root()
.func("assert", assert)?
.func("collectgarbage", collectgarbage)?
.func_raw("dofile", dofile)?
.func("error", error)?
.func("getmetatable", getmetatable)?
.func("ipairs", ipairs)?
.func("__ipairs_next", __ipairs_next)?
.func("load", load)?
.func("loadfile", loadfile)?
.func("next", next)?
.func_raw("pairs", pairs)?
.func_raw("pcall", pcall)?
.func("print", print)?
.func("rawequal", rawequal)?
.func("rawget", rawget)?
.func("rawlen", rawlen)?
.func("rawset", rawset)?
.func("select", select)?
.func("setmetatable", setmetatable)?
.func("tonumber", tonumber)?
.func("tostring", tostring)?
.func("type", r#type)?
.func("warn", warn)?
.func_raw("xpcall", xpcall)?;
Ok(())
}
fn assert(vm: &mut VM) -> Result<Value> {
let v = vm.arg(0)?;
let error = vm.arg_opt(1);
if v.is_falsy() {
return match error {
Some(v) => err!(LuaError::CustomValue(v)),
None => err!(LuaError::AssertFailed),
};
}
let mut ret = vec![v];
if let Some(error) = error {
ret.push(error);
}
ret.extend(vm.arg_split(2));
Ok(Value::Mult(ret))
}
fn collectgarbage(vm: &mut VM) -> Result<Value> {
if vm.frames.iter().any(|f| matches!(f.meta, Some("__gc"))) {
Ok(Value::Nil)
} else {
let opt = match vm.arg_opt(0) {
Some(opt) => opt.to_string_coerce().map_err(|e| vm.arg_error(0, e))?,
None => "collect".as_bytes().to_vec(),
};
let opt = String::from_utf8_lossy(&opt);
match opt.as_ref() {
"collect" => {
vm.collect_garbage();
Ok(Value::empty())
}
"stop" => {
vm.gc_stopped(Some(true));
Ok(Value::empty())
}
"restart" => {
vm.gc_stopped(Some(false));
Ok(Value::empty())
}
"count" => Ok(Value::int(vm.gc_count() as i64)),
"step" => {
let step = vm.arg_int_coerce(1)?;
Ok(Value::Bool(vm.collect_step(step)))
}
"isrunning" => Ok(Value::Bool(!vm.gc_stopped(None))),
"incremental" => Ok(Value::string(vm.gc_mode(GcMode::Incremental).to_string())),
"generational" => Ok(Value::string(vm.gc_mode(GcMode::Generational).to_string())),
s if s.starts_with("set") => {
let param = s.strip_prefix("set").unwrap();
let val = vm.arg_float_coerce(1)?;
let res = vm.gc_param(param, val)?;
Ok(Value::float(res))
}
_ => Err(vm.arg_error(
0,
Error::from_lua(LuaError::CollectGarbageOption(opt.into_owned())),
)),
}
}
}
fn dofile(vm: &mut VM) -> Result<CallResult> {
let filename = vm.arg_string(0)?;
let filename = String::from_utf8_lossy(&filename).into_owned();
let chunk = read_lua_file(&filename)?;
vm.frames.last_mut().unwrap().varargs = vec![
Value::str_bytes(chunk),
Value::string(format!("@{}", filename)),
];
let func = load(vm)?.to_func()?;
vm.call_prepare(
func,
vec![],
Some((
FuncCont {
name: "dofile_finish",
func: Rc::new(dofile_finish),
},
false,
)),
None,
false,
None,
)
}
fn dofile_finish(_: &mut VM, ret: Result<Value>) -> Result<Value> {
ret
}
fn error(vm: &mut VM) -> Result<Value> {
let message = vm.arg_or_nil(0);
let message_type = message.value_type();
let error = match message {
Value::Number(n) => LuaError::CustomNumber(n),
Value::String(str) => LuaError::CustomString(String::from_utf8_lossy(&*str).into_owned()),
Value::Table(tbl) => match tbl.borrow().get_meta() {
Some(meta) => match meta.borrow().get(&Value::str("__tostring")) {
Value::Nil => LuaError::CustomValue(Value::Table(tbl.clone())),
Value::Func(func) => {
let str = vm
.call_recursive(func, vec![Value::Table(tbl.clone())])?
.to_string_coerce()
.map_err(|err| err.map_lua_error(|_| LuaError::ToStringReturnType))?;
LuaError::CustomString(String::from_utf8_lossy(&str).into_owned())
}
val => LuaError::ExpectedType(ValueType::Func, val.value_type()),
},
None => LuaError::CustomValue(Value::Table(tbl.clone())),
},
v => LuaError::CustomValue(v),
};
let level = match vm.arg_opt(1) {
Some(val) => val.to_int_coerce()?,
None => 1,
};
if message_type == ValueType::String && level > 0 {
let empty = Vec::new();
let mut frames = vm.frames.iter().rev().skip(level as usize - 1);
let pos = match frames.next() {
Some(frame) => match &frame.ret {
ReturnType::Lua(code, pc) | ReturnType::RustLua(_, code, pc) => {
code.get_pos(pc.saturating_sub(1))
}
ReturnType::Rust(..) => Pos::default(),
},
None => Pos::default(),
};
let src = match frames.next() {
Some(frame) => match &*frame.func_def {
FuncDef::Defined(func) => func.source.as_ref(),
_ => &empty,
},
None => &empty,
};
err!(src, pos, error)
} else {
err!(empty, error)
}
}
fn getmetatable(vm: &mut VM) -> Result<Value> {
let object = vm.arg(0)?;
Ok(match vm.get_metatable(&object) {
Some(meta) => {
let inner = meta.borrow().get(&Value::str("__metatable"));
if inner == Value::Nil {
Value::Table(meta)
} else {
inner
}
}
None => Value::Nil,
})
}
fn ipairs(vm: &mut VM) -> Result<Value> {
match vm.arg(0)? {
Value::Nil => err!(LuaError::ArgumentExpected),
t => {
let iter = vm.env.borrow().get(&Value::str("__ipairs_next"));
Ok(Value::Mult(vec![iter, t, Value::int(0)]))
}
}
}
fn __ipairs_next(vm: &mut VM) -> Result<Value> {
let tbl = vm.arg(0)?;
let next = match vm.arg(1)? {
Value::Number(Numeric::Integer(n)) => Value::int(n.wrapping_add(1)),
_ => panic!("invalid state"),
};
Ok(Value::Mult(
match vm.get_table_field(tbl, next.clone())? {
Value::Nil => vec![Value::Nil],
val => vec![next, val],
},
))
}
fn load(vm: &mut VM) -> Result<Value> {
let string = match vm.arg(0)? {
Value::String(str) => str.to_vec(),
Value::Func(func) => {
let mut chunk = Vec::new();
loop {
match vm.call_recursive(func.clone(), vec![]) {
Ok(val) => match val.into_single() {
Value::Nil => break,
Value::String(str) => {
if str.is_empty() {
break;
} else {
chunk.extend_from_slice(str.as_slice())
}
}
_ => {
return Ok(Value::Mult(vec![
Value::Nil,
Value::string(LuaError::LoadReader.to_string()),
]))
}
},
Err(err) => {
return Ok(Value::Mult(vec![
Value::Nil,
Value::string(err.to_string()),
]))
}
}
}
chunk
}
v => return err!(LuaError::ExpectedType(ValueType::String, v.value_type())),
};
let name = match vm.arg_or_nil(1) {
Value::Nil => None,
v => Some(v.to_string_coerce()?),
};
let mode = match vm.arg_or_nil(2) {
Value::Nil => vec![b'b', b't'],
v => v.to_string_coerce()?,
};
let env = match vm.arg_opt(3) {
Some(v) => v,
None => Value::Table(vm.env.clone()),
};
Ok(match vm.load_file(&mode, &string, name, env) {
Ok(func) => Value::Func(func),
Err(e) => Value::Mult(vec![Value::Nil, e.into_value()]),
})
}
pub(super) fn loadfile(vm: &mut VM) -> Result<Value> {
let filename = vm.arg(0)?;
let filename = String::from_utf8_lossy(filename.to_string()?).into_owned();
match read_lua_file(&filename) {
Ok(chunk) => {
let mut args = vm.arg_split(1);
args.insert(0, Value::str_bytes(chunk));
args.insert(1, Value::string(format!("@{}", filename)));
vm.frames.last_mut().unwrap().varargs = args;
load(vm)
}
Err(e) => Ok(e.into()),
}
}
fn next(vm: &mut VM) -> Result<Value> {
let table = vm.arg_table(0)?;
let table = table.borrow();
let next = table.next_key(vm.arg_or_nil(1))?;
let val = table.get(&next);
Ok(Value::Mult(vec![next, val]))
}
fn pairs(vm: &mut VM) -> Result<CallResult> {
let t = vm
.frames
.last()
.unwrap()
.varargs
.get(0)
.cloned()
.unwrap_or_default();
match vm
.get_metatable(&t)
.map(|meta| meta.borrow().get(&Value::str("__pairs")))
.unwrap_or(Value::Nil)
{
Value::Nil => {
if t == Value::Nil {
return err!(LuaError::ArgumentBad(
0,
None,
Box::new(LuaError::ArgumentExpected)
));
}
let iter = vm.env.borrow().get(&Value::str("next"));
let ret = Value::Mult(vec![iter, t, Value::Nil]);
vm.debug_return(&ret)?;
vm.frames.pop().unwrap();
Ok(CallResult::Return(ret))
}
v => {
let func = v.to_func()?;
let args = vm.arg_split(0);
vm.call_prepare(
func,
args,
Some((
FuncCont {
name: "pairs_finish",
func: Rc::new(pairs_finish),
},
false,
)),
None,
false,
None,
)
}
}
}
fn pairs_finish(_: &mut VM, ret: Result<Value>) -> Result<Value> {
ret
}
fn pcall(vm: &mut VM) -> Result<CallResult> {
let args = vm.arg_split(0);
match &args[..] {
[] => err!(LuaError::CallInvalid(ValueType::Nil, LuaErrorSrc::None)),
[func, args @ ..] => {
vm.rust_call_depth += 1;
let func = match func.to_func() {
Ok(func) => func,
Err(_) => {
let ret = pcall_finish(
vm,
err!(LuaError::CallInvalid(func.value_type(), LuaErrorSrc::None,)),
)?;
if let Some(ret_reg) = vm.frames.pop().unwrap().ret_reg {
vm.frames.last_mut().unwrap().regs[ret_reg] = ret;
}
return Ok(CallResult::Continue);
}
};
vm.frames.last_mut().unwrap().protected = Some(
vm.thread
.borrow_mut()
.protected
.replace(ProtectedHandler::PCall),
);
vm.call_prepare(
func,
args.to_vec(),
Some((
FuncCont {
name: "pcall_finish",
func: Rc::new(pcall_finish),
},
false,
)),
None,
false,
None,
)
}
}
}
fn pcall_finish(vm: &mut VM, result: Result<Value>) -> Result<Value> {
vm.rust_call_depth -= 1;
match result {
Ok(result) => {
let mut result = result.into_vec();
result.insert(0, Value::Bool(true));
Ok(Value::Mult(result))
}
Err(error) => Ok(Value::Mult(vec![Value::Bool(false), error.into_value()])),
}
}
fn print(vm: &mut VM) -> Result<Value> {
let mut iter = vm.arg_split(0).into_iter();
let mut buf = Vec::new();
if let Some(first) = iter.next() {
first.write(&mut buf)?;
}
for val in iter {
buf.write_all(&[b'\t'])?;
val.write(&mut buf)?;
}
let mut out = vm.stdout.borrow_mut();
out.write_all(&buf)?;
writeln!(out)?;
out.flush()?;
Ok(Value::empty())
}
fn rawequal(vm: &mut VM) -> Result<Value> {
Ok(Value::Bool(vm.arg(0)? == vm.arg(1)?))
}
fn rawget(vm: &mut VM) -> Result<Value> {
let table = vm.arg_table(0)?;
let table = table.borrow();
Ok(table.get(&vm.arg(1)?))
}
fn rawlen(vm: &mut VM) -> Result<Value> {
let len = match vm.arg(0)? {
Value::String(s) => s.len(),
Value::Table(tbl) => tbl.borrow().len(),
v => {
return Err(vm.arg_error(
0,
Error::from_lua(LuaError::ExpectedType(ValueType::String, v.value_type())),
))
}
};
Ok(Value::int(len as i64))
}
fn rawset(vm: &mut VM) -> Result<Value> {
let tbl_ref = vm.arg(0)?.to_table()?;
tbl_ref.borrow_mut().set(vm.arg(1)?, vm.arg(2)?)?;
Ok(Value::Table(tbl_ref))
}
fn select(vm: &mut VM) -> Result<Value> {
let index = match vm.arg(0)? {
Value::Number(n) => n.coerce_int()?,
Value::String(s) => {
if *s == [b'#'] {
return Ok(Value::int(vm.arg_split(1).len() as i64));
}
Numeric::from_str(&s)?.coerce_int()?
}
v => {
return Err(vm.arg_error(
0,
Error::from_lua(LuaError::ExpectedType(ValueType::Number, v.value_type())),
))
}
};
let mut args = vm.arg_split(1);
Ok(Value::Mult(match index.cmp(&0) {
Ordering::Greater => {
let index = index as usize;
if index > args.len() {
vec![]
} else {
args.split_off(index - 1)
}
}
Ordering::Less => {
let index = index.unsigned_abs() as usize;
if index > args.len() {
return err!(LuaError::InvalidIndex);
} else {
args.split_off(args.len() - index)
}
}
Ordering::Equal => return err!(LuaError::InvalidIndex),
}))
}
fn setmetatable(vm: &mut VM) -> Result<Value> {
let tbl = vm.arg_table(0)?;
if let Some(meta) = tbl.borrow().get_meta().clone() {
if Value::Nil != meta.borrow().get(&Value::str("__metatable")) {
return err!(LuaError::MetatableProtected);
}
}
{
match vm.arg(1)? {
Value::Nil => tbl.borrow_mut().set_meta(None),
Value::Table(meta) => {
if meta.borrow().get(&Value::str("__gc")).is_truthy() {
vm.mark_table(tbl.clone());
}
tbl.borrow_mut().set_meta(Some(meta))
}
v => return err!(LuaError::ExpectedType(ValueType::Table, v.value_type())),
}
}
Ok(Value::Table(tbl))
}
fn tonumber(vm: &mut VM) -> Result<Value> {
let base = match vm.arg_opt(1) {
Some(v) => Some(v.to_number_coerce()?.coerce_int()? as u32),
None => None,
};
let num = match vm.arg(0)? {
Value::Number(n) => {
if base.is_some() {
return err!(LuaError::ExpectedType(ValueType::String, ValueType::Number));
} else {
Value::Number(n)
}
}
Value::String(s) => match Numeric::from_str_radix(&s, base) {
Ok(n) => Value::Number(n),
Err(..) => Value::Nil,
},
_ => Value::Nil,
};
Ok(num)
}
pub(in crate::vm) fn tostring(vm: &mut VM) -> Result<Value> {
let v = vm.arg(0)?;
if let Some(meta) = vm.get_metatable(&v) {
match meta.borrow().get(&Value::str("__tostring")) {
Value::Nil => match meta.borrow().get(&Value::str("__name")) {
Value::String(str) => {
let mut str = str.to_vec();
str.write_all(": ".as_bytes())?;
Ok(Value::str_bytes(str))
}
_ => v.write_as_value(),
},
Value::Func(func) => Ok(Value::str_bytes(
vm.call_recursive(func, vec![v])?
.to_string_coerce()
.map_err(|err| err.map_lua_error(|_| LuaError::ToStringReturnType))?,
)),
val => err!(LuaError::ExpectedType(ValueType::Func, val.value_type())),
}
} else {
v.write_as_value()
}
}
fn r#type(vm: &mut VM) -> Result<Value> {
let v = vm.arg(0)?;
Ok(Value::string(format!("{}", v.value_type_basic())))
}
fn warn(vm: &mut VM) -> Result<Value> {
if vm.arg_len() == 0 {
return Err(vm.arg_error(
0,
Error::from_lua(LuaError::ExpectedType(
ValueType::String,
ValueType::Custom("no value".to_string()),
)),
));
}
if vm.arg_len() == 1
&& matches!(
vm.frames
.last()
.and_then(|f| f.varargs.get(0))
.map(Value::to_string),
Some(Ok(&[b'@', ..]))
)
{
let cmd = vm.arg_string(0)?;
match &cmd[..] {
[b'@', b'o', b'n'] => {
vm.warn = true;
}
[b'@', b'o', b'f', b'f'] => {
vm.warn = false;
}
_ => {}
}
} else {
let args = vm.arg_split(0);
let mut buffer = "Lua warning: ".as_bytes().to_vec();
for arg in args {
let str = arg.to_string_coerce()?;
buffer.extend(str);
}
if vm.warn {
stderr().write_all(&buffer)?;
writeln!(stderr())?;
}
}
Ok(Value::empty())
}
fn xpcall(vm: &mut VM) -> Result<CallResult> {
let args = vm.arg_split(0);
match &args[..] {
[] | [_] => err!(LuaError::CallInvalid(ValueType::Nil, LuaErrorSrc::None)),
[func, handler, args @ ..] => {
let func = match func.to_func() {
Ok(func) => func,
Err(_) => {
let ret = pcall_finish(
vm,
err!(LuaError::CallInvalid(func.value_type(), LuaErrorSrc::None,)),
)?;
if let Some(ret_reg) = vm.frames.pop().unwrap().ret_reg {
vm.frames.last_mut().unwrap().regs[ret_reg] = ret;
}
return Ok(CallResult::Continue);
}
};
vm.frames.last_mut().unwrap().protected = Some(
vm.thread
.borrow_mut()
.protected
.replace(ProtectedHandler::XPCall(handler.to_func()?)),
);
vm.rust_call_depth += 1;
vm.call_prepare(
func,
args.to_vec(),
Some((
FuncCont {
name: "pcall_finish",
func: Rc::new(pcall_finish),
},
false,
)),
None,
false,
None,
)
}
}
}