runa/crates/runa-rune/src/execution.rs

use std::cell::{Cell, RefCell};
use std::future::{Future, poll_fn};
use std::num::NonZeroUsize;
use std::pin::Pin;
use std::rc::Rc;
use std::sync::Arc;
use std::task::{Context, Poll};

use rune::runtime::{Args, Unit, VmError};
use rune::{FromValue, Vm};

use crate::ExperimentResult;

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct InvocationProgress {
    pub vm_steps: u64,
    pub instruction_pointer: usize,
    pub call_depth: usize,
}

/// A worker-local execution slice; native callbacks must remain nonblocking.
#[derive(Debug)]
pub enum Slice<T> {
    Completed(T),
    Awaiting,
    BudgetExhausted,
    Failed(Box<dyn std::error::Error>),
}

type ExecutionFuture<T> = Pin<Box<dyn Future<Output = ExperimentResult<T>>>>;

pub const MAX_INSPECTION_FRAMES: usize = 64;

/// Owned instruction addresses only; no Rune stack values escape the worker.
#[derive(Clone)]
pub struct InvocationStack {
    pub unit: Arc<Unit>,
    pub instruction_pointer: usize,
    pub return_addresses: Vec<usize>,
    pub call_depth: usize,
    pub awaiting: bool,
}

impl InvocationStack {
    fn capture(vm: &Vm) -> Self {
        let mut stack = Self {
            unit: Arc::clone(vm.unit()),
            instruction_pointer: vm.ip(),
            return_addresses: Vec::new(),
            call_depth: vm.call_frames().len(),
            awaiting: false,
        };
        stack.update(vm);
        stack
    }

    fn update(&mut self, vm: &Vm) {
        if !Arc::ptr_eq(&self.unit, vm.unit()) {
            self.unit = Arc::clone(vm.unit());
        }
        self.instruction_pointer = vm.ip();
        self.call_depth = vm.call_frames().len();
        self.return_addresses.clear();
        self.return_addresses.extend(
            vm.call_frames()
                .iter()
                .rev()
                .take(MAX_INSPECTION_FRAMES - 1)
                .map(|frame| frame.ip),
        );
        self.awaiting = false;
    }
}

pub struct Invocation<T> {
    future: Option<ExecutionFuture<T>>,
    stepped: Rc<Cell<bool>>,
    progress: Rc<Cell<InvocationProgress>>,
    stack: Rc<RefCell<InvocationStack>>,
}

impl<T: FromValue + 'static> Invocation<T> {
    pub fn start(mut vm: Vm, name: &str, args: impl Args) -> Result<Self, VmError> {
        let execution = vm.execute([name], args)?;
        let entrypoint = execution.vm().ip();
        let mut execution = execution.into_owned();
        // Rune 0.14.2's into_owned transfers the stack but resets the VM IP.
        execution.vm_mut().set_ip(entrypoint);
        let stepped = Rc::new(Cell::new(false));
        let marker = Rc::clone(&stepped);
        let progress = Rc::new(Cell::new(InvocationProgress {
            instruction_pointer: entrypoint,
            ..InvocationProgress::default()
        }));
        let inspector = Rc::clone(&progress);
        let stack = Rc::new(RefCell::new(InvocationStack::capture(execution.vm())));
        let stack_inspector = Rc::clone(&stack);
        let future = Box::pin(async move {
            loop {
                let result = execution.async_step().await;
                let vm = execution.vm();
                let previous = inspector.get();
                stack_inspector.borrow_mut().update(vm);
                inspector.set(InvocationProgress {
                    vm_steps: previous.vm_steps.saturating_add(1),
                    instruction_pointer: vm.ip(),
                    call_depth: vm.call_frames().len(),
                });
                if let Some(value) = result.into_result()? {
                    return Ok(rune::from_value(value)?);
                }
                marker.set(true);
                let mut yielded = false;
                poll_fn(|cx| {
                    if yielded {
                        Poll::Ready(())
                    } else {
                        yielded = true;
                        cx.waker().wake_by_ref();
                        Poll::Pending
                    }
                })
                .await;
            }
        });
        Ok(Self {
            future: Some(future),
            stepped,
            progress,
            stack,
        })
    }

    /// Get instruction telemetry from the most recently completed VM step.
    pub fn progress(&self) -> InvocationProgress {
        self.progress.get()
    }

    pub fn stack(&self) -> InvocationStack {
        self.stack.borrow().clone()
    }

    /// Bound VM steps per call, preserving the exact pending async step.
    pub fn poll_slice(&mut self, cx: &mut Context<'_>, budget: NonZeroUsize) -> Slice<T> {
        let Some(future) = self.future.as_mut() else {
            return Slice::Failed("invocation has already finished".into());
        };
        for _ in 0..budget.get() {
            self.stepped.set(false);
            match future.as_mut().poll(cx) {
                Poll::Ready(result) => {
                    self.future = None;
                    return match result {
                        Ok(value) => Slice::Completed(value),
                        Err(error) => Slice::Failed(error),
                    };
                }
                Poll::Pending if !self.stepped.get() => {
                    self.stack.borrow_mut().awaiting = true;
                    return Slice::Awaiting;
                }
                Poll::Pending => {}
            }
        }
        Slice::BudgetExhausted
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::CompiledProgram;
    use rune::Module;
    use std::sync::Mutex;
    use std::task::{Wake, Waker};

    use std::sync::atomic::{AtomicBool, Ordering};

    #[test]
    fn debug_signatures_and_entry_stack_reads_do_not_execute_or_consume_values()
    -> ExperimentResult<()> {
        use rune::runtime::{InstAddress, debug::DebugArgs};

        let program =
            CompiledProgram::compile("pub fn probe(input) { let local = input + 1; local + 1 }")?;
        let mut vm = program.vm();
        let mut execution = vm.execute(["probe"], (40_i64,))?;
        let vm = execution.vm();
        let (_, signature) = vm
            .unit()
            .debug_info()
            .and_then(|debug| debug.function_at(vm.ip()))
            .ok_or("missing entry signature")?;
        let DebugArgs::Named(names) = &signature.args else {
            return Err("named argument signature unavailable".into());
        };
        assert_eq!(
            names
                .iter()
                .map(|name| name.as_ref())
                .collect::<Vec<&str>>(),
            ["input"]
        );
        for _ in 0..2 {
            assert_eq!(vm.stack().at(InstAddress::ZERO).as_integer::<i64>()?, 40);
        }
        for _ in 0..100 {
            if let Some(value) = execution.step().into_result()? {
                assert_eq!(rune::from_value::<i64>(value)?, 42);
                return Ok(());
            }
        }
        Err("probe did not complete".into())
    }

    #[test]
    fn stack_frames_map_nested_calls_and_return_addresses_to_exact_source() -> ExperimentResult<()>
    {
        let source = "pub fn handle() {\n    outer()\n}\nfn outer() {\n    inner()\n}\nfn inner() {\n    loop {}\n}";
        let program = CompiledProgram::compile(source)?;
        let mut invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        let mut cx = Context::from_waker(Waker::noop());
        for _ in 0..100 {
            assert!(matches!(
                invocation.poll_slice(&mut cx, NonZeroUsize::new(1).unwrap()),
                Slice::BudgetExhausted
            ));
            let inspection =
                program.inspect_stack(source, &invocation.stack(), invocation.progress().vm_steps);
            if inspection.frames.len() == 3
                && inspection.frames[0]
                    .range
                    .is_some_and(|range| range.start.line == 7)
            {
                assert_eq!(
                    inspection
                        .frames
                        .iter()
                        .map(|frame| frame.function.as_deref().unwrap())
                        .collect::<Vec<_>>(),
                    ["inner", "outer", "handle"]
                );
                assert_eq!(
                    inspection
                        .frames
                        .iter()
                        .map(|frame| frame.range.unwrap().start.line)
                        .collect::<Vec<_>>(),
                    [7, 4, 1]
                );
                assert!(
                    inspection.frames[1].source_instruction_pointer.unwrap()
                        < inspection.frames[1].instruction_pointer
                );
                assert_eq!(inspection.call_depth, 2);
                assert!(!inspection.awaiting);
                assert!(!inspection.frames_truncated);
                return Ok(());
            }
        }
        Err("nested source stack was not observed".into())
    }

    #[test]
    fn unavailable_debug_locations_are_not_guessed() -> ExperimentResult<()> {
        let source = "pub fn handle() { loop {} }";
        let program = CompiledProgram::compile(source)?;
        let other = CompiledProgram::compile("pub fn handle() { 1 }")?;
        let mut stack = InvocationStack::capture(&program.vm());
        // A different compiled unit cannot be mapped against this source.
        let invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        let inspection = other.inspect_stack("pub fn handle() { 1 }", &invocation.stack(), 0);
        assert!(inspection.frames.iter().all(|frame| frame.range.is_none()));
        stack.instruction_pointer = usize::MAX;
        stack.return_addresses.clear();
        stack.call_depth = 0;
        let inspection = program.inspect_stack(source, &stack, 0);
        assert_eq!(inspection.frames[0].range, None);
        assert_eq!(inspection.frames[0].function, None);
        assert_eq!(inspection.frames[0].source_instruction_pointer, None);
        Ok(())
    }

    #[test]
    fn deep_stack_inspection_is_bounded_and_reports_truncation() -> ExperimentResult<()> {
        let source = "pub fn handle() { recurse(100) } fn recurse(n) { if n == 0 { loop {} } else { recurse(n - 1) } }";
        let program = CompiledProgram::compile(source)?;
        let mut invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        let mut cx = Context::from_waker(Waker::noop());
        for _ in 0..100 {
            assert!(matches!(
                invocation.poll_slice(&mut cx, NonZeroUsize::new(32).unwrap()),
                Slice::BudgetExhausted
            ));
            let inspection =
                program.inspect_stack(source, &invocation.stack(), invocation.progress().vm_steps);
            if inspection.call_depth > MAX_INSPECTION_FRAMES {
                assert_eq!(inspection.frames.len(), MAX_INSPECTION_FRAMES);
                assert!(inspection.frames_truncated);
                return Ok(());
            }
        }
        Err("deep recursive stack was not observed".into())
    }

    #[test]
    fn shutdown_invocation_preserves_state_argument() -> ExperimentResult<()> {
        let program = CompiledProgram::compile(
            "pub fn handle(ctx, state, sender, message) { let n = 0; while n < 20 { n += 1; } state + message } pub fn on_shutdown(state) { state + 100 }",
        )?;
        assert_eq!(
            rune::from_value::<i64>(program.vm().call(["on_shutdown"], (1i64,))?)?,
            101
        );
        let mut invocation = Invocation::<i64>::start(program.vm(), "on_shutdown", (1i64,))?;
        let mut cx = Context::from_waker(Waker::noop());
        loop {
            match invocation.poll_slice(&mut cx, NonZeroUsize::new(32).unwrap()) {
                Slice::Completed(value) => {
                    assert_eq!(value, 101);
                    break;
                }
                Slice::Failed(error) => return Err(error),
                _ => {}
            }
        }
        Ok(())
    }

    struct Signal {
        waker: Mutex<Option<Waker>>,
        ready: AtomicBool,
        dropped: AtomicBool,
    }
    struct PendingOperation(Arc<Signal>);

    impl Future for PendingOperation {
        type Output = i64;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<i64> {
            if self.0.ready.load(Ordering::SeqCst) {
                Poll::Ready(7)
            } else {
                *self.0.waker.lock().unwrap() = Some(cx.waker().clone());
                Poll::Pending
            }
        }
    }

    impl Drop for PendingOperation {
        fn drop(&mut self) {
            self.0.dropped.store(true, Ordering::SeqCst);
        }
    }

    struct WakeCounter(std::sync::atomic::AtomicUsize);
    impl Wake for WakeCounter {
        fn wake(self: Arc<Self>) {
            self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }
    }

    #[test]
    fn budgets_await_wake_resume_and_cancel() -> ExperimentResult<()> {
        let signal = Arc::new(Signal {
            waker: Mutex::new(None),
            ready: AtomicBool::new(false),
            dropped: AtomicBool::new(false),
        });
        let mut module = Module::new();
        let native_signal = Arc::clone(&signal);
        module
            .function("wait", move || PendingOperation(Arc::clone(&native_signal)))
            .build()?;
        let source = "pub async fn handle() { let n = 0; while n < 10 { n += 1; } let x = wait().await; while n < 20 { n += 1; } n + x }";
        let program = CompiledProgram::compile_with_module(source, module)?;
        let counter = Arc::new(WakeCounter(std::sync::atomic::AtomicUsize::new(0)));
        let waker = Waker::from(Arc::clone(&counter));
        let mut cx = Context::from_waker(&waker);
        let budget = NonZeroUsize::new(1).unwrap();
        let mut invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        assert!(matches!(
            invocation.poll_slice(&mut cx, budget),
            Slice::BudgetExhausted
        ));
        for _ in 0..1_000 {
            if matches!(invocation.poll_slice(&mut cx, budget), Slice::Awaiting) {
                break;
            }
        }
        let registered = signal
            .waker
            .lock()
            .unwrap()
            .take()
            .expect("await registered waker");
        let stack = invocation.stack();
        assert!(stack.awaiting);
        let inspection = program.inspect_stack(source, &stack, invocation.progress().vm_steps);
        let range = inspection.frames[0].range.unwrap();
        assert_eq!(range.start.line, 0);
        assert_eq!(
            &source[range.start.character..range.end.character],
            "wait().await"
        );
        assert!(matches!(
            invocation.poll_slice(&mut cx, budget),
            Slice::Awaiting
        ));
        signal.ready.store(true, Ordering::SeqCst);
        registered.wake();
        assert!(counter.0.load(std::sync::atomic::Ordering::Relaxed) > 0);
        let mut completed = false;
        let mut exhausted_after_await = false;
        for _ in 0..1_000 {
            match invocation.poll_slice(&mut cx, budget) {
                Slice::Completed(value) => {
                    assert_eq!(value, 27);
                    completed = true;
                    break;
                }
                Slice::Failed(error) => return Err(error),
                Slice::BudgetExhausted => exhausted_after_await = true,
                Slice::Awaiting => return Err("ready operation remained pending".into()),
            }
        }
        assert!(completed);
        assert!(exhausted_after_await);

        signal.ready.store(false, Ordering::SeqCst);
        signal.dropped.store(false, Ordering::SeqCst);
        let mut cancelled = Invocation::<i64>::start(program.vm(), "handle", ())?;
        for _ in 0..1_000 {
            if matches!(cancelled.poll_slice(&mut cx, budget), Slice::Awaiting) {
                break;
            }
        }
        assert!(!signal.dropped.load(Ordering::SeqCst));
        drop(cancelled);
        assert!(signal.dropped.load(Ordering::SeqCst));
        Ok(())
    }

    #[test]
    fn errors_and_finished_invocations_have_typed_outcomes() -> ExperimentResult<()> {
        let program = CompiledProgram::compile("pub fn handle() { 1 / 0 }")?;
        let mut invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        let mut cx = Context::from_waker(Waker::noop());
        assert!(matches!(
            invocation.poll_slice(&mut cx, NonZeroUsize::new(100).unwrap()),
            Slice::Failed(_)
        ));
        assert!(matches!(
            invocation.poll_slice(&mut cx, NonZeroUsize::new(100).unwrap()),
            Slice::Failed(_)
        ));
        assert!(Invocation::<i64>::start(program.vm(), "missing", ()).is_err());
        Ok(())
    }
}