use {
crate::rt::async_support::waitable::{WaitableOp, WaitableOperation},
crate::rt::async_support::ReturnCode,
crate::rt::Cleanup,
std::{
alloc::Layout,
fmt,
future::{Future, IntoFuture},
marker,
pin::Pin,
ptr,
sync::atomic::{AtomicU32, Ordering::Relaxed},
task::{Context, Poll},
},
};
#[doc(hidden)]
pub struct FutureVtable<T> {
pub layout: Layout,
pub lower: unsafe fn(value: T, dst: *mut u8),
pub dealloc_lists: unsafe fn(dst: *mut u8),
pub lift: unsafe fn(dst: *mut u8) -> T,
pub start_write: unsafe extern "C" fn(future: u32, val: *const u8) -> u32,
pub start_read: unsafe extern "C" fn(future: u32, val: *mut u8) -> u32,
pub cancel_write: unsafe extern "C" fn(future: u32) -> u32,
pub cancel_read: unsafe extern "C" fn(future: u32) -> u32,
pub drop_writable: unsafe extern "C" fn(future: u32),
pub drop_readable: unsafe extern "C" fn(future: u32),
pub new: unsafe extern "C" fn() -> u64,
}
pub unsafe fn future_new<T>(
default: fn() -> T,
vtable: &'static FutureVtable<T>,
) -> (FutureWriter<T>, FutureReader<T>) {
unsafe {
let handles = (vtable.new)();
let reader = handles as u32;
let writer = (handles >> 32) as u32;
rtdebug!("future.new() = [{writer}, {reader}]");
(
FutureWriter::new(writer, default, vtable),
FutureReader::new(reader, vtable),
)
}
}
pub struct FutureWriter<T: 'static> {
handle: u32,
vtable: &'static FutureVtable<T>,
should_write_default_value: bool,
default: fn() -> T,
}
impl<T> FutureWriter<T> {
#[doc(hidden)]
pub unsafe fn new(handle: u32, default: fn() -> T, vtable: &'static FutureVtable<T>) -> Self {
Self {
handle,
default,
should_write_default_value: true,
vtable,
}
}
pub fn write(self, value: T) -> FutureWrite<T> {
FutureWrite {
op: WaitableOperation::new((self, value)),
}
}
}
impl<T> fmt::Debug for FutureWriter<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("FutureWriter")
.field("handle", &self.handle)
.finish()
}
}
impl<T> Drop for FutureWriter<T> {
fn drop(&mut self) {
if self.should_write_default_value {
let clone = FutureWriter {
handle: self.handle,
default: self.default,
should_write_default_value: false,
vtable: self.vtable,
};
crate::rt::async_support::spawn(async move {
let value = (clone.default)();
let _ = clone.write(value).await;
});
} else {
unsafe {
rtdebug!("future.drop-writable({})", self.handle);
(self.vtable.drop_writable)(self.handle);
}
}
}
}
pub struct FutureWrite<T: 'static> {
op: WaitableOperation<FutureWriteOp<T>>,
}
struct FutureWriteOp<T>(marker::PhantomData<T>);
enum WriteComplete<T> {
Written,
Dropped(T),
Cancelled(T),
}
unsafe impl<T> WaitableOp for FutureWriteOp<T>
where
T: 'static,
{
type Start = (FutureWriter<T>, T);
type InProgress = (FutureWriter<T>, Option<Cleanup>);
type Result = (WriteComplete<T>, FutureWriter<T>);
type Cancel = FutureWriteCancel<T>;
fn start((writer, value): Self::Start) -> (u32, Self::InProgress) {
let (ptr, cleanup) = Cleanup::new(writer.vtable.layout);
let code = unsafe {
(writer.vtable.lower)(value, ptr);
(writer.vtable.start_write)(writer.handle, ptr)
};
rtdebug!("future.write({}, {ptr:?}) = {code:#x}", writer.handle);
(code, (writer, cleanup))
}
fn start_cancelled((writer, value): Self::Start) -> Self::Cancel {
FutureWriteCancel::Cancelled(value, writer)
}
fn in_progress_update(
(mut writer, cleanup): Self::InProgress,
code: u32,
) -> Result<Self::Result, Self::InProgress> {
let ptr = cleanup
.as_ref()
.map(|c| c.ptr.as_ptr())
.unwrap_or(ptr::null_mut());
match code {
super::BLOCKED => Err((writer, cleanup)),
super::DROPPED | super::CANCELLED => {
let value = unsafe { (writer.vtable.lift)(ptr) };
let status = if code == super::DROPPED {
writer.should_write_default_value = false;
WriteComplete::Dropped(value)
} else {
WriteComplete::Cancelled(value)
};
Ok((status, writer))
}
super::COMPLETED => {
writer.should_write_default_value = false;
unsafe {
(writer.vtable.dealloc_lists)(ptr);
}
Ok((WriteComplete::Written, writer))
}
other => unreachable!("unexpected code {other:?}"),
}
}
fn in_progress_waitable((writer, _): &Self::InProgress) -> u32 {
writer.handle
}
fn in_progress_cancel((writer, _): &Self::InProgress) -> u32 {
let code = unsafe { (writer.vtable.cancel_write)(writer.handle) };
rtdebug!("future.cancel-write({}) = {code:#x}", writer.handle);
code
}
fn result_into_cancel((result, writer): Self::Result) -> Self::Cancel {
match result {
WriteComplete::Written => FutureWriteCancel::AlreadySent,
WriteComplete::Dropped(val) => FutureWriteCancel::Dropped(val),
WriteComplete::Cancelled(val) => FutureWriteCancel::Cancelled(val, writer),
}
}
}
impl<T: 'static> Future for FutureWrite<T> {
type Output = Result<(), FutureWriteError<T>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.pin_project()
.poll_complete(cx)
.map(|(result, _writer)| match result {
WriteComplete::Written => Ok(()),
WriteComplete::Dropped(value) | WriteComplete::Cancelled(value) => {
Err(FutureWriteError { value })
}
})
}
}
impl<T: 'static> FutureWrite<T> {
fn pin_project(self: Pin<&mut Self>) -> Pin<&mut WaitableOperation<FutureWriteOp<T>>> {
unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().op) }
}
pub fn cancel(self: Pin<&mut Self>) -> FutureWriteCancel<T> {
self.pin_project().cancel()
}
}
pub struct FutureWriteError<T> {
pub value: T,
}
impl<T> fmt::Debug for FutureWriteError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("FutureWriteError").finish_non_exhaustive()
}
}
impl<T> fmt::Display for FutureWriteError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
"read end dropped".fmt(f)
}
}
impl<T> std::error::Error for FutureWriteError<T> {}
#[derive(Debug)]
pub enum FutureWriteCancel<T: 'static> {
AlreadySent,
Dropped(T),
Cancelled(T, FutureWriter<T>),
}
pub struct FutureReader<T: 'static> {
handle: AtomicU32,
vtable: &'static FutureVtable<T>,
}
impl<T> fmt::Debug for FutureReader<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("FutureReader")
.field("handle", &self.handle)
.finish()
}
}
impl<T> FutureReader<T> {
#[doc(hidden)]
pub fn new(handle: u32, vtable: &'static FutureVtable<T>) -> Self {
Self {
handle: AtomicU32::new(handle),
vtable,
}
}
#[doc(hidden)]
pub fn take_handle(&self) -> u32 {
let ret = self.opt_handle().unwrap();
self.handle.store(u32::MAX, Relaxed);
ret
}
fn handle(&self) -> u32 {
self.opt_handle().unwrap()
}
fn opt_handle(&self) -> Option<u32> {
match self.handle.load(Relaxed) {
u32::MAX => None,
other => Some(other),
}
}
}
impl<T> IntoFuture for FutureReader<T> {
type Output = T;
type IntoFuture = FutureRead<T>;
fn into_future(self) -> Self::IntoFuture {
FutureRead {
op: WaitableOperation::new(self),
}
}
}
impl<T> Drop for FutureReader<T> {
fn drop(&mut self) {
let Some(handle) = self.opt_handle() else {
return;
};
unsafe {
rtdebug!("future.drop-readable({handle})");
(self.vtable.drop_readable)(handle);
}
}
}
pub struct FutureRead<T: 'static> {
op: WaitableOperation<FutureReadOp<T>>,
}
struct FutureReadOp<T>(marker::PhantomData<T>);
enum ReadComplete<T> {
Value(T),
Cancelled,
}
unsafe impl<T> WaitableOp for FutureReadOp<T>
where
T: 'static,
{
type Start = FutureReader<T>;
type InProgress = (FutureReader<T>, Option<Cleanup>);
type Result = (ReadComplete<T>, FutureReader<T>);
type Cancel = Result<T, FutureReader<T>>;
fn start(reader: Self::Start) -> (u32, Self::InProgress) {
let (ptr, cleanup) = Cleanup::new(reader.vtable.layout);
let code = unsafe { (reader.vtable.start_read)(reader.handle(), ptr) };
rtdebug!("future.read({}, {ptr:?}) = {code:#x}", reader.handle());
(code, (reader, cleanup))
}
fn start_cancelled(state: Self::Start) -> Self::Cancel {
Err(state)
}
fn in_progress_update(
(reader, cleanup): Self::InProgress,
code: u32,
) -> Result<Self::Result, Self::InProgress> {
match ReturnCode::decode(code) {
ReturnCode::Blocked => Err((reader, cleanup)),
ReturnCode::Cancelled(0) => Ok((ReadComplete::Cancelled, reader)),
ReturnCode::Completed(0) => {
let ptr = cleanup
.as_ref()
.map(|c| c.ptr.as_ptr())
.unwrap_or(ptr::null_mut());
let value = unsafe { (reader.vtable.lift)(ptr) };
Ok((ReadComplete::Value(value), reader))
}
other => panic!("unexpected code {other:?}"),
}
}
fn in_progress_waitable((reader, _): &Self::InProgress) -> u32 {
reader.handle()
}
fn in_progress_cancel((reader, _): &Self::InProgress) -> u32 {
let code = unsafe { (reader.vtable.cancel_read)(reader.handle()) };
rtdebug!("future.cancel-read({}) = {code:#x}", reader.handle());
code
}
fn result_into_cancel((value, reader): Self::Result) -> Self::Cancel {
match value {
ReadComplete::Value(value) => Ok(value),
ReadComplete::Cancelled => Err(reader),
}
}
}
impl<T: 'static> Future for FutureRead<T> {
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.pin_project()
.poll_complete(cx)
.map(|(result, _reader)| match result {
ReadComplete::Value(val) => val,
ReadComplete::Cancelled => panic!("cannot poll after cancelling"),
})
}
}
impl<T> FutureRead<T> {
fn pin_project(self: Pin<&mut Self>) -> Pin<&mut WaitableOperation<FutureReadOp<T>>> {
unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().op) }
}
pub fn cancel(self: Pin<&mut Self>) -> Result<T, FutureReader<T>> {
self.pin_project().cancel()
}
}