Skip to main content

rune/runtime/
vm_execution.rs

1use core::fmt;
2use core::future::Future;
3use core::mem::{replace, take};
4use core::pin::{pin, Pin};
5use core::task::{ready, Context, Poll, RawWaker, RawWakerVTable, Waker};
6
7use crate::alloc::prelude::*;
8use crate::async_vm_try;
9use crate::runtime::budget::Budget;
10use crate::runtime::{budget, Awaited};
11use crate::shared::AssertSend;
12use crate::sync::Arc;
13
14use super::{
15    Address, GeneratorState, Globals, Output, RuntimeContext, Unit, Value, Vm, VmDiagnostics,
16    VmError, VmErrorKind, VmHalt, VmHaltInfo,
17};
18
19static COMPLETE_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
20    |_| RawWaker::new(&(), &COMPLETE_WAKER_VTABLE),
21    |_| {},
22    |_| {},
23    |_| {},
24);
25
26// SAFETY: This waker does nothing.
27static COMPLETE_WAKER: Waker =
28    unsafe { Waker::from_raw(RawWaker::new(&(), &COMPLETE_WAKER_VTABLE)) };
29
30/// The state of an execution. We keep track of this because it's important to
31/// correctly interact with functions that yield (like generators and streams)
32/// by initially just calling the function, then by providing a value pushed
33/// onto the stack.
34#[derive(Debug, Clone, Copy, PartialEq)]
35#[non_exhaustive]
36pub(crate) enum ExecutionState {
37    /// The initial state of an execution.
38    Initial,
39    /// execution is waiting.
40    Resumed(Output),
41    /// Suspended execution.
42    Suspended,
43    /// Execution exited.
44    Exited(Option<Address>),
45}
46
47impl fmt::Display for ExecutionState {
48    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49        match self {
50            ExecutionState::Initial => write!(f, "initial"),
51            ExecutionState::Resumed(out) => write!(f, "resumed({out})"),
52            ExecutionState::Suspended => write!(f, "suspended"),
53            ExecutionState::Exited(..) => write!(f, "exited"),
54        }
55    }
56}
57
58#[derive(TryClone)]
59#[try_clone(crate)]
60pub(crate) struct VmExecutionState {
61    pub(crate) context: Option<Arc<RuntimeContext>>,
62    pub(crate) unit: Option<Arc<Unit>>,
63    pub(crate) globals: Option<Globals>,
64}
65
66/// The execution environment for a virtual machine.
67///
68/// When an execution is dropped, the stack of the stack of the head machine
69/// will be cleared.
70pub struct VmExecution<T> {
71    /// The current head vm which holds the execution.
72    vm: T,
73    /// The state of an execution.
74    state: ExecutionState,
75    /// Indicates the current stack of suspended contexts.
76    states: Vec<VmExecutionState>,
77}
78
79impl<T> VmExecution<T> {
80    /// Construct an execution from a virtual machine.
81    #[inline]
82    pub(crate) fn new(vm: T) -> Self {
83        Self {
84            vm,
85            state: ExecutionState::Initial,
86            states: Vec::new(),
87        }
88    }
89
90    /// Get a reference to the current virtual machine.
91    #[inline]
92    pub fn vm(&self) -> &Vm
93    where
94        T: AsRef<Vm>,
95    {
96        self.vm.as_ref()
97    }
98
99    /// Get a mutable reference the current virtual machine.
100    #[inline]
101    pub fn vm_mut(&mut self) -> &mut Vm
102    where
103        T: AsMut<Vm>,
104    {
105        self.vm.as_mut()
106    }
107}
108
109impl<T> VmExecution<T>
110where
111    T: AsMut<Vm>,
112{
113    /// Coerce the current execution into a generator if appropriate.
114    ///
115    /// ```
116    /// use rune::Vm;
117    /// use rune::sync::Arc;
118    ///
119    /// let mut sources = rune::sources! {
120    ///     entry => {
121    ///         pub fn main() {
122    ///             yield 1;
123    ///             yield 2;
124    ///         }
125    ///     }
126    /// };
127    ///
128    /// let unit = rune::prepare(&mut sources).build()?;
129    /// let unit = Arc::try_new(unit)?;
130    /// let mut vm = Vm::without_runtime(unit)?;
131    ///
132    /// let mut generator = vm.execute(["main"], ())?.into_generator();
133    ///
134    /// let mut n = 1i64;
135    ///
136    /// while let Some(value) = generator.next()? {
137    ///     let value: i64 = rune::from_value(value)?;
138    ///     assert_eq!(value, n);
139    ///     n += 1;
140    /// }
141    /// # Ok::<_, rune::support::Error>(())
142    /// ```
143    pub fn into_generator(self) -> VmGenerator<T> {
144        VmGenerator {
145            execution: Some(self),
146        }
147    }
148
149    /// Coerce the current execution into a stream if appropriate.
150    ///
151    /// ```
152    /// use rune::Vm;
153    /// use rune::sync::Arc;
154    ///
155    /// # futures_executor::block_on(async move {
156    /// let mut sources = rune::sources! {
157    ///     entry => {
158    ///         pub async fn main() {
159    ///             yield 1;
160    ///             yield 2;
161    ///         }
162    ///     }
163    /// };
164    ///
165    /// let unit = rune::prepare(&mut sources).build()?;
166    /// let unit = Arc::try_new(unit)?;
167    /// let mut vm = Vm::without_runtime(unit)?;
168    ///
169    /// let mut stream = vm.execute(["main"], ())?.into_stream();
170    ///
171    /// let mut n = 1i64;
172    ///
173    /// while let Some(value) = stream.next().await? {
174    ///     let value: i64 = rune::from_value(value)?;
175    ///     assert_eq!(value, n);
176    ///     n += 1;
177    /// }
178    /// # Ok::<_, rune::support::Error>(())
179    /// # })?;
180    /// # Ok::<_, rune::support::Error>(())
181    /// ```
182    pub fn into_stream(self) -> VmStream<T> {
183        VmStream {
184            execution: Some(self),
185        }
186    }
187}
188
189impl<T> VmExecution<T>
190where
191    T: AsMut<Vm>,
192{
193    /// Synchronously complete the current execution.
194    ///
195    /// # Errors
196    ///
197    /// If anything except the completion of the execution is encountered, this
198    /// will result in an error.
199    ///
200    /// To handle other outcomes and more configurability see
201    /// [`VmExecution::resume`].
202    pub fn complete(&mut self) -> Result<Value, VmError> {
203        self.resume().complete()?.into_complete()
204    }
205
206    /// Asynchronously complete the current execution.
207    ///
208    /// # Errors
209    ///
210    /// If anything except the completion of the execution is encountered, this
211    /// will result in an error.
212    ///
213    /// To handle other outcomes and more configurability see
214    /// [`VmExecution::resume`].
215    pub async fn async_complete(&mut self) -> Result<Value, VmError> {
216        self.resume().await?.into_complete()
217    }
218
219    /// Resume the current execution.
220    ///
221    /// To complete this operation synchronously, use [`VmResume::complete`].
222    ///
223    /// ## Resume with a value
224    ///
225    /// To resume an execution with a value, use [`VmResume::with_value`]. This
226    /// requires that the execution has yielded first, otherwise an error will
227    /// be produced.
228    ///
229    /// ## Resume with diagnostics
230    ///
231    /// To associated [`VmDiagnostics`] with the execution, use
232    /// [`VmResume::with_diagnostics`].
233    pub fn resume(&mut self) -> VmResume<'_, 'static, T> {
234        VmResume {
235            execution: self,
236            diagnostics: None,
237            awaited: None,
238            init: Some(Value::empty()),
239        }
240    }
241
242    /// End execution and perform debug checks.
243    pub(crate) fn end(&mut self) -> Result<Value, VmError> {
244        let ExecutionState::Exited(addr) = self.state else {
245            return Err(VmError::new(VmErrorKind::ExpectedExitedExecutionState {
246                actual: self.state,
247            }));
248        };
249
250        let value = match addr {
251            Some(addr) => self.vm.as_mut().stack().at(addr).clone(),
252            None => Value::unit(),
253        };
254
255        debug_assert!(self.states.is_empty(), "Execution states should be empty");
256        Ok(value)
257    }
258
259    /// Push a virtual machine state onto the execution.
260    #[tracing::instrument(skip_all)]
261    pub(crate) fn push_state(&mut self, state: VmExecutionState) -> Result<(), VmError> {
262        tracing::trace!("pushing suspended state");
263        let vm = self.vm.as_mut();
264        let context = state.context.map(|c| replace(vm.context_mut(), c));
265        let unit = state.unit.map(|u| replace(vm.unit_mut(), u));
266        let globals = state.globals.map(|g| replace(vm.globals_mut(), g));
267        self.states.try_push(VmExecutionState {
268            context,
269            unit,
270            globals,
271        })?;
272        Ok(())
273    }
274
275    /// Pop a virtual machine state from the execution and transfer the top of
276    /// the stack from the popped machine.
277    #[tracing::instrument(skip_all)]
278    fn pop_state(&mut self) -> Result<(), VmError> {
279        tracing::trace!("popping suspended state");
280
281        let state = self.states.pop().ok_or(VmErrorKind::NoRunningVm)?;
282        let vm = self.vm.as_mut();
283
284        if let Some(context) = state.context {
285            *vm.context_mut() = context;
286        }
287
288        if let Some(unit) = state.unit {
289            *vm.unit_mut() = unit;
290        }
291
292        if let Some(globals) = state.globals {
293            *vm.globals_mut() = globals;
294        }
295
296        Ok(())
297    }
298}
299
300impl VmExecution<&mut Vm> {
301    /// Convert the current execution into one which owns its virtual machine.
302    pub fn into_owned(self) -> VmExecution<Vm> {
303        let stack = take(self.vm.stack_mut());
304        let head = Vm::with_stack(self.vm.context().clone(), self.vm.unit().clone(), stack)
305            .with_globals(self.vm.globals().clone());
306
307        VmExecution {
308            vm: head,
309            states: self.states,
310            state: self.state,
311        }
312    }
313}
314
315/// A wrapper that makes [`VmExecution`] [`Send`].
316///
317/// This is accomplished by preventing any [`Value`] from escaping the [`Vm`].
318/// As long as this is maintained, it is safe to send the execution across,
319/// threads, and therefore schedule the future associated with the execution on
320/// a thread pool like Tokio's through [tokio::spawn].
321///
322/// [tokio::spawn]: https://docs.rs/tokio/0/tokio/runtime/struct.Runtime.html#method.spawn
323pub struct VmSendExecution(pub(crate) VmExecution<Vm>);
324
325// Safety: we wrap all APIs around the [VmExecution], preventing values from
326// escaping from contained virtual machine.
327unsafe impl Send for VmSendExecution {}
328
329impl VmSendExecution {
330    /// Complete the current execution with support for async instructions.
331    ///
332    /// This requires that the result of the Vm is converted into a
333    /// [crate::FromValue] that also implements [Send],  which prevents non-Send
334    /// values from escaping from the virtual machine.
335    pub fn complete(mut self) -> impl Future<Output = Result<Value, VmError>> + Send + 'static {
336        let future = async move { self.0.resume().await.and_then(VmOutcome::into_complete) };
337
338        // Safety: we wrap all APIs around the [VmExecution], preventing values
339        // from escaping from contained virtual machine.
340        unsafe { AssertSend::new(future) }
341    }
342
343    /// Alias for [`VmSendExecution::complete`].
344    #[deprecated = "Use `VmSendExecution::complete`"]
345    pub fn async_complete(self) -> impl Future<Output = Result<Value, VmError>> + Send + 'static {
346        self.complete()
347    }
348
349    /// Complete the current execution with support for async instructions.
350    ///
351    /// This requires that the result of the Vm is converted into a
352    /// [crate::FromValue] that also implements [Send],  which prevents non-Send
353    /// values from escaping from the virtual machine.
354    pub fn complete_with_diagnostics(
355        mut self,
356        diagnostics: &mut dyn VmDiagnostics,
357    ) -> impl Future<Output = Result<Value, VmError>> + Send + '_ {
358        let future = async move {
359            self.0
360                .resume()
361                .with_diagnostics(diagnostics)
362                .await
363                .and_then(VmOutcome::into_complete)
364        };
365
366        // Safety: we wrap all APIs around the [VmExecution], preventing values
367        // from escaping from contained virtual machine.
368        unsafe { AssertSend::new(future) }
369    }
370
371    /// Alias for [`VmSendExecution::complete_with_diagnostics`].
372    #[deprecated = "Use `VmSendExecution::complete_with_diagnostics`"]
373    pub fn async_complete_with_diagnostics(
374        self,
375        diagnostics: &mut dyn VmDiagnostics,
376    ) -> impl Future<Output = Result<Value, VmError>> + Send + '_ {
377        self.complete_with_diagnostics(diagnostics)
378    }
379}
380
381impl<T> TryClone for VmExecution<T>
382where
383    T: AsRef<Vm> + AsMut<Vm> + TryClone,
384{
385    #[inline]
386    fn try_clone(&self) -> Result<Self, rune_alloc::Error> {
387        Ok(Self {
388            vm: self.vm.try_clone()?,
389            state: self.state,
390            states: self.states.try_clone()?,
391        })
392    }
393}
394
395/// The outcome of completing an execution through a [`VmResume`] operation.
396#[non_exhaustive]
397pub enum VmOutcome {
398    /// A value has been produced by the execution returning.
399    Complete(Value),
400    /// A value has been yielded by the execution.
401    Yielded(Value),
402    /// The execution has been limited.
403    Limited,
404}
405
406impl VmOutcome {
407    /// Convert the outcome into a [`GeneratorState`].
408    ///
409    /// # Errors
410    ///
411    /// If the execution is not in a state compatible with producing a generator
412    /// state, such as having been completed or yielded, this will produce an
413    /// error.
414    pub fn into_generator_state(self) -> Result<GeneratorState, VmError> {
415        match self {
416            VmOutcome::Complete(value) => Ok(GeneratorState::Complete(value)),
417            VmOutcome::Yielded(value) => Ok(GeneratorState::Yielded(value)),
418            VmOutcome::Limited => Err(VmError::new(VmErrorKind::Halted {
419                halt: VmHaltInfo::Limited,
420            })),
421        }
422    }
423
424    /// Convert the outcome into a completed value.
425    ///
426    /// # Errors
427    ///
428    /// If the execution hasn't returned, this will produce an error.
429    pub fn into_complete(self) -> Result<Value, VmError> {
430        match self {
431            VmOutcome::Complete(value) => Ok(value),
432            VmOutcome::Yielded(..) => Err(VmError::new(VmErrorKind::Halted {
433                halt: VmHaltInfo::Yielded,
434            })),
435            VmOutcome::Limited => Err(VmError::new(VmErrorKind::Halted {
436                halt: VmHaltInfo::Limited,
437            })),
438        }
439    }
440}
441
442/// An execution that has been resumed.
443///
444/// This can either be completed as a future, which allows the execution to
445/// perform asynchronous operations, or it can be completed by calling
446/// [`VmResume::complete`] which will produce an error in case asynchronous
447/// operations that need to be suspended are encountered.
448pub struct VmResume<'this, 'diag, T> {
449    execution: &'this mut VmExecution<T>,
450    diagnostics: Option<&'diag mut dyn VmDiagnostics>,
451    init: Option<Value>,
452    awaited: Option<Awaited>,
453}
454
455impl<'this, 'diag, T> VmResume<'this, 'diag, T> {
456    /// Associated a budget with the resumed execution.
457    pub fn with_budget(self, budget: usize) -> Budget<Self> {
458        budget::with(budget, self)
459    }
460
461    /// Associate a value with the resumed execution.
462    ///
463    /// This is necessary to provide a value for a generator which has yielded.
464    pub fn with_value(self, value: Value) -> VmResume<'this, 'diag, T> {
465        Self {
466            init: Some(value),
467            ..self
468        }
469    }
470
471    /// Associate diagnostics with the execution.
472    pub fn with_diagnostics<'a>(
473        self,
474        diagnostics: &'a mut dyn VmDiagnostics,
475    ) -> VmResume<'this, 'a, T> {
476        VmResume {
477            execution: self.execution,
478            diagnostics: Some(diagnostics),
479            init: self.init,
480            awaited: self.awaited,
481        }
482    }
483}
484
485impl<'this, 'diag, T> VmResume<'this, 'diag, T>
486where
487    T: AsMut<Vm>,
488{
489    /// Try to synchronously complete the run, returning the generator state it produced.
490    ///
491    /// This will error if the execution is suspended through awaiting.
492    pub fn complete(self) -> Result<VmOutcome, VmError> {
493        let this = pin!(self);
494        let mut cx = Context::from_waker(&COMPLETE_WAKER);
495
496        match this.poll(&mut cx) {
497            Poll::Ready(result) => result,
498            Poll::Pending => Err(VmError::new(VmErrorKind::Halted {
499                halt: VmHaltInfo::Awaited,
500            })),
501        }
502    }
503}
504
505impl<'this, 'diag, T> Future for VmResume<'this, 'diag, T>
506where
507    T: AsMut<Vm>,
508{
509    type Output = Result<VmOutcome, VmError>;
510
511    #[inline]
512    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
513        // SAFETY: We are ensuring that we never move this value or any
514        // projected fields.
515        let this = unsafe { Pin::get_unchecked_mut(self) };
516
517        if let Some(value) = this.init.take() {
518            let state = replace(&mut this.execution.state, ExecutionState::Suspended);
519
520            if let ExecutionState::Resumed(out) = state {
521                let vm = this.execution.vm.as_mut();
522                async_vm_try!(vm.stack_mut().store(out, value));
523            }
524        }
525
526        loop {
527            let vm = this.execution.vm.as_mut();
528
529            if let Some(awaited) = &mut this.awaited {
530                let awaited = unsafe { Pin::new_unchecked(awaited) };
531                async_vm_try!(ready!(awaited.poll(cx, vm)));
532                this.awaited = None;
533            }
534
535            let result = vm.run(match this.diagnostics {
536                Some(ref mut value) => Some(&mut **value),
537                None => None,
538            });
539
540            match async_vm_try!(VmError::with_vm(result, vm)) {
541                VmHalt::Exited(addr) => {
542                    this.execution.state = ExecutionState::Exited(addr);
543                }
544                VmHalt::Awaited(awaited) => {
545                    this.awaited = Some(awaited);
546                    continue;
547                }
548                VmHalt::VmCall(vm_call) => {
549                    async_vm_try!(vm_call.into_execution(this.execution));
550                    continue;
551                }
552                VmHalt::Yielded(addr, out) => {
553                    let value = match addr {
554                        Some(addr) => vm.stack().at(addr).clone(),
555                        None => Value::unit(),
556                    };
557
558                    this.execution.state = ExecutionState::Resumed(out);
559                    return Poll::Ready(Ok(VmOutcome::Yielded(value)));
560                }
561                VmHalt::Limited => {
562                    return Poll::Ready(Ok(VmOutcome::Limited));
563                }
564            }
565
566            if this.execution.states.is_empty() {
567                let value = async_vm_try!(this.execution.end());
568                return Poll::Ready(Ok(VmOutcome::Complete(value)));
569            }
570
571            async_vm_try!(this.execution.pop_state());
572        }
573    }
574}
575
576/// A [`VmExecution`] that can be used with a generator api.
577pub struct VmGenerator<T> {
578    execution: Option<VmExecution<T>>,
579}
580
581impl<T> VmGenerator<T>
582where
583    T: AsMut<Vm>,
584{
585    /// Get the next value produced by this generator.
586    ///
587    /// See [`VmExecution::into_generator`].
588    pub fn next(&mut self) -> Result<Option<Value>, VmError> {
589        let Some(execution) = &mut self.execution else {
590            return Ok(None);
591        };
592
593        let outcome = execution.resume().complete()?;
594
595        match outcome {
596            VmOutcome::Complete(_) => {
597                self.execution = None;
598                Ok(None)
599            }
600            VmOutcome::Yielded(value) => Ok(Some(value)),
601            VmOutcome::Limited => Err(VmError::new(VmErrorKind::Halted {
602                halt: VmHaltInfo::Limited,
603            })),
604        }
605    }
606
607    /// Resume the generator with a value and get the next [`GeneratorState`].
608    ///
609    /// See [`VmExecution::into_generator`].
610    pub fn resume(&mut self, value: Value) -> Result<GeneratorState, VmError> {
611        let execution = self
612            .execution
613            .as_mut()
614            .ok_or(VmErrorKind::GeneratorComplete)?;
615
616        let outcome = execution.resume().with_value(value).complete()?;
617
618        match outcome {
619            VmOutcome::Complete(value) => {
620                self.execution = None;
621                Ok(GeneratorState::Complete(value))
622            }
623            VmOutcome::Yielded(value) => Ok(GeneratorState::Yielded(value)),
624            VmOutcome::Limited => Err(VmError::new(VmErrorKind::Halted {
625                halt: VmHaltInfo::Limited,
626            })),
627        }
628    }
629}
630
631/// A [`VmExecution`] that can be used with a stream api.
632pub struct VmStream<T> {
633    execution: Option<VmExecution<T>>,
634}
635
636impl<T> VmStream<T>
637where
638    T: AsMut<Vm>,
639{
640    /// Get the next value produced by this stream.
641    ///
642    /// See [`VmExecution::into_stream`].
643    pub async fn next(&mut self) -> Result<Option<Value>, VmError> {
644        let Some(execution) = &mut self.execution else {
645            return Ok(None);
646        };
647
648        match execution.resume().await? {
649            VmOutcome::Complete(value) => {
650                self.execution = None;
651                Ok(Some(value))
652            }
653            VmOutcome::Yielded(..) => Ok(None),
654            VmOutcome::Limited => Err(VmError::new(VmErrorKind::Halted {
655                halt: VmHaltInfo::Limited,
656            })),
657        }
658    }
659
660    /// Resume the stream with a value and return the next [`GeneratorState`].
661    ///
662    /// See [`VmExecution::into_stream`].
663    pub async fn resume(&mut self, value: Value) -> Result<GeneratorState, VmError> {
664        let execution = self
665            .execution
666            .as_mut()
667            .ok_or(VmErrorKind::GeneratorComplete)?;
668
669        match execution.resume().with_value(value).await? {
670            VmOutcome::Complete(value) => {
671                self.execution = None;
672                Ok(GeneratorState::Complete(value))
673            }
674            VmOutcome::Yielded(value) => Ok(GeneratorState::Yielded(value)),
675            VmOutcome::Limited => Err(VmError::new(VmErrorKind::Halted {
676                halt: VmHaltInfo::Limited,
677            })),
678        }
679    }
680}