luau_vm/libs/table/

mod.rs

use core::ptr;

use luau_common::ByteSlice;

use crate::Table;
use crate::VmResult;
use crate::call::ThreadStack;
use crate::debug::DebugRuntime;
use crate::gc::GcBarrier;
use crate::gc::GcObject;
use crate::handle::RawHandle;
use crate::native::{NativeCallContext, NativeCallResult, NativeFunction};
use crate::state::ThreadState;
use crate::string::StringRuntime;
use crate::table::TableRuntime;
use crate::thread::stack::RawStackAccess;
use crate::thread::{LuaStringBuilder, LuaStringBuilderStorage, Thread};
use crate::types::{LUA_TFUNCTION, LUA_TNUMBER, LUA_TSTRING, LUA_TTABLE};
use crate::vm::VmOperations;

mod sort;

use sort::table_sort;

static TABLE_FUNCS: [NativeFunction; 17] = [
    NativeFunction {
        name: "concat",
        function: table_concat,
    },
    NativeFunction {
        name: "foreach",
        function: table_for_each,
    },
    NativeFunction {
        name: "foreachi",
        function: table_for_each_i,
    },
    NativeFunction {
        name: "getn",
        function: table_getn,
    },
    NativeFunction {
        name: "maxn",
        function: table_maxn,
    },
    NativeFunction {
        name: "insert",
        function: table_insert,
    },
    NativeFunction {
        name: "remove",
        function: table_remove,
    },
    NativeFunction {
        name: "sort",
        function: table_sort,
    },
    NativeFunction {
        name: "pack",
        function: table_pack,
    },
    NativeFunction {
        name: "unpack",
        function: table_unpack,
    },
    NativeFunction {
        name: "move",
        function: table_move,
    },
    NativeFunction {
        name: "create",
        function: table_create,
    },
    NativeFunction {
        name: "find",
        function: table_find,
    },
    NativeFunction {
        name: "clear",
        function: table_clear,
    },
    NativeFunction {
        name: "freeze",
        function: table_freeze,
    },
    NativeFunction {
        name: "isfrozen",
        function: table_is_frozen,
    },
    NativeFunction {
        name: "clone",
        function: table_clone,
    },
];

fn table_argument(thread: &Thread, argument: i32) -> VmResult<Table> {
    unsafe { thread.check_type(argument, LUA_TTABLE)? };
    Ok(unsafe { thread.to_object(argument).unwrap_unchecked().table_value() })
}

/// `moveelements`
unsafe fn move_elements(
    thread: &Thread,
    src_index: i32,
    dst_index: i32,
    first: i32,
    last: i32,
    target: i32,
) -> VmResult {
    let src = table_argument(thread, src_index)?;
    let dst = table_argument(thread, dst_index)?;

    unsafe {
        if dst.as_ptr().as_ref().unwrap_unchecked().readonly != 0 {
            return thread.readonly_error().map_err(Into::into);
        }

        let count = last - first + 1;
        let src_size = src.as_ptr().as_ref().unwrap_unchecked().size_array;
        let dst_size = dst.as_ptr().as_ref().unwrap_unchecked().size_array;

        if (first as u32).wrapping_sub(1) < src_size as u32
            && (target as u32).wrapping_sub(1) < dst_size as u32
            && (first as u32).wrapping_sub(1).wrapping_add(count as u32) <= src_size as u32
            && (target as u32).wrapping_sub(1).wrapping_add(count as u32) <= dst_size as u32
        {
            if count > 0 {
                let src_array = src.array_cursor().add((first - 1) as usize).as_ptr();
                let dst_array = dst.array_cursor().add((target - 1) as usize).as_ptr();
                ptr::copy(src_array, dst_array, count as usize);
            }

            let dst_object: GcObject = dst.into();
            if dst_object.is_black() {
                thread.barrier_back(
                    dst_object,
                    &raw mut dst.as_ptr().as_mut().unwrap_unchecked().gc_list,
                );
            }
        } else if target > last || target <= first || dst != src {
            for i in 0..count {
                thread.raw_geti(src_index, first + i)?;
                thread.raw_seti(dst_index, target + i)?;
            }
        } else {
            for i in (0..count).rev() {
                thread.raw_geti(src_index, first + i)?;
                thread.raw_seti(dst_index, target + i)?;
            }
        }
    }
    Ok(())
}

/// `addfield`
unsafe fn add_field(
    thread: &Thread,
    buffer: &mut LuaStringBuilder<'_, '_>,
    index: i32,
    table: Table,
) -> VmResult {
    unsafe {
        if (index as u32).wrapping_sub(1)
            < table.as_ptr().as_ref().unwrap_unchecked().size_array as u32
        {
            let entry = table.array_slot((index - 1) as usize);
            if entry.is_string() {
                buffer.push_bytes(entry.string_value().as_bytes())?;
                return Ok(());
            }
        }

        let value_type = thread.raw_geti(1, index)?;
        if value_type != LUA_TSTRING && value_type != LUA_TNUMBER {
            let type_name = thread.lua_type_name(-1);
            let message = luau_printf::sprintf!(
                "invalid value (%s) at index %d in table for 'concat'",
                type_name.as_bstr(),
                index
            );
            return crate::error!(thread, &message).map_err(Into::into);
        }

        buffer.push_stack_value()?;
    }
    Ok(())
}

/// `maxn`
fn table_maxn(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let max = unsafe {
        thread.check_type(1, LUA_TTABLE)?;

        let table = thread.to_object(1).unwrap_unchecked().table_value();
        let mut max = 0.0f64;

        for i in 0..table.as_ptr().as_ref().unwrap_unchecked().size_array {
            if !table.array_slot(i as usize).is_nil() {
                max = (i + 1) as f64;
            }
        }

        for i in 0..table.node_count() {
            let node = table.node(i as i32);
            if !node.value_unchecked().is_nil() && node.key().tt() == LUA_TNUMBER {
                let value = node.key().number_value();
                if value > max {
                    max = value;
                }
            }
        }

        max
    };
    ctx.push_number(max)?;
    Ok(1)
}

/// `getn`
fn table_getn(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.check_type(1, LUA_TTABLE)? };
    ctx.push_integer(unsafe { thread.obj_len(1) })?;
    Ok(1)
}

/// `tinsert`
fn table_insert(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        let n = thread.obj_len(1);

        let pos = match thread.get_top() {
            2 => n + 1,
            3 => {
                let pos = thread.check_integer(2)?;
                if (1..=n).contains(&pos) {
                    move_elements(thread, 1, 1, pos, n, pos + 1)?;
                }
                pos
            }
            _ => {
                return crate::error!(thread, "wrong number of arguments to 'insert'")
                    .map_err(Into::into);
            }
        };

        thread.raw_seti(1, pos)?;
        Ok(0)
    }
}

/// `tremove`
fn table_remove(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        let n = thread.obj_len(1);
        let pos = thread.opt_integer(2, n)?;

        if !(1..=n).contains(&pos) {
            return Ok(0);
        }

        thread.raw_geti(1, pos)?;
        move_elements(thread, 1, 1, pos + 1, n, pos)?;
        thread.push_nil()?;
        thread.raw_seti(1, n)?;
        Ok(1)
    }
}

/// `tmove`
fn table_move(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let dst_index = unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        let first = thread.check_integer(2)?;
        let last = thread.check_integer(3)?;
        let target = thread.check_integer(4)?;
        let dst_index = if thread.is_none_or_nil(5) != 0 { 1 } else { 5 };
        thread.check_type(dst_index, LUA_TTABLE)?;

        if last >= first {
            if first <= 0 && last >= i32::MAX + first {
                return thread
                    .lua_arg_error(3, "too many elements to move")
                    .map_err(Into::into);
            }

            let count = last - first + 1;
            if target > i32::MAX - count + 1 {
                return thread
                    .lua_arg_error(4, "destination wrap around")
                    .map_err(Into::into);
            }

            let dst = thread.to_object(dst_index).unwrap_unchecked().table_value();
            if dst.as_ptr().as_ref().unwrap_unchecked().readonly != 0 {
                return thread.readonly_error().map_err(Into::into);
            }

            if target > 0
                && (target - 1) <= dst.as_ptr().as_ref().unwrap_unchecked().size_array
                && (target - 1 + count) > dst.as_ptr().as_ref().unwrap_unchecked().size_array
            {
                thread.resize_array(dst, target - 1 + count)?;
            }

            move_elements(thread, 1, dst_index, first, last, target)?;
        }

        dst_index
    };

    unsafe { thread.push_value(dst_index)? };
    Ok(1)
}

/// `tconcat`
fn table_concat(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let separator = thread.opt_string(2)?.unwrap_or(b"".as_bstr());
        thread.check_type(1, LUA_TTABLE)?;
        let mut index = thread.opt_integer(3, 1)?;
        let last = if thread.is_none_or_nil(4) != 0 {
            thread.obj_len(1)
        } else {
            thread.check_integer(4)?
        };
        let table = thread.to_object(1).unwrap_unchecked().table_value();
        let mut buffer_storage = LuaStringBuilderStorage::uninit();
        let mut buffer = LuaStringBuilder::new(thread, &mut buffer_storage);

        while index < last {
            add_field(thread, &mut buffer, index, table)?;
            if !separator.is_empty() {
                buffer.push_bytes(separator)?;
            }
            index += 1;
        }

        if index == last {
            add_field(thread, &mut buffer, index, table)?;
        }

        buffer.finish()?;
        Ok(1)
    }
}

/// `foreachi`
fn table_for_each_i(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        thread.check_type(2, LUA_TFUNCTION)?;
        let n = thread.obj_len(1);
        for index in 1..=n {
            thread.push_value(2)?;
            thread.push_integer(index)?;
            thread.raw_geti(1, index)?;
            thread.call(2, 1)?;

            if thread.is_nil(-1) == 0 {
                return Ok(1);
            }

            thread.pop(1);
        }

        Ok(0)
    }
}

/// `foreach`
fn table_for_each(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        thread.check_type(2, LUA_TFUNCTION)?;
        thread.push_nil()?;
        while thread.next(1)? != 0 {
            thread.push_value(2)?;
            thread.push_value(-3)?;
            thread.push_value(-3)?;
            thread.call(2, 1)?;

            if thread.is_nil(-1) == 0 {
                return Ok(1);
            }

            thread.pop(2);
        }

        Ok(0)
    }
}

/// `tpack`
fn table_pack(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let n = thread.get_top();
        thread.create_table(n as usize, 1)?;

        let table = thread.to_object(-1).unwrap_unchecked().table_value();
        let array = table.array_cursor();
        let base = thread.stack_base();
        for i in 0..n as usize {
            array
                .add(i)
                .value_unchecked()
                .set_obj(base.add(i).value_unchecked());
        }

        let key = thread.intern_string(b"n".as_bstr())?;
        let node_cursor = thread.set_str(table, key)?;
        node_cursor
            .node_unchecked()
            .value_unchecked()
            .set_number(n as f64);
        Ok(1)
    }
}

/// `tunpack`
fn table_unpack(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        let table = thread.to_object(1).unwrap_unchecked().table_value();
        let start = thread.opt_integer(2, 1)?;
        let end = if thread.is_none_or_nil(3) != 0 {
            thread.obj_len(1)
        } else {
            thread.check_integer(3)?
        };

        if start > end {
            return Ok(0);
        }

        let n = (end as u32).wrapping_sub(start as u32) as i32 + 1;
        if n <= 0 || thread.check_stack(n) == 0 {
            return crate::error!(thread, "too many results to unpack").map_err(Into::into);
        }

        if start == 1 && n <= table.as_ptr().as_ref().unwrap_unchecked().size_array {
            let top = thread.stack_top();
            for i in 0..n as usize {
                top.add(i).value_unchecked().set_obj(table.array_slot(i));
            }
            thread.expand_stack_limit(top.add(n as usize));
            thread.set_stack_top(top.add(n as usize));
        } else {
            for i in start..end {
                thread.raw_geti(1, i)?;
            }
            thread.raw_geti(1, end)?;
        }

        Ok(n as usize)
    }
}

/// `tcreate`
fn table_create(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let size = thread.check_integer(1)?;
        if size < 0 {
            return thread
                .lua_arg_error(1, "size out of range")
                .map_err(Into::into);
        }

        if thread.is_none_or_nil(2) == 0 {
            let value = thread.stack_base().add(1).value_unchecked();
            thread.create_table(size as usize, 0)?;
            let table = thread.to_object(-1).unwrap_unchecked().table_value();
            let array = table.array_cursor();
            for i in 0..size as usize {
                array.add(i).value_unchecked().set_obj(value);
            }
        } else {
            thread.create_table(size as usize, 0)?;
        }

        Ok(1)
    }
}

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

        let init = thread.opt_integer(3, 1)?;
        if init < 1 {
            return thread
                .lua_arg_error(3, "index out of range")
                .map_err(Into::into);
        }

        let table = thread.to_object(1).unwrap_unchecked().table_value();
        let needle = thread.stack_base().add(1).value_unchecked();

        for index in init.. {
            let entry = table.get_num(index);
            if entry.is_nil() {
                break;
            }

            let equal = if entry.tt() == needle.tt() {
                thread.equal_value(entry, needle)? != 0
            } else {
                false
            };

            if equal {
                thread.push_integer(index)?;
                return Ok(1);
            }
        }

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

/// `tclear`
fn table_clear(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.check_type(1, LUA_TTABLE)? };
    unsafe { thread.clear_table(1)? };
    Ok(0)
}

/// `tfreeze`
fn table_freeze(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        if thread.get_readonly(1) != 0 {
            return thread
                .lua_arg_error(1, "table is already frozen")
                .map_err(Into::into);
        }
        if thread.get_metafield(1, "__metatable")? != 0 {
            return thread
                .lua_arg_error(1, "table has a protected metatable")
                .map_err(Into::into);
        }

        thread.set_readonly(1, 1);
        thread.push_value(1)?;
        Ok(1)
    }
}

/// `tisfrozen`
fn table_is_frozen(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.check_type(1, LUA_TTABLE)? };
    unsafe { thread.push_boolean(thread.get_readonly(1))? };
    Ok(1)
}

/// `tclone`
fn table_clone(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;
        if thread.get_metafield(1, "__metatable")? != 0 {
            return thread
                .lua_arg_error(1, "table has a protected metatable")
                .map_err(Into::into);
        }

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

impl Thread {
    /// `luaopen_table`
    pub unsafe fn open_table(&self) -> NativeCallResult {
        unsafe { self.register(Some(super::LUA_TABLIB_NAME), &TABLE_FUNCS[..])? };

        // Lua 5.1 compatibility global.
        unsafe {
            self.push_native_function(table_unpack, Some("unpack"))?;
            self.set_global("unpack")?;
        }

        Ok(1)
    }
}