luau 0.732.0

use core::cell::Cell;
use std::sync::Arc;

use luau_vm::Thread as VmThread;
use luau_vm::state::{GcInterrupt as VmGcInterrupt, GcPhase as VmGcPhase, InterruptRequest};
use luau_vm::{VmErrorResult, VmResult};

use crate::callback::{raise_callback_error, raise_callback_vm_error};
use crate::error::{Error, Result};
use crate::lua::runtime::RuntimeData;
use crate::lua::{Lua, LuaRef, StackInfo};

/// The control action requested at an execution interrupt point.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[non_exhaustive]
pub enum InterruptAction {
    /// Continue execution.
    #[default]
    Continue,
    /// Yield the active thread.
    Yield,
    /// Break the active thread.
    Break,
}

/// Controls when an interrupt handler is called.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[non_exhaustive]
pub enum InterruptMode {
    /// Call the handler at every interrupt point.
    #[default]
    Continuous,
    /// Call the handler only after its handle is requested.
    Requested,
}

/// A thread-safe handle for requesting interrupt delivery.
#[must_use = "retain the handle to request interrupts"]
#[derive(Clone, Debug)]
pub struct InterruptHandle {
    request: Arc<InterruptRequest>,
}

/// The collector phase active immediately before a garbage-collection step.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum GcPhase {
    /// The collector is between cycles.
    Pause,
    /// The collector is tracing reachable objects.
    Propagate,
    /// The collector is revisiting objects before the atomic phase.
    PropagateAgain,
    /// The collector is completing the mark phase.
    Atomic,
    /// The collector is reclaiming unreachable objects.
    Sweep,
}

/// The point at which a garbage-collection observer was invoked.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum GcInterruptStage {
    /// Immediately before a collector step.
    BeforeStep,
    /// Immediately after a collector step.
    AfterStep {
        /// The collector phase active before the step.
        previous_phase: GcPhase,
    },
}

/// Context supplied at a Luau execution safepoint.
#[derive(Clone, Copy)]
pub struct ExecutionInterruptContext<'callback> {
    lua: LuaRef<'callback>,
    deferred: &'callback Cell<bool>,
}

/// Context supplied during pattern matching.
#[derive(Clone, Copy)]
pub struct PatternInterruptContext<'callback> {
    lua: LuaRef<'callback>,
    deferred: &'callback Cell<bool>,
}

/// Context supplied around a garbage-collection step.
#[derive(Clone, Copy)]
pub struct GcInterruptContext<'callback> {
    lua: LuaRef<'callback>,
    stage: GcInterruptStage,
    deferred: &'callback Cell<bool>,
}

/// A coherent set of callbacks for VM interrupt points.
///
/// Continuous handlers run at every interrupt point. A request-driven handler
/// runs only after its state's [`InterruptHandle`] is requested, which can be
/// done from another thread. While a callback is running, nested execution,
/// pattern, and garbage-collection interrupts are suppressed.
///
/// The first interrupt point consumes a request. Default methods defer it so an
/// unimplemented callback cannot consume a request intended for another point.
/// An implemented callback can defer delivery through its context.
pub trait InterruptHandler: 'static {
    /// Returns how this handler should be delivered.
    ///
    /// This is read when the handler is installed.
    fn mode(&self) -> InterruptMode {
        InterruptMode::Continuous
    }

    /// Called at Luau execution safepoints.
    fn execution(&self, context: ExecutionInterruptContext<'_>) -> Result<InterruptAction> {
        context.defer_interrupt();
        Ok(InterruptAction::Continue)
    }

    /// Called periodically during string pattern matching.
    fn pattern(&self, context: PatternInterruptContext<'_>) -> Result<()> {
        context.defer_interrupt();
        Ok(())
    }

    /// Called immediately before and after garbage-collection steps.
    fn garbage_collection(&self, context: GcInterruptContext<'_>) {
        context.defer_interrupt();
    }
}

type ExecutionCallback =
    Box<dyn for<'callback> Fn(ExecutionInterruptContext<'callback>) -> Result<InterruptAction>>;
type PatternCallback = Box<dyn for<'callback> Fn(PatternInterruptContext<'callback>) -> Result<()>>;
type GcCallback = Box<dyn for<'callback> Fn(GcInterruptContext<'callback>)>;

/// Closure-based builder for an [`InterruptHandler`].
#[derive(Default)]
pub struct InterruptHooks {
    mode: InterruptMode,
    execution: Option<ExecutionCallback>,
    pattern: Option<PatternCallback>,
    garbage_collection: Option<GcCallback>,
}

impl InterruptHooks {
    /// Creates an empty interrupt callback set.
    pub const fn new() -> Self {
        Self {
            mode: InterruptMode::Continuous,
            execution: None,
            pattern: None,
            garbage_collection: None,
        }
    }

    /// Sets how these callbacks should be delivered.
    #[must_use]
    pub const fn set_mode(mut self, mode: InterruptMode) -> Self {
        self.mode = mode;
        self
    }

    /// Sets the execution interrupt callback.
    #[must_use]
    pub fn on_execution<F>(mut self, callback: F) -> Self
    where
        F: for<'callback> Fn(ExecutionInterruptContext<'callback>) -> Result<InterruptAction>
            + 'static,
    {
        self.execution = Some(Box::new(callback));
        self
    }

    /// Sets the pattern interrupt callback.
    #[must_use]
    pub fn on_pattern<F>(mut self, callback: F) -> Self
    where
        F: for<'callback> Fn(PatternInterruptContext<'callback>) -> Result<()> + 'static,
    {
        self.pattern = Some(Box::new(callback));
        self
    }

    /// Sets the garbage-collection observer.
    #[must_use]
    pub fn on_garbage_collection<F>(mut self, callback: F) -> Self
    where
        F: for<'callback> Fn(GcInterruptContext<'callback>) + 'static,
    {
        self.garbage_collection = Some(Box::new(callback));
        self
    }
}

impl InterruptHandler for InterruptHooks {
    fn mode(&self) -> InterruptMode {
        self.mode
    }

    fn execution(&self, context: ExecutionInterruptContext<'_>) -> Result<InterruptAction> {
        match &self.execution {
            Some(callback) => callback(context),
            None => {
                context.defer_interrupt();
                Ok(InterruptAction::Continue)
            }
        }
    }

    fn pattern(&self, context: PatternInterruptContext<'_>) -> Result<()> {
        match &self.pattern {
            Some(callback) => callback(context),
            None => {
                context.defer_interrupt();
                Ok(())
            }
        }
    }

    fn garbage_collection(&self, context: GcInterruptContext<'_>) {
        if let Some(callback) = &self.garbage_collection {
            callback(context);
        } else {
            context.defer_interrupt();
        }
    }
}

impl Lua {
    /// Returns a handle for requesting a request-driven interrupt handler.
    pub fn interrupt_handle(&self) -> InterruptHandle {
        self.runtime.callbacks().interrupt_handle()
    }

    /// Installs the interrupt callbacks for this state.
    ///
    /// Use [`Lua::interrupt_handle`] to trigger a handler whose mode is
    /// [`InterruptMode::Requested`].
    pub fn set_interrupt_handler(&mut self, handler: impl InterruptHandler) {
        self.runtime
            .callbacks_mut()
            .set_interrupt_handler(Box::new(handler));
        self.install_callbacks();
    }

    /// Removes the interrupt callbacks from this state.
    pub fn remove_interrupt_handler(&mut self) {
        self.runtime.callbacks_mut().remove_interrupt_handler();
        self.install_callbacks();
    }
}

impl InterruptHandle {
    pub(in crate::hooks) fn new() -> Self {
        Self {
            request: Arc::new(InterruptRequest::new()),
        }
    }

    /// Requests delivery at the next interrupt point.
    ///
    /// Multiple requests made before delivery are coalesced.
    pub fn request(&self) {
        self.request.request();
    }

    pub(in crate::hooks) fn clear(&self) {
        self.request.clear();
    }

    pub(in crate::hooks) fn as_ptr(&self) -> *const InterruptRequest {
        Arc::as_ptr(&self.request)
    }
}

impl<'callback> ExecutionInterruptContext<'callback> {
    pub(in crate::hooks) const fn new(
        lua: LuaRef<'callback>,
        deferred: &'callback Cell<bool>,
    ) -> Self {
        Self { lua, deferred }
    }

    /// Returns a borrowed reference to the active Luau state.
    pub const fn lua(&self) -> LuaRef<'callback> {
        self.lua
    }

    /// Returns whether the active thread can yield.
    pub fn is_yieldable(&self) -> bool {
        self.lua.is_yieldable()
    }

    /// Defers a requested delivery to a later interrupt point.
    ///
    /// This has no effect on a continuous handler.
    pub fn defer_interrupt(&self) {
        self.deferred.set(true);
    }
}

impl<'callback> PatternInterruptContext<'callback> {
    pub(in crate::hooks) const fn new(
        lua: LuaRef<'callback>,
        deferred: &'callback Cell<bool>,
    ) -> Self {
        Self { lua, deferred }
    }

    /// Returns a borrowed reference to the active Luau state.
    pub const fn lua(&self) -> LuaRef<'callback> {
        self.lua
    }

    /// Defers a requested delivery to a later interrupt point.
    ///
    /// This has no effect on a continuous handler.
    pub fn defer_interrupt(&self) {
        self.deferred.set(true);
    }
}

impl<'callback> GcInterruptContext<'callback> {
    pub(in crate::hooks) const fn new(
        lua: LuaRef<'callback>,
        stage: GcInterruptStage,
        deferred: &'callback Cell<bool>,
    ) -> Self {
        Self {
            lua,
            stage,
            deferred,
        }
    }

    /// Returns the point around the collection step at which this callback ran.
    pub const fn stage(&self) -> GcInterruptStage {
        self.stage
    }

    /// Inspects the active stack frame at `level`.
    pub fn inspect_stack<R>(
        &self,
        level: usize,
        inspect: impl FnOnce(&StackInfo<'_>) -> R,
    ) -> Result<Option<R>> {
        self.lua.inspect_stack(level, inspect)
    }

    /// Defers a requested delivery to a later interrupt point.
    ///
    /// This has no effect on a continuous handler.
    pub fn defer_interrupt(&self) {
        self.deferred.set(true);
    }
}

pub(in crate::hooks) fn execution_interrupt(thread: &VmThread) -> VmResult {
    let runtime = RuntimeData::from_thread(thread);
    let result = runtime.callbacks().invoke_interrupt(|handler, delivery| {
        runtime.with_thread(thread, || {
            let lua = LuaRef::new(thread, runtime);
            let action = match handler.execution(ExecutionInterruptContext::new(lua, delivery)) {
                Ok(action) => action,
                Err(error) => return raise_callback_error(thread, runtime, error),
            };

            match action {
                InterruptAction::Continue => Ok(()),
                InterruptAction::Yield if lua.is_yieldable() => unsafe {
                    thread.yield_current(0).map(|_| ())
                },
                InterruptAction::Yield => raise_callback_error(
                    thread,
                    runtime,
                    Error::runtime(
                        "interrupt handler attempted to yield from a non-yieldable boundary",
                    ),
                ),
                InterruptAction::Break => unsafe { thread.break_current().map(|_| ()) },
            }
        })
    });
    let Some(result) = result else {
        return Ok(());
    };
    result
}

pub(in crate::hooks) fn pattern_interrupt(thread: &VmThread) -> VmErrorResult {
    let runtime = RuntimeData::from_thread(thread);
    let result = runtime.callbacks().invoke_interrupt(|handler, delivery| {
        runtime.with_thread(thread, || {
            let lua = LuaRef::new(thread, runtime);
            handler
                .pattern(PatternInterruptContext::new(lua, delivery))
                .or_else(|error| raise_callback_vm_error(thread, runtime, error))
        })
    });
    let Some(result) = result else {
        return Ok(());
    };
    result
}

pub(in crate::hooks) fn gc_interrupt(thread: &VmThread, event: VmGcInterrupt) -> VmErrorResult {
    let runtime = RuntimeData::from_thread(thread);
    let stage = match event {
        VmGcInterrupt::BeforeStep => GcInterruptStage::BeforeStep,
        VmGcInterrupt::AfterStep { previous_phase } => GcInterruptStage::AfterStep {
            previous_phase: map_gc_phase(previous_phase),
        },
    };
    let result = runtime.callbacks().invoke_interrupt(|handler, delivery| {
        runtime.with_thread(thread, || {
            handler.garbage_collection(GcInterruptContext::new(
                LuaRef::new(thread, runtime),
                stage,
                delivery,
            ))
        })
    });
    let Some(()) = result else {
        return Ok(());
    };
    Ok(())
}

const fn map_gc_phase(phase: VmGcPhase) -> GcPhase {
    match phase {
        VmGcPhase::Pause => GcPhase::Pause,
        VmGcPhase::Propagate => GcPhase::Propagate,
        VmGcPhase::PropagateAgain => GcPhase::PropagateAgain,
        VmGcPhase::Atomic => GcPhase::Atomic,
        VmGcPhase::Sweep => GcPhase::Sweep,
    }
}