luau 0.732.0

use core::fmt;
use luau_vm::VmErrorResult;
use luau_vm::native::{NativeCallContext, NativeCallResult};
use luau_vm::thread::{StackGuard, Thread as VmThread};

use crate::error::Error;
use crate::lua::runtime::RuntimeData;
use crate::table::Table;
use crate::thread::Thread;
use crate::value::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, LuaType, Value, ValueRef};

mod bind;
pub(crate) mod callback;
mod debug;

pub use debug::{CoverageInfo, FunctionInfo};

/// A handle to a Luau function.
pub struct Function<'lua> {
    handle: FunctionHandle<'lua>,
}

enum FunctionHandle<'lua> {
    Rooted(ValueRef<'lua>),
    BorrowedStack {
        source_thread: &'lua VmThread,
        reference_thread: &'lua VmThread,
        runtime: &'lua RuntimeData,
        index: i32,
        pointer: *const (),
    },
}

impl<'lua> Function<'lua> {
    pub(crate) fn from_ref(
        thread: &'lua VmThread,
        runtime: &'lua RuntimeData,
        reference: i32,
        pointer: *const (),
    ) -> Self {
        Self {
            handle: FunctionHandle::Rooted(ValueRef::new(thread, runtime, reference, pointer)),
        }
    }

    pub(crate) unsafe fn from_stack(thread: &Thread<'lua>, index: i32) -> VmErrorResult<Self> {
        debug_assert_ne!(unsafe { thread.as_vm().is_function(index) }, 0);
        Ok(Self {
            handle: FunctionHandle::Rooted(ValueRef::from_stack(thread, index)?),
        })
    }

    /// Borrows a function from a stack slot whose owner prevents mutation.
    ///
    /// # Safety
    ///
    /// `index` must contain a function and remain below the source thread's
    /// stack top, unchanged, for `'lua`. The source thread and VM must remain
    /// live for `'lua`.
    pub(crate) unsafe fn from_borrowed_stack(thread: &'lua Thread<'_>, index: i32) -> Self {
        let source_thread = thread.as_vm();
        debug_assert_ne!(unsafe { source_thread.is_function(index) }, 0);
        Self {
            handle: FunctionHandle::BorrowedStack {
                source_thread,
                reference_thread: thread.reference_thread(),
                runtime: thread.runtime(),
                index,
                pointer: unsafe { source_thread.to_pointer(index) },
            },
        }
    }

    /// Creates another rooted handle to this function.
    pub fn try_clone(&self) -> Result<Self, Error> {
        match &self.handle {
            FunctionHandle::Rooted(reference) => Ok(Self {
                handle: FunctionHandle::Rooted(reference.try_clone()?),
            }),
            FunctionHandle::BorrowedStack { .. } => unsafe {
                let thread = self.thread();
                let vm_thread = thread.as_vm();
                let _stack = StackGuard::new(vm_thread);
                self.push_to(vm_thread)?;
                Self::from_stack(&thread, -1)
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
            },
        }
    }

    pub(crate) fn push_to(&self, target: impl AsRef<VmThread>) -> Result<(), Error> {
        match &self.handle {
            FunctionHandle::Rooted(reference) => reference.push_to(target),
            FunctionHandle::BorrowedStack {
                source_thread,
                index,
                ..
            } => unsafe {
                let source_thread = *source_thread;
                let target = target.as_ref();
                if !source_thread.same_vm(target) {
                    return Err(Error::foreign_lua_handle());
                }

                if *source_thread == *target {
                    source_thread
                        .push_value(*index)
                        .map_err(|exit| Error::from_thread_exit(source_thread, exit))?;
                } else {
                    let _stack = StackGuard::new(source_thread);
                    source_thread
                        .push_value(*index)
                        .map_err(|exit| Error::from_thread_exit(source_thread, exit))?;
                    source_thread
                        .x_move(target, 1)
                        .map_err(|exit| Error::from_thread_exit(source_thread, exit))?;
                }
                Ok(())
            },
        }
    }

    pub(crate) fn thread(&self) -> Thread<'lua> {
        Thread::new(self.reference_thread(), self.runtime())
    }

    fn reference_thread(&self) -> &'lua VmThread {
        match &self.handle {
            FunctionHandle::Rooted(reference) => reference.reference_thread(),
            FunctionHandle::BorrowedStack {
                reference_thread, ..
            } => reference_thread,
        }
    }

    fn runtime(&self) -> &'lua RuntimeData {
        match &self.handle {
            FunctionHandle::Rooted(reference) => reference.runtime(),
            FunctionHandle::BorrowedStack { runtime, .. } => runtime,
        }
    }

    pub(crate) fn pointer(&self) -> *const () {
        match &self.handle {
            FunctionHandle::Rooted(reference) => reference.pointer(),
            FunctionHandle::BorrowedStack { pointer, .. } => *pointer,
        }
    }

    /// Returns a pointer that uniquely identifies this function.
    ///
    /// The pointer is intended for hashing and diagnostics and cannot be
    /// converted back into a function.
    pub fn to_pointer(&self) -> *const () {
        self.pointer()
    }

    /// Returns this function's environment.
    ///
    /// Native functions do not have an environment and return `None`.
    pub fn environment(&self) -> Result<Option<Table<'lua>>, Error> {
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(vm_thread)?;
            if vm_thread.is_native_function(-1) != 0 {
                return Ok(None);
            }

            vm_thread
                .get_fenv(-1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            if vm_thread.is_table(-1) == 0 {
                return Ok(None);
            }

            Table::from_stack(&thread, -1)
                .map(Some)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Sets the environment used to resolve this Luau function's globals.
    ///
    /// The target environment's safe-environment optimization is cleared so
    /// constants cached under the previous environment cannot be reused.
    /// Native functions are unchanged and return `false`.
    pub fn set_environment(&self, environment: &Table<'_>) -> Result<bool, Error> {
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(vm_thread)?;
            if vm_thread.is_native_function(-1) != 0 {
                return Ok(false);
            }

            let function_index = vm_thread.get_top();
            environment.push_to(&thread)?;
            vm_thread.set_safe_env(-1, 0);
            let changed = vm_thread.set_fenv(function_index) != 0;
            if changed {
                self.runtime().invalidate_managed_safe_env();
            }
            Ok(changed)
        }
    }

    /// Creates a deep clone of this function.
    ///
    /// Luau functions have their prototype and upvalues copied. Native
    /// functions return another handle to the same function.
    pub fn deep_clone(&self) -> Result<Self, Error> {
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(vm_thread)?;
            if vm_thread.is_native_function(-1) != 0 {
                return self.try_clone();
            }

            vm_thread
                .clone_function(-1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            Function::from_stack(&thread, -1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Creates a coroutine whose body is this function.
    pub fn into_thread(self) -> Result<Thread<'lua>, Error> {
        unsafe {
            let reference_thread = self.thread();
            let vm_thread = reference_thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            let thread = vm_thread
                .new_thread()
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            let thread_pointer = vm_thread.to_pointer(-1);
            let thread_reference = vm_thread
                .ref_value(-1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;

            let Err(error) = ({
                let target = Thread::new(&thread, reference_thread.runtime());
                self.push_to(&target)
            }) else {
                return Ok(Thread::from_ref(
                    self.reference_thread(),
                    self.runtime(),
                    thread,
                    thread_reference,
                    thread_pointer,
                ));
            };

            vm_thread.unref_value(thread_reference);
            Err(error)
        }
    }

    /// Calls this function with the given arguments.
    pub fn call<R>(&self, args: impl IntoLuaMulti<'lua>) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let stack = StackGuard::new(vm_thread);
            self.push_to(vm_thread)?;

            let arg_count = i32::try_from(args.push_into_stack_multi(&thread)?)
                .map_err(|_| Error::StackError)?;

            vm_thread
                .protected_call(arg_count, luau_vm::thread::LUA_MULTRET, 0)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;

            let result_count = vm_thread.get_top() - stack.top();
            R::from_stack_multi(&thread, stack.top(), result_count)
        }
    }
}

impl LuaType for Function<'_> {
    fn push_type_key(thread: impl AsRef<VmThread>) -> Result<(), Error> {
        let thread = thread.as_ref();
        unsafe {
            thread
                .push_native_function(type_key_nop, None)
                .map_err(|exit| Error::from_thread_exit(thread, exit))
        }
    }
}

fn type_key_nop(_: NativeCallContext<'_>) -> NativeCallResult {
    Ok(0)
}

impl<'lua, 'function> IntoLua<'lua> for Function<'function>
where
    'function: 'lua,
{
    fn into_lua(self, _: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        Ok(Value::Function(self))
    }

    unsafe fn push_into_stack(self, thread: &Thread<'lua>) -> Result<(), Error> {
        self.push_to(thread)
    }
}

impl<'lua, 'function> IntoLua<'lua> for &Function<'function>
where
    'function: 'lua,
{
    fn into_lua(self, _: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        self.try_clone().map(Value::Function)
    }

    unsafe fn push_into_stack(self, thread: &Thread<'lua>) -> Result<(), Error> {
        self.push_to(thread)
    }
}

impl<'lua> FromLua<'lua> for Function<'lua> {
    fn from_lua(value: Value<'lua>, _: crate::LuaRef<'lua>) -> Result<Self, Error> {
        match value {
            Value::Function(function) => Ok(function),
            value => Err(Error::from_lua_conversion(
                value.type_name(),
                "function",
                None,
            )),
        }
    }
}

impl PartialEq for Function<'_> {
    fn eq(&self, other: &Self) -> bool {
        (unsafe { self.reference_thread().same_vm(other.reference_thread()) })
            && self.pointer() == other.pointer()
    }
}

impl Eq for Function<'_> {}

impl fmt::Debug for Function<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_tuple("Function")
            .field(&self.pointer())
            .finish()
    }
}