rune/runtime/
future.rs

1use core::fmt;
2use core::future;
3use core::pin::Pin;
4use core::ptr::NonNull;
5use core::task::{Context, Poll};
6
7use crate::alloc::alloc::Global;
8use crate::alloc::{self, Box};
9use crate::runtime::{ToValue, Value, VmError, VmErrorKind};
10use crate::Any;
11
12use pin_project::pin_project;
13
14/// A virtual table for a type-erased future.
15struct Vtable {
16    poll: unsafe fn(*mut (), cx: &mut Context<'_>) -> Poll<Result<Value, VmError>>,
17    drop: unsafe fn(*mut ()),
18}
19
20/// A type-erased future that can only be unsafely polled in combination with
21/// the virtual machine that created it.
22#[derive(Any)]
23#[rune(crate)]
24#[rune(item = ::std::future)]
25pub struct Future {
26    future: Option<NonNull<()>>,
27    vtable: &'static Vtable,
28}
29
30impl Future {
31    /// Construct a new wrapped future.
32    pub(crate) fn new<T, O>(future: T) -> alloc::Result<Self>
33    where
34        T: 'static + future::Future<Output = Result<O, VmError>>,
35        O: ToValue,
36    {
37        let (future, Global) = Box::into_raw_with_allocator(Box::try_new(future)?);
38
39        let future = unsafe { NonNull::new_unchecked(future).cast() };
40
41        Ok(Self {
42            future: Some(future),
43            vtable: &Vtable {
44                poll: |future, cx| unsafe {
45                    match Pin::new_unchecked(&mut *future.cast::<T>()).poll(cx) {
46                        Poll::Pending => Poll::Pending,
47                        Poll::Ready(result) => match result {
48                            Ok(result) => match result.to_value() {
49                                Ok(value) => Poll::Ready(Ok(value)),
50                                Err(err) => Poll::Ready(Err(err.into())),
51                            },
52                            Err(err) => Poll::Ready(Err(err)),
53                        },
54                    }
55                },
56                drop: |future| unsafe {
57                    _ = Box::from_raw_in(future.cast::<T>(), Global);
58                },
59            },
60        })
61    }
62
63    /// Check if future is completed.
64    ///
65    /// This will prevent it from being used in a select expression.
66    pub fn is_completed(&self) -> bool {
67        self.future.is_none()
68    }
69}
70
71impl future::Future for Future {
72    type Output = Result<Value, VmError>;
73
74    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<Value, VmError>> {
75        unsafe {
76            let this = self.get_unchecked_mut();
77
78            let Some(future) = this.future else {
79                return Poll::Ready(Err(VmError::new(VmErrorKind::FutureCompleted)));
80            };
81
82            match (this.vtable.poll)(future.as_ptr(), cx) {
83                Poll::Ready(result) => {
84                    this.future = None;
85                    (this.vtable.drop)(future.as_ptr());
86                    Poll::Ready(result)
87                }
88                Poll::Pending => Poll::Pending,
89            }
90        }
91    }
92}
93
94impl Drop for Future {
95    fn drop(&mut self) {
96        unsafe {
97            if let Some(future) = self.future.take() {
98                (self.vtable.drop)(future.as_ptr());
99            }
100        }
101    }
102}
103
104impl fmt::Debug for Future {
105    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
106        fmt.debug_struct("Future")
107            .field("is_completed", &self.future.is_none())
108            .finish_non_exhaustive()
109    }
110}
111
112/// Future wrapper used to keep track of associated data.
113#[pin_project]
114pub struct SelectFuture<T, F> {
115    data: T,
116    #[pin]
117    future: F,
118}
119
120impl<T, F> SelectFuture<T, F> {
121    /// Construct a new select future.
122    pub fn new(data: T, future: F) -> Self {
123        Self { data, future }
124    }
125}
126
127impl<T, F> future::Future for SelectFuture<T, F>
128where
129    T: Copy,
130    F: future::Future<Output = Result<Value, VmError>>,
131{
132    type Output = Result<(T, Value), VmError>;
133
134    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
135        let this = self.project();
136        let result = this.future.poll(cx);
137
138        match result {
139            Poll::Ready(result) => match result {
140                Ok(value) => Poll::Ready(Ok((*this.data, value))),
141                Err(error) => Poll::Ready(Err(error)),
142            },
143            Poll::Pending => Poll::Pending,
144        }
145    }
146}
rune/runtime/future.rs

rune/runtime/
future.rs
