How does the compiler convert the `.await` to the invocation of `poll`?

struct S;
impl Future for S{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
       std::task::Poll::Pending
    }
}

async fn run(){
   let s = S;
   s.await;  // #1
}

IMO, The .await syntax is sugar that keeps calling the poll on S to change the machine state, which means poll is the underlying operation. However, the poll takes a receiver expression whose type should be Pin<& mut S> in this example. I wonder how does the compiler converts the .await expression to actually call the associated poll?

The generated impl Future code for async fn run() will do the equivalent of:

  1. Move s to a temporary variable (thus ensuring that no other code can be using s in any way at all).
  2. Call Pin::new_unchecked() on a mutable borrow of that temporary variable. (This is sound in the same way pin! is sound — nothing else can name that temporary, so nothing else can cause its value to be moved out, and variables on the stack are always dropped before they are deallocated.)
  3. Call poll() on that.

But there's not exactly any Rust code that corresponds to the contents of an async block — it's defined at a lower level than that, so it's not really any library function calls — just taking a pointer to the data, and specifying the type of that pointer as being Pin<&mut S>.

11 Likes

The key point to remember here is that async blocks and functions return an anonymous object that implements Future. This object can safely forward the context and pinning guarantees that it is given when its poll() method is called to any other Future::poll calls that it needs to defer to.

As a side note, this means that async blocks can only forward and rearrange the execution of futures. Somewhere down the line, there needs to be a manually-written Future that integrates with the I/O system and triggers the waker at an appropriate time.

If you wanted to write the transformation manually, the closest version would look something like this:

use std::mem::MaybeUninit;
use std::future::Future;
use std::task::{Context, Poll};
use std::pin::Pin;

pub fn run2()->impl Future<Output = ()> {
    RunFuture::call()
}

struct RunFuture {
    yield_point: u32,
    s: MaybeUninit<S>,
    _pin: std::marker::PhantomPinned
}

impl RunFuture {
    fn call()->Self {
        RunFuture { yield_point: 0, s: MaybeUninit::uninit(), _pin: std::marker::PhantomPinned }
    }
}

impl Future for RunFuture {
    type Output = ();
    
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>)->Poll<()> {
        // Safety: We never move `*unpinned` ourselves, and
        //         the unpinned reference doesn't escape code we control
        let unpinned = unsafe { self.get_unchecked_mut() };
        loop {
            match unpinned.yield_point {
                0 => {
                    unpinned.s.write(S);
                }
                1 => {
                    match unsafe {
                        Pin::new_unchecked(unpinned.s.assume_init_mut()).poll(cx)
                    } {
                        Poll::Pending => { return Poll::Pending }
                        Poll::Ready(_out) => {
                            unsafe { unpinned.s.assume_init_drop(); }
                            return Poll::Ready(());
                        }
                    };
                }
                _ => unreachable!()
            };
            unpinned.yield_point += 1;
        }
    }
}
8 Likes

Could I construe that every expr.await is conceptually equivalent to

let mut unique_temporary_name = expr;
let mut unique_temporary_name = unsafe {Pin::new_unchecked(& mut unique_temporary_name)};
match unique_temporary_name.as_mut().poll(cx){ 
   // ...
}

right?

Let’s try to gain some insights by looking at the generated HIR from your code example / run function. The playground does have an option to output HIR, though the formatting of the human-readable output lacks a bit. Let me try to fix the formatting, and it shows:

async fn run(){
   let s = S;
   s.await;  // #1
}

becoming

async fn run() -> /*impl Trait*/
   |mut _task_context: ResumeTy| {
      {
         let _t = {
            let s = S;
            match #[lang = "into_future"](s) {
               mut __awaitee => loop {
                  match unsafe {
                        #[lang = "poll"](#[lang = "new_unchecked"](&mut __awaitee),
                        #[lang = "get_context"](_task_context))
                     } {
                     #[lang = "Ready"] { 0: result } => break result,
                     #[lang = "Pending"] {} => { }
                  } // #1
                  _task_context = (yield ());
               },
            };
         };
         _t
      }
   }

Now we just have to map all those lang items back to their respective entries in the standard library…

async fn run() -> /*impl Trait*/
   |mut _task_context: ResumeTy| {
      {
         let _t = {
            let s = S;
            match IntoFuture::into_future(s) {
               mut __awaitee => loop {
                  match unsafe {
                        Future::poll(Pin::new_unchecked(&mut __awaitee),
                        std::future::get_context(_task_context))
                     } {
                     Poll::Ready(result) => break result,
                     Poll::Pending => { }
                  } // #1
                  _task_context = (yield ());
               },
            };
         };
         _t
      }
   }

Well, and we need to make this look a bit more like valid Rust syntax again…

fn run() -> impl Future<Output = ()> {
    |mut _task_context: ResumeTy| {
        {
            let _t = {
                let s = S;
                match IntoFuture::into_future(s) {
                    mut __awaitee => loop {
                        match unsafe {
                            Future::poll(
                                Pin::new_unchecked(&mut __awaitee),
                                std::future::get_context(_task_context),
                            )
                        } {
                            Poll::Ready(result) => break result,
                            Poll::Pending => {}
                        } // #1
                        _task_context = (yield ());
                    },
                };
            };
            _t
        }
    }
}

What’s with the usage of a closure? And oh look, there’s a “yield”? Can we make this compile somehow anyways? Well, we kind-of can! Using (unstable) coroutines (formerly “generators”), the feature that async internally is built on anyways:

fn run() -> impl Coroutine<ResumeTy, Yield = (), Return = ()> {
    |mut _task_context: ResumeTy| {
        {
            let _t = {
                let s = S;
                match IntoFuture::into_future(s) {
                    mut __awaitee => loop {
                        match unsafe {
                            Future::poll(
                                Pin::new_unchecked(&mut __awaitee),
                                std::future::get_context(_task_context),
                            )
                        } {
                            Poll::Ready(result) => break result,
                            Poll::Pending => {}
                        } // #1
                        _task_context = (yield ());
                    },
                };
            };
            _t
        }
    }
}

Compiles now: Rust Playground

Just… there must be some implicit compiler magic to turn such a Coroutine<ResumeTy, Yield = (), Return = Out> back into Future<Output = Out> for async fn or async blocks. No worries, we can do it ourselves! Well… almost. Looks like there’s no publicly accessible way to create this ResumeTy wrapper. We’ll ignore this and just transmute into it for the purposes of this demonstration.

fn run() -> impl Future<Output = ()> {
    CoFut(|mut _task_context: ResumeTy| {
        {
            let _t = {
                let s = S;
                match IntoFuture::into_future(s) {
                    mut __awaitee => loop {
                        match unsafe {
                            Future::poll(
                                Pin::new_unchecked(&mut __awaitee),
                                std::future::get_context(_task_context),
                            )
                        } {
                            Poll::Ready(result) => break result,
                            Poll::Pending => {}
                        } // #1
                        _task_context = (yield ());
                    },
                };
            };
            _t
        }
    })
}
struct CoFut<C>(C);

impl<C: Coroutine<ResumeTy, Yield = ()>> Future for CoFut<C> {
    type Output = C::Return;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        unsafe {
            let this: Pin<&mut C> = Pin::new_unchecked(&mut Pin::into_inner_unchecked(self).0);

            let ctx: NonNull<Context<'_>> = NonNull::from(ctx);

            match this.resume(mem::transmute(ctx)) {
                CoroutineState::Yielded(()) => Poll::Pending,
                CoroutineState::Complete(r) => Poll::Ready(r),
            }
        }
    }
}

Rust Playground

7 Likes

So with the above HIR desugaring in mind, your proposed “conceptually equivalent” code would need to be updated mainly by adding a loop, and then the match arms can also be filled out. So in the spirit of what you’ve started with, you’d change it to something like

let mut unique_temporary_name = expr;
let mut unique_temporary_name = unsafe {Pin::new_unchecked(& mut unique_temporary_name)};
loop {
    match unique_temporary_name.as_mut().poll(cx){ 
        Poll::Ready(()) => break,
        Poll::Pending => (),
    }
    cx = yield; // `yield` point, to wait for pending future
    // once `yield` returns, it give us the new up-to-date value
    // of `cx` which is updated to be used for the next `poll`
    // of this or another Future
}

The full desugaring additionally supports

  • better behavior around temporary lifetimes and destruction
    • the trick to use match EXPR { var => { … } } (instead of { let var = EXPR; … }) ensures that temporaries of EXPR stay alive throughout the code in “…”
    • and using match here, instead of moving it to a let in the surrounding block, also ensures the future is dropped at the end of the statements surrounding the original .await
  • return values (propagated through the break result to become the value of the whole loop and thus the outer match)
  • not just Future but IntoFuture types are supported

The question of how the whole async desugars further into a state machine then of course also needs to “desugar” the yields and the whole coroutine. The inner workings of what the result of such a transformation could be doing is what @2e71828 explored above.

4 Likes