rune/runtime/
future.rs

Help
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, VmErrorKind, VmResult};
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<VmResult<Value>>,
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 = VmResult<O>>,
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(VmResult::Ok(result)) => match result.to_value() {
48                            Ok(value) => Poll::Ready(VmResult::Ok(value)),
49                            Err(err) => Poll::Ready(VmResult::Err(err.into())),
50                        },
51                        Poll::Ready(VmResult::Err(err)) => Poll::Ready(VmResult::Err(err)),
52                    }
53                },
54                drop: |future| unsafe {
55                    _ = Box::from_raw_in(future.cast::<T>(), Global);
56                },
57            },
58        })
59    }
60
61    /// Check if future is completed.
62    ///
63    /// This will prevent it from being used in a select expression.
64    pub fn is_completed(&self) -> bool {
65        self.future.is_none()
66    }
67}
68
69impl future::Future for Future {
70    type Output = VmResult<Value>;
71
72    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<VmResult<Value>> {
73        unsafe {
74            let this = self.get_unchecked_mut();
75
76            let Some(future) = this.future else {
77                return Poll::Ready(VmResult::err(VmErrorKind::FutureCompleted));
78            };
79
80            match (this.vtable.poll)(future.as_ptr(), cx) {
81                Poll::Ready(result) => {
82                    this.future = None;
83                    (this.vtable.drop)(future.as_ptr());
84                    Poll::Ready(result)
85                }
86                Poll::Pending => Poll::Pending,
87            }
88        }
89    }
90}
91
92impl Drop for Future {
93    fn drop(&mut self) {
94        unsafe {
95            if let Some(future) = self.future.take() {
96                (self.vtable.drop)(future.as_ptr());
97            }
98        }
99    }
100}
101
102impl fmt::Debug for Future {
103    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
104        fmt.debug_struct("Future")
105            .field("is_completed", &self.future.is_none())
106            .finish_non_exhaustive()
107    }
108}
109
110/// Future wrapper used to keep track of associated data.
111#[pin_project]
112pub struct SelectFuture<T, F> {
113    data: T,
114    #[pin]
115    future: F,
116}
117
118impl<T, F> SelectFuture<T, F> {
119    /// Construct a new select future.
120    pub fn new(data: T, future: F) -> Self {
121        Self { data, future }
122    }
123}
124
125impl<T, F> future::Future for SelectFuture<T, F>
126where
127    T: Copy,
128    F: future::Future<Output = VmResult<Value>>,
129{
130    type Output = VmResult<(T, Value)>;
131
132    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
133        let this = self.project();
134        let result = this.future.poll(cx);
135
136        match result {
137            Poll::Ready(result) => match result {
138                VmResult::Ok(value) => Poll::Ready(VmResult::Ok((*this.data, value))),
139                VmResult::Err(error) => Poll::Ready(VmResult::Err(error)),
140            },
141            Poll::Pending => Poll::Pending,
142        }
143    }
144}
rune/runtime/future.rs

rune/runtime/
future.rs
