How to compose two async function in Rust - asynchronous

Question

I am trying to write a higher order function that compose two async function.
i am basically looking for the async version this
fn compose<A, B, C, G, F>(f: F, g: G) -> impl Fn(A) -> C
where
F: Fn(A) -> B,
G: Fn(B) -> C,
{
move |x| g(f(x))
}
This my attempt so far.
fn compose_future<A, B, C, G, F>(f: F, g: G) -> (impl Fn(A) -> impl Future<C>)
where
F: Fn(A) -> impl Future<B>,
G: Fn(B) -> impl Future<C>,
{
move |x| async { g(f(x).await).await }
}
and i get the following error
error[E0562]: `impl Trait` not allowed outside of function and inherent method return types
--> src\channel.rs:13:17
|
13 | F: Fn(A) -> impl Future<B>,
| ^^^^^^^^^^^^^^
Is it possible to accomplish this ?

Answer 1

I'm not sure it is possible to do it that simple with impl Trait-s. The one solution I can come up with is old-fashioned future types usage without async-await feature. TLDR: full playground. Async-await uses a generators which internally holds a state machine, so we need to define it manually:
enum State<In, F, FutOutF, G, FutOutG> {
Initial(In, F, G), // Out composed type created
FirstAwait(FutOutF, G), // Composed type waits for the first future
SecondAwait(FutOutG), // and for the second
// here can be a `Completed` state, but it simpler
// to handle it with `Option<..>` in our future itself
}
Then define a composed type itself:
struct Compose<In, Out, F, FutOutF, G, FutOutG> {
state: Option<State<In, F, FutOutF, G, FutOutG>>,
_t: PhantomData<Out>,
}
// And "entry-point" would be something like that:
fn compose_fut<In, Out, F, FutOutF, G, FutOutG>(
i: In,
f: F,
g: G,
) -> Compose<In, Out, F, FutOutF, G, FutOutG> {
Compose {
state: Some(State::Initial(i, f, g)),
_t: PhantomData,
}
}
Then comes the most complex part - impl Future itself, here a base impl declaration without implementation:
impl<In, Mid, Out, F, FutOutF, G, FutOutG> Future for Compose<In, Out, F, FutOutF, G, FutOutG>
where
FutOutF: Future<Output = Mid>,
F: FnOnce(In) -> FutOutF,
FutOutG: Future<Output = Out>,
G: FnOnce(Mid) -> FutOutG,
{
type Output = Out;
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
// here comes the magic
}
}
Values transformed as following: In -> Mid -> Out, where F and G are our composed functions and their output are FutOutF and FutOutG accordingly. And finally Future::poll implementation:
let this = unsafe { self.get_unchecked_mut() };
let state = this.state.take();
match state {
None => Poll::Pending, // invalid state
Some(State::Initial(i, f, g)) => {
let fut = f(i);
this.state = Some(State::FirstAwait(fut, g));
cx.waker().wake_by_ref();
Poll::Pending
}
Some(State::FirstAwait(mut fut, g)) => {
let val = match unsafe { Pin::new_unchecked(&mut fut) }.poll(cx) {
Poll::Ready(v) => v,
Poll::Pending => {
this.state = Some(State::FirstAwait(fut, g));
return Poll::Pending;
}
};
let fut = g(val);
this.state = Some(State::SecondAwait(fut));
cx.waker().wake_by_ref();
Poll::Pending
}
Some(State::SecondAwait(mut fut)) => {
match unsafe { Pin::new_unchecked(&mut fut) }.poll(cx) {
Poll::Ready(v) => Poll::Ready(v),
Poll::Pending => {
this.state = Some(State::SecondAwait(fut));
Poll::Pending
}
}
}
}
I avoid any library to make it "plain", usually unsafe parts are handled with pin-project or futures::pin_mut. The state management is fairly complex, so I suggest to re-check the implementation, there might be mistakes.