@@ -8,14 +8,15 @@ extern crate std;
88use core:: sync:: atomic:: { AtomicBool , Ordering } ;
99use std:: boxed:: Box ;
1010use std:: collections:: HashMap ;
11+ use std:: ffi:: c_void;
1112use std:: fmt:: { self , Debug , Display } ;
1213use std:: future:: Future ;
1314use std:: mem;
1415use std:: pin:: Pin ;
1516use std:: ptr;
1617use std:: string:: String ;
1718use std:: sync:: Arc ;
18- use std:: task:: { Context , Poll , Wake , Waker } ;
19+ use std:: task:: { Context , Poll , Wake } ;
1920use std:: vec:: Vec ;
2021
2122use futures:: channel:: oneshot;
@@ -36,6 +37,7 @@ macro_rules! rtdebug {
3637}
3738
3839mod abi_buffer;
40+ mod cabi;
3941mod future_support;
4042mod stream_support;
4143mod waitable;
@@ -60,21 +62,32 @@ struct FutureState {
6062tasks : Option < FuturesUnordered < BoxFuture > > ,
6163 /// The waitable set containing waitables created by this task, if any.
6264waitable_set : Option < u32 > ,
63- /// A map of waitables to the corresponding waker and completion code.
64- ///
65- /// This is primarily filled in and managed by `WaitableOperation<S>`. The
66- /// waker here comes straight from `std::task::Context` and the pointer is
67- /// otherwise stored within the `WaitableOperation<S>` The raw pointer here
68- /// has a disconnected lifetime with each future but the management of the
69- /// internal states with respect to drop should always ensure that this is
70- /// only ever pointing to active waitable operations.
71- ///
72- /// When a waitable notification is received the corresponding entry in this
73- /// map is removed, the status code is filled in, and the waker is notified.
74- wakers : HashMap < u32 , ( Waker , * mut Option < u32 > ) > ,
65+
66+ /// State of all waitables in `waitable_set`, and the ptr/callback they're
67+ /// associated with.
68+ waitables : HashMap < u32 , ( * mut c_void , unsafe extern "C" fn ( * mut c_void , u32 ) ) > ,
69+
70+ /// Raw structure used to pass to `cabi::wasip3_task_set`
71+ wasip3_task : cabi:: wasip3_task ,
7572}
7673
7774impl FutureState {
75+ fn new ( future : BoxFuture ) -> FutureState {
76+ FutureState {
77+ todo : 0 ,
78+ tasks : Some ( [ future] . into_iter ( ) . collect ( ) ) ,
79+ waitable_set : None ,
80+ waitables : HashMap :: new ( ) ,
81+ wasip3_task : cabi:: wasip3_task {
82+ // This pointer is filled in before calling `wasip3_task_set`.
83+ ptr : ptr:: null_mut ( ) ,
84+ version : cabi:: WASIP3_TASK_V1 ,
85+ waitable_register,
86+ waitable_unregister,
87+ } ,
88+ }
89+ }
90+
7891 fn get_or_create_waitable_set ( & mut self ) -> u32 {
7992 * self . waitable_set . get_or_insert_with ( waitable_set_new)
8093 }
@@ -88,7 +101,32 @@ impl FutureState {
88101 }
89102
90103 fn remaining_work ( & self ) -> bool {
91- self . todo > 0 || !self . wakers . is_empty ( )
104+ self . todo > 0 || !self . waitables . is_empty ( )
105+ }
106+ }
107+
108+ unsafe extern "C" fn waitable_register (
109+ ptr : * mut c_void ,
110+ waitable : u32 ,
111+ callback : unsafe extern "C" fn ( * mut c_void , u32 ) ,
112+ callback_ptr : * mut c_void ,
113+ ) -> * mut c_void {
114+ let ptr = ptr. cast :: < FutureState > ( ) ;
115+ assert ! ( !ptr. is_null( ) ) ;
116+ ( * ptr) . add_waitable ( waitable) ;
117+ match ( * ptr) . waitables . insert ( waitable, ( callback_ptr, callback) ) {
118+ Some ( ( prev, _) ) => prev,
119+ None => ptr:: null_mut ( ) ,
120+ }
121+ }
122+
123+ unsafe extern "C" fn waitable_unregister ( ptr : * mut c_void , waitable : u32 ) -> * mut c_void {
124+ let ptr = ptr. cast :: < FutureState > ( ) ;
125+ assert ! ( !ptr. is_null( ) ) ;
126+ ( * ptr) . remove_waitable ( waitable) ;
127+ match ( * ptr) . waitables . remove ( & waitable) {
128+ Some ( ( prev, _) ) => prev,
129+ None => ptr:: null_mut ( ) ,
92130 }
93131}
94132
@@ -145,6 +183,22 @@ unsafe fn poll(state: *mut FutureState) -> Poll<()> {
145183 }
146184 }
147185
186+ // Finish our `wasip3_task` by initializing its self-referential pointer,
187+ // and then register it for the duration of this function with
188+ // `wasip3_task_set`. The previous value of `wasip3_task_set` will get
189+ // restored when this function returns.
190+ struct ResetTask ( * mut cabi:: wasip3_task ) ;
191+ impl Drop for ResetTask {
192+ fn drop ( & mut self ) {
193+ unsafe {
194+ cabi:: wasip3_task_set ( self . 0 ) ;
195+ }
196+ }
197+ }
198+ ( * state) . wasip3_task . ptr = state. cast ( ) ;
199+ let prev = cabi:: wasip3_task_set ( & mut ( * state) . wasip3_task ) ;
200+ let _reset = ResetTask ( prev) ;
201+
148202 loop {
149203 if let Some ( futures) = ( * state) . tasks . as_mut ( ) {
150204 let old = CURRENT ;
@@ -191,16 +245,9 @@ pub fn first_poll<T: 'static>(
191245 future : impl Future < Output = T > + ' static ,
192246 fun : impl FnOnce ( & T ) + ' static ,
193247) -> i32 {
194- let state = Box :: into_raw ( Box :: new ( FutureState {
195- todo : 0 ,
196- tasks : Some (
197- [ Box :: pin ( future. map ( |v| fun ( & v) ) ) as BoxFuture ]
198- . into_iter ( )
199- . collect ( ) ,
200- ) ,
201- waitable_set : None ,
202- wakers : HashMap :: new ( ) ,
203- } ) ) ;
248+ let state = Box :: into_raw ( Box :: new ( FutureState :: new ( Box :: pin (
249+ future. map ( |v| fun ( & v) ) ,
250+ ) ) ) ) ;
204251 let done = unsafe { poll ( state) . is_ready ( ) } ;
205252 unsafe { callback_code ( state, done) }
206253}
@@ -339,9 +386,8 @@ unsafe fn callback_with_state(
339386 "EVENT_{{STREAM,FUTURE}}_{{READ,WRITE}}({event0:#x}, {event1:#x}, {event2:#x})"
340387 ) ;
341388 ( * state) . remove_waitable ( event1 as u32 ) ;
342- let ( waker, code) = ( * state) . wakers . remove ( & ( event1 as u32 ) ) . unwrap ( ) ;
343- * code = Some ( event2 as u32 ) ;
344- waker. wake ( ) ;
389+ let ( ptr, callback) = ( * state) . waitables . remove ( & ( event1 as u32 ) ) . unwrap ( ) ;
390+ callback ( ptr, event2 as u32 ) ;
345391
346392 let done = poll ( state) . is_ready ( ) ;
347393 callback_code ( state, done)
@@ -469,16 +515,7 @@ pub fn spawn(future: impl Future<Output = ()> + 'static) {
469515// TODO: refactor so `'static` bounds aren't necessary
470516pub fn block_on < T : ' static > ( future : impl Future < Output = T > + ' static ) -> T {
471517 let ( tx, mut rx) = oneshot:: channel ( ) ;
472- let state = & mut FutureState {
473- todo : 0 ,
474- tasks : Some (
475- [ Box :: pin ( future. map ( move |v| drop ( tx. send ( v) ) ) ) as BoxFuture ]
476- . into_iter ( )
477- . collect ( ) ,
478- ) ,
479- waitable_set : None ,
480- wakers : HashMap :: new ( ) ,
481- } ;
518+ let state = & mut FutureState :: new ( Box :: pin ( future. map ( move |v| drop ( tx. send ( v) ) ) ) as BoxFuture ) ;
482519 loop {
483520 match unsafe { poll ( state) } {
484521 Poll :: Ready ( ( ) ) => break rx. try_recv ( ) . unwrap ( ) . unwrap ( ) ,
0 commit comments