futures-rs icon indicating copy to clipboard operation
futures-rs copied to clipboard

Idea: `poll_map_fn` as an analogous to `poll_fn`, but pass through the lifetime of the poll function

Open yshui opened this issue 3 years ago • 1 comments

Motivation

poll_fn is insufficient in cases where the return of the poll borrows from its arguments.

Say (the usual Context parameter omitted for brevity):

struct AsyncBufReader;

impl AsyncBufReader {
    fn poll_fill_buf<'a>(self: Pin<&'a mut Self>) -> Poll<&'a [u8]> {
        todo!()
    }
}

Say the poll_fill_buf function borrows from a internal buffer of AsyncBufReader.

Now let's try to use poll_fill_buf with poll_fn:

let reader = AsyncBufReader;
let buf = block_on(poll_fn(|_| {
    Pin::new(&mut reader).poll_fill_buf()
}));

we got:

error: captured variable cannot escape `FnMut` closure body
  --> src/main.rs:17:9
   |
15 |     let mut reader = AsyncBufReader;
   |         ---------- variable defined here
16 |     let buf = block_on(poll_fn(|_| {
   |                                  - inferred to be a `FnMut` closure
17 |         Pin::new(&mut reader).poll_fill_buf()
   |         ^^^^^^^^^^^^^^------^^^^^^^^^^^^^^^^^
   |         |             |
   |         |             variable captured here
   |         returns a reference to a captured variable which escapes the closure body
   |
   = note: `FnMut` closures only have access to their captured variables while they are executing...
   = note: ...therefore, they cannot allow references to captured variables to escape

reader is moved into the closure and we can't return a borrow of it.

Let's try again:

let mut reader = AsyncBufReader;
let reader = &mut reader;
let buf = block_on(poll_fn(|_| {
    Pin::new(reader).poll_fill_buf()
}));

Now we get:

error[[E0507]](https://doc.rust-lang.org/stable/error-index.html#E0507): cannot move out of `reader`, a captured variable in an `FnMut` closure
  --> src/main.rs:18:18
   |
16 |       let reader = &mut reader;
   |           ------ captured outer variable
17 |       let buf = block_on(poll_fn(|_| {
   |  ________________________________-
18 | |         Pin::new(reader).poll_fill_buf()
   | |                  ^^^^^^ move occurs because `reader` has type `&mut AsyncBufReader`, which does not implement the `Copy` trait
19 | |     }));
   | |_____- captured by this `FnMut` closure

Because poll_fill_buf takes Pin<&mut Self>, the reference we created got moved into the function, which is not allowed from FnMut

Obviously we also can't move reader = AsyncBufRead into the closure. So we are at a dead end.

Proposal

pub fn poll_map_fn<'a, F, T, O>(polled: Pin<&'a mut T>, f: F) -> PollMapFn<'a, F, T, O>
where
    F: Fn(Pin<&'a mut T>, &mut ::std::task::Context<'_>) -> Poll<O> + Unpin;

This function takes an object and a poll function to be called on this object. Instead of requiring the function to capture the object to be polled, the polled object is passed into the function every time it's called. The function is allowed to return a type O which borrows from T.

The above example using poll_map_fn would look like:

let mut reader = AsyncBufReader;
pin_mut!(reader);
let buf = block_on(poll_map_fn(reader, |reader, _| {
    reader.poll_fill_buf()
}));

Soundness?

This function clearly can't be implemented in safe Rust. PollMapFn::poll have to repeatedly produce a Pin<&mut T> despite it being consumed by F each time it's called.

I think this is safe, because if F returns Pending, then the Pin<&mut T> has to have been dropped (F is a Fn so it's impossible for it to tuck the reference away somewhere), so it's fine for PollMapFn to reproduce it somehow.

But maybe there is some obvious unsoundness I am missing here.

Alternatives

Modify std::task::Poll so Pending can give back the Pin<&mut Self>?

yshui avatar Sep 10 '22 19:09 yshui

I realised I need the Fn bound to be:

F: Fn(Pin<&mut T>, &mut ::std::task::Context<'_>) -> Poll<O> + Unpin;

instead. To make sure F can't store the &mut T somewhere. But then there is no easy way to express that O borrows from &mut T. Discussion: https://rust-lang.zulipchat.com/#narrow/stream/122651-general/topic/HRTB.20and.20borrowing.20closure/near/298601730

yshui avatar Sep 13 '22 15:09 yshui