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
26static COMPLETE_WAKER: Waker =
28 unsafe { Waker::from_raw(RawWaker::new(&(), &COMPLETE_WAKER_VTABLE)) };
29
30#[derive(Debug, Clone, Copy, PartialEq)]
35#[non_exhaustive]
36pub(crate) enum ExecutionState {
37 Initial,
39 Resumed(Output),
41 Suspended,
43 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
66pub struct VmExecution<T> {
71 vm: T,
73 state: ExecutionState,
75 states: Vec<VmExecutionState>,
77}
78
79impl<T> VmExecution<T> {
80 #[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 #[inline]
92 pub fn vm(&self) -> &Vm
93 where
94 T: AsRef<Vm>,
95 {
96 self.vm.as_ref()
97 }
98
99 #[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 pub fn into_generator(self) -> VmGenerator<T> {
144 VmGenerator {
145 execution: Some(self),
146 }
147 }
148
149 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 pub fn complete(&mut self) -> Result<Value, VmError> {
203 self.resume().complete()?.into_complete()
204 }
205
206 pub async fn async_complete(&mut self) -> Result<Value, VmError> {
216 self.resume().await?.into_complete()
217 }
218
219 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 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 #[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 #[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 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
315pub struct VmSendExecution(pub(crate) VmExecution<Vm>);
324
325unsafe impl Send for VmSendExecution {}
328
329impl VmSendExecution {
330 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 unsafe { AssertSend::new(future) }
341 }
342
343 #[deprecated = "Use `VmSendExecution::complete`"]
345 pub fn async_complete(self) -> impl Future<Output = Result<Value, VmError>> + Send + 'static {
346 self.complete()
347 }
348
349 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 unsafe { AssertSend::new(future) }
369 }
370
371 #[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#[non_exhaustive]
397pub enum VmOutcome {
398 Complete(Value),
400 Yielded(Value),
402 Limited,
404}
405
406impl VmOutcome {
407 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 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
442pub 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 pub fn with_budget(self, budget: usize) -> Budget<Self> {
458 budget::with(budget, self)
459 }
460
461 pub fn with_value(self, value: Value) -> VmResume<'this, 'diag, T> {
465 Self {
466 init: Some(value),
467 ..self
468 }
469 }
470
471 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 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 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
576pub struct VmGenerator<T> {
578 execution: Option<VmExecution<T>>,
579}
580
581impl<T> VmGenerator<T>
582where
583 T: AsMut<Vm>,
584{
585 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 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
631pub struct VmStream<T> {
633 execution: Option<VmExecution<T>>,
634}
635
636impl<T> VmStream<T>
637where
638 T: AsMut<Vm>,
639{
640 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 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}