luau 0.732.0

use core::marker::PhantomData;

use luau_vm::internal::RawHandle;
use luau_vm::thread::StackGuard;

use super::Table;
use crate::error::Error;
use crate::value::{FromLua, IntoLua, Value};

/// An iterator over consecutive integer keys in a [`Table`], starting at 1.
pub struct TableSequence<'table, 'lua, V> {
    table: &'table Table<'lua>,
    index: usize,
    _marker: PhantomData<V>,
}

impl<'lua> Table<'lua> {
    /// Returns the result of Luau's length operator.
    ///
    /// This may invoke the `__len` metamethod. Use [`Table::raw_len`] to read
    /// the table's raw sequence boundary.
    pub fn len(&self) -> Result<i64, Error> {
        unsafe {
            let thread = self.reference.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            self.push_to(&thread)?;
            let length = vm_thread
                .len(-1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            if !length.is_finite() || length.fract() != 0.0 {
                return Err(Error::runtime("object length is not an integer"));
            }
            num_traits::cast(length)
                .ok_or_else(|| Error::runtime("object length is not an integer"))
        }
    }

    /// Returns the raw sequence boundary without invoking metamethods.
    pub fn raw_len(&self) -> usize {
        self.raw_len_cached()
    }

    /// Returns whether both the array and record parts are empty.
    pub fn is_empty(&self) -> bool {
        self.is_empty_cached()
    }

    /// Appends a value using Luau length and assignment semantics.
    pub fn push(&self, value: impl IntoLua<'lua>) -> Result<(), Error> {
        self.set(
            self.len()?
                .checked_add(1)
                .ok_or_else(Error::index_out_of_bounds)?,
            value,
        )
    }

    /// Removes and returns the value at the Luau sequence boundary.
    pub fn pop<V>(&self) -> Result<V, Error>
    where
        V: FromLua<'lua>,
    {
        let index = self.len()?;
        let value = self.get(index)?;
        self.set(index, Value::Nil)?;
        Ok(value)
    }

    /// Removes a key, shifting later sequence values down for array indices.
    pub fn remove(&self, key: impl IntoLua<'lua>) -> Result<(), Error> {
        let thread = self.thread();
        let key = key.into_lua(thread.lua_ref())?;
        let Some(index) = key.as_table_array_index()? else {
            return self.set(key, Value::Nil);
        };

        let size = self.len()?;
        if index > size {
            return Err(Error::index_out_of_bounds());
        }

        for current in index..size {
            let value: Value<'lua> = self.get(current + 1)?;
            self.set(current, value)?;
        }
        self.set(size, Value::Nil)
    }

    /// Sets an integer key without invoking metamethods.
    pub fn raw_seti(&self, index: usize, value: impl IntoLua<'lua>) -> Result<(), Error> {
        let index = raw_array_index(index)?;
        self.invalidate_safe_env();
        unsafe {
            let thread = self.reference.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            value.push_into_stack(&thread)?;
            vm_thread
                .raw_seti(-2, index)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Appends a value without invoking metamethods.
    pub fn raw_push(&self, value: impl IntoLua<'lua>) -> Result<(), Error> {
        self.raw_seti(self.raw_len_cached() + 1, value)
    }

    /// Removes and returns the value at the raw sequence boundary.
    pub fn raw_pop<V>(&self) -> Result<V, Error>
    where
        V: FromLua<'lua>,
    {
        let index = raw_array_index(self.raw_len_cached())?;
        self.invalidate_safe_env();
        unsafe {
            let thread = self.reference.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            let table_index = vm_thread.get_top();
            vm_thread
                .raw_geti(table_index, index)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            vm_thread
                .push_nil()
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            vm_thread
                .raw_seti(table_index, index)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            V::from_stack(&thread, -1)
        }
    }

    /// Inserts a value at an array index without invoking metamethods.
    pub fn raw_insert(&self, index: usize, value: impl IntoLua<'lua>) -> Result<(), Error> {
        let size = self.raw_len_cached();
        if index < 1 || index > size + 1 {
            return Err(Error::index_out_of_bounds());
        }
        self.invalidate_safe_env();

        unsafe {
            let thread = self.reference.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            let table_index = vm_thread.get_top();
            for current in (index..=size).rev() {
                vm_thread
                    .raw_geti(table_index, raw_array_index(current)?)
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
                vm_thread
                    .raw_seti(table_index, raw_array_index(current + 1)?)
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            }
            value.push_into_stack(&thread)?;
            vm_thread
                .raw_seti(table_index, raw_array_index(index)?)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Removes a key without invoking metamethods.
    ///
    /// Array indices shift later sequence values down.
    pub fn raw_remove(&self, key: impl IntoLua<'lua>) -> Result<(), Error> {
        let thread = self.thread();
        let key = key.into_lua(thread.lua_ref())?;
        let Some(index) = key.as_table_array_index()? else {
            return self.raw_set(key, Value::Nil);
        };

        let size = raw_array_index(self.raw_len_cached())?;
        if index > i64::from(size) {
            return Err(Error::index_out_of_bounds());
        }
        let index = i32::try_from(index).map_err(|_| Error::index_out_of_bounds())?;
        self.invalidate_safe_env();

        unsafe {
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            let table_index = vm_thread.get_top();
            for current in index..size {
                vm_thread
                    .raw_geti(table_index, current + 1)
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
                vm_thread
                    .raw_seti(table_index, current)
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            }
            vm_thread
                .push_nil()
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            vm_thread
                .raw_seti(table_index, size)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Iterates over consecutive values starting at key 1 and ending at `nil`.
    pub fn sequence_values<V>(&self) -> TableSequence<'_, 'lua, V>
    where
        V: FromLua<'lua>,
    {
        TableSequence {
            table: self,
            index: 1,
            _marker: PhantomData,
        }
    }

    /// Calls a function for each consecutive sequence value.
    pub fn for_each_value<V>(
        &self,
        mut function: impl FnMut(V) -> Result<(), Error>,
    ) -> Result<(), Error>
    where
        V: FromLua<'lua>,
    {
        for value in self.sequence_values() {
            function(value?)?;
        }
        Ok(())
    }

    fn raw_len_cached(&self) -> usize {
        unsafe { self.table.getn() as usize }
    }

    fn is_empty_cached(&self) -> bool {
        unsafe {
            for index in 0..(*self.table.as_ptr()).size_array as usize {
                if !self.table.array_slot(index).is_nil() {
                    return false;
                }
            }

            if !self.table.has_dummy_node() {
                for index in 0..self.table.node_count() {
                    if !self.table.node(index as i32).value_unchecked().is_nil() {
                        return false;
                    }
                }
            }

            true
        }
    }
}

impl<'lua, V> Iterator for TableSequence<'_, 'lua, V>
where
    V: FromLua<'lua>,
{
    type Item = Result<V, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        let result = (|| unsafe {
            let thread = self.table.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.table.push_to(&thread)?;
            let table_index = vm_thread.get_top();
            let tag = vm_thread
                .raw_geti(table_index, raw_array_index(self.index)?)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            if tag == luau_vm::types::LUA_TNIL {
                return Ok(None);
            }

            self.index += 1;
            V::from_stack(&thread, -1).map(Some)
        })();

        match result {
            Ok(Some(value)) => Some(Ok(value)),
            Ok(None) => None,
            Err(error) => Some(Err(error)),
        }
    }
}

fn raw_array_index(index: usize) -> Result<i32, Error> {
    i32::try_from(index).map_err(|_| Error::index_out_of_bounds())
}