luau_vm/libs/

base.rs

use luau_common::ByteSlice;
use std::io::{self, Write};

use crate::VmResult;
use crate::debug::LuaDebug;
use crate::native::{NativeCallContext, NativeCallResult, NativeFunction, RawNativeFunction};
use crate::state::LUA_OK;
use crate::thread::{LUA_GLOBALS_INDEX, LUA_MULTRET, LUA_TNONE, Thread};
use crate::types::{LUA_TBOOLEAN, LUA_TFUNCTION, LUA_TNIL, LUA_TNUMBER, LUA_TSTRING, LUA_TTABLE};
use crate::userdata::USERDATA_TAG_PROXY;
use crate::userdata::UserdataRuntime;

static BASE_FUNCS: [NativeFunction; 19] = [
    NativeFunction {
        name: "assert",
        function: base_assert,
    },
    NativeFunction {
        name: "error",
        function: base_error,
    },
    NativeFunction {
        name: "gcinfo",
        function: base_gcinfo,
    },
    NativeFunction {
        name: "getfenv",
        function: base_get_fenv,
    },
    NativeFunction {
        name: "getmetatable",
        function: base_get_metatable,
    },
    NativeFunction {
        name: "next",
        function: base_next,
    },
    NativeFunction {
        name: "newproxy",
        function: base_newproxy,
    },
    NativeFunction {
        name: "print",
        function: base_print,
    },
    NativeFunction {
        name: "rawequal",
        function: base_raw_equal,
    },
    NativeFunction {
        name: "rawget",
        function: base_raw_get,
    },
    NativeFunction {
        name: "rawset",
        function: base_raw_set,
    },
    NativeFunction {
        name: "rawlen",
        function: base_raw_len,
    },
    NativeFunction {
        name: "select",
        function: base_select,
    },
    NativeFunction {
        name: "setfenv",
        function: base_set_fenv,
    },
    NativeFunction {
        name: "setmetatable",
        function: base_set_metatable,
    },
    NativeFunction {
        name: "tonumber",
        function: base_tonumber,
    },
    NativeFunction {
        name: "tostring",
        function: base_tostring,
    },
    NativeFunction {
        name: "type",
        function: base_type,
    },
    NativeFunction {
        name: "typeof",
        function: base_typeof,
    },
];

fn parse_unsigned_radix(bytes: &[u8], base: u32) -> Option<u64> {
    let mut bytes = bytes;
    while let Some(first) = bytes.first() {
        if !first.is_ascii_whitespace() {
            break;
        }
        bytes = &bytes[1..];
    }

    let digits_end = bytes
        .iter()
        .position(|byte| byte.is_ascii_whitespace())
        .unwrap_or(bytes.len());
    let digits = &bytes[..digits_end];
    if digits.is_empty() {
        return None;
    }

    let mut trailing = &bytes[digits_end..];
    while let Some(first) = trailing.first() {
        if !first.is_ascii_whitespace() {
            return None;
        }
        trailing = &trailing[1..];
    }

    let mut value = 0u64;
    for &byte in digits {
        let digit = match byte {
            b'0'..=b'9' => u32::from(byte - b'0'),
            b'a'..=b'z' => u32::from(byte - b'a') + 10,
            b'A'..=b'Z' => u32::from(byte - b'A') + 10,
            _ => return None,
        };

        if digit >= base {
            return None;
        }

        value = value.checked_mul(u64::from(base))?;
        value = value.checked_add(u64::from(digit))?;
    }

    Some(value)
}

/// `getfunc`
unsafe fn get_func(thread: &Thread, allow_default_level: bool) -> VmResult {
    unsafe {
        if thread.type_of(1) == LUA_TFUNCTION {
            thread.push_value(1)?;
            return Ok(());
        }

        let level = if allow_default_level {
            thread.opt_integer(1, 1)?
        } else {
            thread.check_integer(1)?
        };

        if level < 0 {
            return thread
                .lua_arg_error(1, "level must be non-negative")
                .map_err(Into::into);
        }

        let mut ar = LuaDebug::default();
        if thread.get_info(level, "f", &mut ar)? == 0 {
            return thread.lua_arg_error(1, "invalid level").map_err(Into::into);
        }

        if thread.type_of(-1) == LUA_TNIL {
            return crate::error!(
                thread,
                "no function environment for tail call at level %d",
                level
            )
            .map_err(Into::into);
        }
    }
    Ok(())
}

/// `luaB_assert`
fn base_assert(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_any(1)?;
        if thread.to_boolean(1) == 0 {
            let message = thread
                .opt_string(2)?
                .unwrap_or(b"assertion failed!".as_bstr());
            return crate::error!(thread, message).map_err(Into::into);
        }

        Ok(thread.get_top() as usize)
    }
}

/// `luaB_print`
fn base_print(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let count = thread.get_top();
        let mut stdout = io::stdout().lock();

        for index in 1..=count {
            let string = thread.lua_to_string(index)?;
            if index > 1 {
                let _ = stdout.write_all(b"\t");
            }
            let _ = stdout.write_all(string.as_bytes());
            thread.pop(1);
        }

        let _ = stdout.write_all(b"\n");
        Ok(0)
    }
}

/// `luaB_error`
fn base_error(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let level = thread.opt_integer(2, 1)?;
        thread.set_top(1)?;

        if thread.is_string(1) != 0 && level > 0 {
            thread.push_where(level)?;
            thread.push_value(1)?;
            thread.concat(2)?;
        }

        thread.error().map_err(Into::into)
    }
}

/// `luaB_gcinfo`
fn base_gcinfo(ctx: NativeCallContext) -> NativeCallResult {
    let count = unsafe { ctx.raw_thread().gc(crate::LUA_GC_COUNT, 0)? };
    ctx.push_integer(count)?;
    Ok(1)
}

/// `luaB_getfenv`
fn base_get_fenv(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        get_func(thread, true)?;

        if thread.is_native_function(-1) != 0 {
            thread.push_value(LUA_GLOBALS_INDEX)?;
        } else {
            thread.get_fenv(-1)?;
        }

        thread.set_safe_env(-1, 0);
    }
    Ok(1)
}

/// `luaB_getmetatable`
fn base_get_metatable(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_any(1)?;

        if thread.get_metatable(1)? == 0 {
            thread.push_nil()?;
            return Ok(1);
        }

        let _ = thread.get_metafield(1, "__metatable")?;
        Ok(1)
    }
}

/// `luaB_next`
fn base_next(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        thread.set_top(2)?;

        if thread.next(1)? != 0 {
            Ok(2)
        } else {
            thread.push_nil()?;
            Ok(1)
        }
    }
}

/// `luaB_inext`
fn base_inext(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let index = thread.check_integer(2)? + 1;
        thread.check_type(1, LUA_TTABLE)?;
        thread.push_integer(index)?;
        thread.raw_geti(1, index)?;
        Ok(if thread.type_of(-1) == LUA_TNIL { 0 } else { 2 })
    }
}

/// `luaB_ipairs`
fn base_ipairs(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.check_type(1, LUA_TTABLE)? };
    unsafe {
        thread.push_value(crate::thread::upvalue_index(1))?;
        thread.push_value(1)?;
        thread.push_integer(0)?;
    }
    Ok(3)
}

/// `luaB_pairs`
fn base_pairs(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.check_type(1, LUA_TTABLE)? };
    unsafe {
        thread.push_value(crate::thread::upvalue_index(1))?;
        thread.push_value(1)?;
        thread.push_nil()?;
    }
    Ok(3)
}

/// `luaB_newproxy`
fn base_newproxy(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let tag = thread.type_of(1);
        ctx.arg(1).expected(
            tag == LUA_TNONE || tag == LUA_TNIL || tag == LUA_TBOOLEAN,
            "nil or boolean",
        )?;

        let needs_metatable = thread.to_boolean(1) != 0;
        let _ = thread.new_userdata_tagged_internal(0, USERDATA_TAG_PROXY as i32)?;

        if needs_metatable {
            thread.create_table(0, 0)?;
            thread.set_metatable(-2)?;
        }
    }

    Ok(1)
}

/// `luaB_rawequal`
fn base_raw_equal(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let result = unsafe {
        thread.check_any(1)?;
        thread.check_any(2)?;
        thread.raw_equal(1, 2)
    };
    ctx.push_boolean(result != 0)?;
    Ok(1)
}

/// `luaB_rawget`
fn base_raw_get(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        thread.check_any(2)?;
        thread.set_top(2)?;
        thread.raw_get(1);
    }
    Ok(1)
}

/// `luaB_rawset`
fn base_raw_set(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        thread.check_any(2)?;
        thread.check_any(3)?;
        thread.set_top(3)?;
        thread.raw_set(1)?;
    }
    Ok(1)
}

/// `luaB_rawlen`
fn base_raw_len(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let tag = unsafe { thread.type_of(1) };
    ctx.arg(1).expected(
        tag == LUA_TTABLE || tag == LUA_TSTRING,
        "table or string expected",
    )?;
    ctx.push_integer(unsafe { thread.obj_len(1) })?;
    Ok(1)
}

/// `luaB_select`
fn base_select(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let count = thread.get_top();

        if thread.type_of(1) == LUA_TSTRING && matches!(thread.check_string(1)?.first(), Some(b'#'))
        {
            thread.push_integer(count - 1)?;
            return Ok(1);
        }

        let mut index = thread.check_integer(1)?;
        if index < 0 {
            index += count;
        } else if index > count {
            index = count;
        }

        if index < 1 {
            return thread
                .lua_arg_error(1, "index out of range")
                .map_err(Into::into);
        }

        Ok((count - index) as usize)
    }
}

/// `luaB_setfenv`
fn base_set_fenv(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(2, LUA_TTABLE)?;
        get_func(thread, false)?;
        thread.push_value(2)?;
        thread.set_safe_env(-1, 0);

        if thread.type_of(1) == LUA_TNUMBER && thread.to_number(1) == Some(0.0) {
            thread.push_thread()?;
            thread.insert(-2);
            thread.set_fenv(-2);
            return Ok(0);
        }

        if thread.is_native_function(-2) != 0 || thread.set_fenv(-2) == 0 {
            return crate::error!(
                thread,
                "%s",
                "'setfenv' cannot change environment of given object"
            )
            .map_err(Into::into);
        }
    }

    Ok(1)
}

/// `luaB_setmetatable`
fn base_set_metatable(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let second_type = thread.type_of(2);
        thread.check_type(1, LUA_TTABLE)?;
        ctx.arg(2).expected(
            second_type == LUA_TNIL || second_type == LUA_TTABLE,
            "nil or table",
        )?;

        if thread.get_metafield(1, "__metatable")? != 0 {
            return crate::error!(thread, "cannot change a protected metatable")
                .map_err(Into::into);
        }

        thread.set_top(2)?;
        thread.set_metatable(1)?;
    }
    Ok(1)
}

/// `luaB_tonumber`
fn base_tonumber(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let base = thread.opt_integer(2, 10)?;

        if base == 10 {
            if let Some(value) = thread.to_number(1) {
                thread.push_number(value)?;
                return Ok(1);
            }

            thread.check_any(1)?;
        } else {
            let string = thread.check_string(1)?;
            if !(2..=36).contains(&base) {
                return thread
                    .lua_arg_error(2, "base out of range")
                    .map_err(Into::into);
            }

            if let Some(value) = parse_unsigned_radix(string, base as u32) {
                thread.push_number(value as f64)?;
                return Ok(1);
            }
        }

        thread.push_nil()?;
    }
    Ok(1)
}

/// `luaB_type`
fn base_type(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let name = unsafe {
        thread.check_any(1)?;
        thread.type_name(thread.type_of(1))
    };
    ctx.push_string(name)?;
    Ok(1)
}

/// `luaB_typeof`
fn base_typeof(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let name = unsafe {
        thread.check_any(1)?;
        thread.lua_type_name(1)
    };
    ctx.push_string(name)?;
    Ok(1)
}

/// `luaB_pcally`
fn base_pcally(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_any(1)?;

        thread.protected_call_yieldable(thread.get_top() - 1, LUA_MULTRET, 0)
    }
}

/// `luaB_pcallcont`
fn base_pcall_cont(ctx: NativeCallContext, status: i32) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.raw_check_stack(1)?;

        if status == LUA_OK {
            thread.push_boolean(1)?;
            thread.insert(1);
            Ok(thread.get_top() as usize)
        } else {
            thread.push_boolean(0)?;
            thread.insert(-2);
            Ok(2)
        }
    }
}

/// `luaB_xpcally`
fn base_xpcally(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(2, LUA_TFUNCTION)?;
        thread.push_value(1)?;
        thread.push_value(2)?;
        thread.replace(1);
        thread.replace(2);

        thread.protected_call_yieldable(thread.get_top() - 2, LUA_MULTRET, 1)
    }
}

/// `luaB_xpcallcont`
fn base_xpcall_cont(ctx: NativeCallContext, status: i32) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        if status == LUA_OK {
            thread.raw_check_stack(1)?;
            thread.push_boolean(1)?;
            thread.replace(1);
            return Ok(thread.get_top() as usize);
        }

        thread.raw_check_stack(1)?;
        thread.push_boolean(0)?;
        thread.insert(-2);
        Ok(2)
    }
}

/// `luaB_tostring`
fn base_tostring(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.check_any(1)? };
    let _ = unsafe { thread.lua_to_string(1)? };
    Ok(1)
}

/// `auxopen`
unsafe fn aux_open(
    thread: &Thread,
    name: &'static str,
    function: RawNativeFunction,
    upvalue: RawNativeFunction,
) -> NativeCallResult {
    unsafe {
        thread.push_native_closure_k(upvalue, None, 0, None)?;
        thread.push_native_closure_k(function, Some(name), 1, None)?;
        thread.raw_set_field(-2, name)?;
    }
    Ok(0)
}

impl Thread {
    /// `luaopen_base`
    pub unsafe fn open_base(&self) -> NativeCallResult {
        unsafe {
            self.push_value(LUA_GLOBALS_INDEX)?;
            self.set_global("_G")?;

            self.register(Some("_G"), &BASE_FUNCS[..])?;
            self.push_string("Luau")?;
            self.set_global("_VERSION")?;

            aux_open(self, "ipairs", base_ipairs, base_inext)?;
            aux_open(self, "pairs", base_pairs, base_next)?;

            self.push_native_closure_k(base_pcally, Some("pcall"), 0, Some(base_pcall_cont))?;
            self.raw_set_field(-2, "pcall")?;

            self.push_native_closure_k(base_xpcally, Some("xpcall"), 0, Some(base_xpcall_cont))?;
            self.raw_set_field(-2, "xpcall")?;
        }

        Ok(1)
    }
}