How to have a future share a variable with an inner future

help
#1

I want to have two nested futures share access to a variable. This is fairly easy to do with regular functions (just pass a borrow), and is also possible with an explicit state machine (access the variable inside the inner struct), but when the inner future is an async fn (which I would really like to do because of the ergonomics), it seems impossible to pass variables to it on poll, and the compiler wants to keep layout details of the state machine struct to itself so it isn't easy to just access the value in the struct. Here are three attempts to accomplish this, with varying levels of failure:

#[macro_use] extern crate futures;

use std::{future::Future, pin::Pin, task::{Context, Poll}};
async fn f() {}

/// Use two `async fn`'s. This fails because `fut` has no members that
/// `outer` can access. We can't pass `&mut foo` to `inner` either because
/// it will take ownership of the mut borrow making `foo` inaccessible as long as
/// `fut` exists.
fn attempt_1() -> impl Future<Output=u32> {
    async fn inner() -> u32 {
        let mut foo = 0;
        f().await;
        foo = 1;
        f().await;
        foo = 2;
        f().await;
        foo
    }

    async fn outer() -> u32 {
        let fut = inner();
        fut.await;
        fut.foo = 3; // fail
        fut.await;
        fut.foo
    }

    outer()
}

/// Use a hand rolled struct for the outer future. We have better control
/// over the layout this way, but we still can't pass `foo` to the inner future
/// because it would cause `Outer` to be self referential (which should be fine
/// but I haven't managed to grok the pin API).
fn attempt_2() -> impl Future<Output=u32> {

    async fn inner(foo: &mut u32) -> u32 {
        f().await;
        *foo = 1;
        f().await;
        *foo = 2;
        f().await;
        *foo
    }

    struct Outer<F>(u32, F);

    impl<F: Future<Output=u32>> Future for Outer<F> {
        type Output = u32;
        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
            let Outer(foo, fut) = unsafe { self.get_unchecked_mut() };
            *foo += 1;
            ready!(unsafe { Pin::new_unchecked(fut) }.poll(cx, /* foo */ ));
            Poll::Ready(*foo)
        }
    }

    let mut foo = 0;
    Outer(foo, inner(&mut foo)) // fail, we want the &mut foo to be a reference to the first foo
}

/// This one "works", by resorting to pointers and with an extra heap allocation.
fn attempt_3() -> impl Future<Output=u32> {

    async fn inner(foo: *mut u32) -> u32 {
        f().await;
        unsafe { *foo = 1; }
        f().await;
        unsafe { *foo = 2; }
        f().await;
        unsafe { *foo }
    }

    struct Outer<F>(Pin<Box<u32>>, F);

    impl<F: Future<Output=u32>> Future for Outer<F> {
        type Output = u32;
        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
            let Outer(foo, fut) = unsafe { self.get_unchecked_mut() };
            **foo += 1;
            ready!(unsafe { Pin::new_unchecked(fut) }.poll(cx, /* foo */ ));
            Poll::Ready(**foo)
        }
    }

    let mut foo = Box::pin(0);
    let ptr: *mut _ = &mut *foo;
    Outer(foo, inner(ptr)) // ok, but lotsa unsafe
}
#2

One option is to pass the reference on every call to poll, using a custom poll method rather than an impl of Future.

#3

But an async fn is going to implement Future and nothing more, right? The goal here is to avoid having to manually write the state machine for inner, which is big and complicated in my real application. Outer is just a wrapper so I can do stuff before/after poll. If inner implements a custom poll method, then it must be a hand rolled struct, unless there is a way to use generics to wrap it as with Outer, but in this case we are back to the same issue.

#4

You can greatly reduce the unsafe of that last example by deciding that what you are trying to do leads to non-trivial exclusive accesses (the whole concept loses a lot of its meaning when Pinning is involved), so you could just embrace that shared mutability aspect:

async fn attempt_3_safe () -> u32
{
    use ::core::cell::Cell as Mut;

    async fn inner(foo: &Mut <u32>) -> u32
    {
        f().await;
        foo.set(1);
        f().await;
        foo.set(2);
        f().await;
        foo.get()
    }
    
    let foo = &Mut::new(0);
    let fut = inner(foo);
    ::futures::pin_mut!(fut);
    ::futures::future::poll_fn(move |cx| Poll::Ready({
        foo.set(foo.get() + 1);
        ::futures::ready!(
            fut.as_mut().poll(cx)
        )
    })).await
}

The other alternative would be to hand-roll a Generator that takes a &mut Context<'_>, &mut u32 each time it is polled for the inner one, and then hand-roll a Generator-to-Future conversion on the outer one.