diff --git a/hercules_cg/src/rt.rs b/hercules_cg/src/rt.rs index beb83f51691a91ac8e9c377709f55de8f7be117d..d3013239f5f78ce8e63181c181c0d5a8cbf77f81 100644 --- a/hercules_cg/src/rt.rs +++ b/hercules_cg/src/rt.rs @@ -594,7 +594,7 @@ impl<'a> RTContext<'a> { (_, true) => { write!(block, "{}", self.clone_arc(id).unwrap())?; format!( - "*async_call_{}.lock().await = Some(::async_std::task::spawn(async move {{ ", + "*async_call_{}.lock().await = ::hercules_rt::__FutureSlotWrapper::new(::async_std::task::spawn(async move {{ ", id.idx(), ) } @@ -1106,7 +1106,7 @@ impl<'a> RTContext<'a> { if is_async_call { write!( w, - "let mut async_call_{} = ::std::sync::Arc::new(::async_std::sync::Mutex::new(None));", + "let mut async_call_{} = ::std::sync::Arc::new(::async_std::sync::Mutex::new(::hercules_rt::__FutureSlotWrapper::empty()));", idx, )?; } else { @@ -1382,10 +1382,7 @@ impl<'a> RTContext<'a> { && func.schedules[id.idx()].contains(&Schedule::AsyncCall) { assert!(!lhs); - format!( - "async_call_{}.lock().await.as_mut().unwrap().await", - id.idx(), - ) + format!("async_call_{}.lock().await.inspect().await", id.idx(),) } else { format!("node_{}", id.idx()) } diff --git a/hercules_rt/src/lib.rs b/hercules_rt/src/lib.rs index 090a38a02cbbcd46253452f76a0b71681363c833..d19a0a5a16e4438e8746ca312428759ea28ed556 100644 --- a/hercules_rt/src/lib.rs +++ b/hercules_rt/src/lib.rs @@ -1,10 +1,10 @@ #![feature(once_cell_try)] use std::alloc::{alloc, dealloc, GlobalAlloc, Layout, System}; +use std::future::Future; use std::marker::PhantomData; use std::ptr::{copy_nonoverlapping, write_bytes, NonNull}; use std::slice::{from_raw_parts, from_raw_parts_mut}; - use std::sync::OnceLock; /* @@ -426,6 +426,46 @@ impl __RawPtrSendSync { unsafe impl Send for __RawPtrSendSync {} unsafe impl Sync for __RawPtrSendSync {} +#[derive(Clone, Debug)] +pub struct __FutureSlotWrapper<T, F> +where + T: Copy, + F: Future<Output = T>, +{ + future: Option<F>, + slot: Option<T>, +} + +impl<T, F> __FutureSlotWrapper<T, F> +where + T: Copy, + F: Future<Output = T>, +{ + pub fn empty() -> Self { + Self { + future: None, + slot: None, + } + } + + pub fn new(f: F) -> Self { + Self { + future: Some(f), + slot: None, + } + } + + pub async fn inspect(&mut self) -> T { + if let Some(slot) = self.slot { + slot + } else { + let result = self.future.take().unwrap().await; + self.slot = Some(result); + result + } + } +} + /* * A HerculesBox holds memory that can be on any device and provides a common interface to moving * data where it is needed.