How do I implement `.or()` for `Future<Option>`?

Question

I'm trying to implement .or for futures: return the same Option if result.is_some() and another Option type otherwise:

trait FutureOr {
    type Output;

    fn or(self, f: Self::Output) -> Self::Output;
}

impl <Fut, T> FutureOr for Fut where Fut: std::future::Future<Output = Option<T>> {
    type Output = Fut;

    fn or(self, f: Fut) -> Self::Output {

        self.then(|res| match res {
            Some(res) => async { Some(res) },
            None => async { f.await }
        })
    }
}

However, this doesn't compile:

`match` arms have incompatible types
expected opaque type `impl futures::Future<Output = std::option::Option<T>>` (`async` block)
   found opaque type `impl futures::Future<Output = std::option::Option<T>>` (`async` block)

Obviously, the types are the same, but from what I gather, the compiler is complaining about the fact that the opaque types might differ, i.e. the runtime type the generic will assume (?), could differ. So how can I indicate to the compiler I don't need their opaque types to differ?

I tried to:

• return the Future of the second match arm directly instead of wrapping it in an async block, and wrapping the first arm with std::future::ready(), but to my surprise this doesn't return a std::future::Future but a std::future::Ready, incompatible with the Future;
• wrapping the match statement in an async block, thus avoiding having two different concrete futures there, but the compiler warn this does not return a Future but a Then<Future>, and I don't know how to collect that into a Future (into_future() does not help, still returns Then<Future>, but why?).

Link to playground

---

UPDATE: OK, I think I am almost there. Compiler doesn't want the types to be impl:

error: unconstrained opaque type
  --> src/main.rs:19:15
   |
19 |     type Or = impl std::future::Future<Output = Option<T>>;
   |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   |
   = note: `Or` must be used in combination with a concrete type within the same module

---

UPDATE: Could get rid of impl by having generics, but compiler doesn't like that either:

error[E0207]: the type parameter `OrFut` is not constrained by the impl trait, self type, or predicates
  --> src/main.rs:15:11
   |
15 | impl<Fut, OrFut, OutFut, T> FutureOr for Fut
   |           ^^^^^ unconstrained type parameter

error[E0207]: the type parameter `OutFut` is not constrained by the impl trait, self type, or predicates
  --> src/main.rs:15:18
   |
15 | impl<Fut, OrFut, OutFut, T> FutureOr for Fut
   |                  ^^^^^^ unconstrained type parameter

Maybe I need to move the implementation to the function, trying that now...

---

UPDATE: Closest I could come. Only thing that seems left is understanding how to return a Future instead of a AndThen<Future>. This can probably be resolved by using the async_trait crate, letting the crate do the async de-sugaring for us, but I don't want to pull that in just because of a single method trait...

Answer 1

As written, your code is trying to treat two different futures (f, and the return value of fn or()) as the same type (Self::Output). This can't work, because every distinct implementation of Future is a different type. And, as you already discovered, it's not possible to name the type of an async block in a way that's useful for traits (until type_alias_impl_trait is stabilized).

Here is a rewrite that uses boxed futures, the usual workaround for TAIT. We carefully distinguish between the 3 future types: Self, Other, and the returned future.

use futures::{future::BoxFuture, Future};

trait FutureOr: Future + 'static {
    fn or<Other>(self, f: Other) -> BoxFuture<'static, Self::Output>
    where
        Other: Future<Output = Self::Output> + Send + 'static;
}

impl<Fut, T> FutureOr for Fut
where
    T: Send + 'static,
    Fut: std::future::Future<Output = Option<T>> + Send + 'static,
{
    fn or<Other>(self, f: Other) -> BoxFuture<'static, Self::Output>
    where
        Other: Future<Output = Self::Output> + Send + 'static,
    {
        Box::pin(async {
            match self.await {
                Some(res) => Some(res),
                None => f.await,
            }
        })
    }
}

The Send + 'static bounds are unfortunately needed since boxed futures have to pick a specific choice of lifetime and Send-or-!Send.

Once we have type_alias_impl_trait in a future stable version of Rust, we can avoid the boxing and the extra bounds:

#![feature(type_alias_impl_trait)]
use std::future::Future;

trait FutureOr<Other>: Future
where
    Other: Future<Output = Self::Output>,
{
    type Ored: Future<Output = Self::Output>;

    fn or(self, f: Other) -> Self::Ored;
}

type OptOred<Fut, Other, T>
where
    Fut: Future<Output = Option<T>>,
= impl Future<Output = Option<T>>;

impl<Fut, Other, T> FutureOr<Other> for Fut
where
    Fut: std::future::Future<Output = Option<T>>,
    Other: std::future::Future<Output = Self::Output>,
{
    type Ored = OptOred<Fut, Other, T>;

    fn or(self, f: Other) -> Self::Ored {
        async {
            match self.await {
                Some(res) => Some(res),
                None => f.await,
            }
        }
    }
}

type Ored is a special type alias which gets assigned the type of the async block future in fn or() because we returned that future from a function whose return type is that type alias. Note that Other now has to be a type parameter on the trait so that type Ored can refer to it, whereas in the first version the details were hidden inside dyn Future and Other didn't need to be mentioned.

Finally, we can still have a clean generic FutureOr on stable, at the price of manually writing a future. In this case, I've chosen to use the pin_project library to avoid this needing any newly written unsafe code, but it's still fairly complex:

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

trait FutureOr<Other>: Future
where
    Other: Future<Output = Self::Output>,
{
    type Ored: Future<Output = Self::Output>;

    fn or(self, f: Other) -> Self::Ored;
}

impl<Fut, Other, T> FutureOr<Other> for Fut
where
    Fut: std::future::Future<Output = Option<T>>,
    Other: std::future::Future<Output = Self::Output>,
{
    type Ored = OptOred<Fut, Other>;

    fn or(self, b: Other) -> Self::Ored {
        OptOred {
            skip_a: false,
            a: self,
            b,
        }
    }
}

#[pin_project::pin_project]
struct OptOred<A, B> {
    skip_a: bool,
    #[pin]
    a: A,
    #[pin]
    b: B,
}

impl<A, B, T> Future for OptOred<A, B>
where
    A: Future<Output = Option<T>>,
    B: Future<Output = Option<T>>,
{
    type Output = A::Output;

    fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<<Self as Future>::Output> {
        let proj = self.project();
        if !*proj.skip_a {
            match proj.a.poll(context) {
                Poll::Pending => Poll::Pending,
                r @ Poll::Ready(Some(_)) => r,
                Poll::Ready(None) => {
                    *proj.skip_a = true; // don't poll `a` again
                    proj.b.poll(context)
                }
            }
        } else {
            proj.b.poll(context)
        }
    }
}

Disclaimer: I haven't tested this code, only checked that it compiles.

Answer 2

The only way to get this to work today on stable Rust without boxing the returned future is to implement your own future. This is a bit of work, mostly due to pinning gymnastics. We need an additional generic type to represent the type of the provided "default value" future -- using Self::Output won't work because you're wanting to call this method with any type of future, but Self::Output is the type of future produced by or().

This type parameter has to go on the trait itself and not the or method, because the type will be used as part of the Output associated type in the implementation.

trait FutureOr<FDef> {
    type Output;

    fn or(self, f: FDef) -> Self::Output;
}

Our custom future type needs an instance of Self and FDef. Once Self has produced a value and if that value is None, we need to switch to polling the provided f future instead, so we need a way to signal to ourselves that the Self future has been polled to completion. We'll do this by wrapping it in an Option.

struct FutureOrFuture<T, F, FDef>
where F: Future<Output=Option<T>>,
    FDef: Future<Output=Option<T>>
{
    fut: Option<F>,
    def: FDef,
}

For the pinning gymnastics, we need a way to map Pin<&mut Self> to pinned wrappers around these fields. We'll write some helpers to make this ergonomic within poll().

impl<T, F, FDef> FutureOrFuture<T, F, FDef>
where F: Future<Output=Option<T>>,
    FDef: Future<Output=Option<T>>
{
    fn get_fut(self: Pin<&mut Self>) -> Option<Pin<&mut F>> {
        unsafe { self.get_unchecked_mut() }.fut.as_mut()
            .map(|r| unsafe { Pin::new_unchecked(r) })
    }
    
    fn get_def(self: Pin<&mut Self>) -> Pin<&mut FDef> {
        unsafe { self.map_unchecked_mut(|v| &mut v.def) }
    }
}

Finally, we can implement Future. We'll poll fut when it's Some, falling back on def otherwise.

impl<T, F, FDef> Future for FutureOrFuture<T, F, FDef>
where F: Future<Output=Option<T>>,
    FDef: Future<Output=Option<T>>
{
    type Output = Option<T>;
    
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.as_mut().get_fut() {
            Some(fut) => match fut.poll(ctx) {
                Poll::Ready(None) => {
                    // Clear self.fut; it's now consumed.
                    unsafe { self.as_mut().get_unchecked_mut() }.fut = None;
                    self.get_def().poll(ctx)
                },
                // Pending and Ready(Some(_))
                v => v,
            }
            None => self.get_def().poll(ctx)
        }
    }
}

Now the or() implementation just delegates to this future type.

impl<Fut, T, FDef> FutureOr<FDef> for Fut
where Fut: Future<Output=Option<T>>,
    FDef: Future<Output=Option<T>>
{
    type Output = FutureOrFuture<T, Self, FDef>;

    fn or(self, f: FDef) -> Self::Output {
        FutureOrFuture { fut: Some(self), def: f }
    }
}

(Playground)

---

Alternatively, if you can live with the overhead of boxing, the whole thing becomes a bit simpler since you can use .await -- but then you have to futz around with lifetimes. (You don't have to in the approach given above because any lifetimes are part of the generic type arguments of the resulting future type, so those lifetimes flow through by themselves. Boxing erases the compile-time types and so we have to do some work to get the lifetimes properly communicated.)

Note that this approach (returning a pinned box with lifetimes sprinkled in) is exactly what async_trait does for you.

The "incompatible match arms" issue is resolved here by wrapping the whole match block in async and awaiting f inside of it.

trait FutureOr {
    fn or<'a, 'b, 'c, 'd, T>(self, f: impl Future<Output=Self::Output> + Send + 'a)
    -> Pin<Box<dyn Future<Output=Option<T>> + Send + 'd>>
    where T: Send + 'b,
        Self: Future<Output=Option<T>> + Send + 'c,
        'a: 'd,
        'b: 'd,
        'c: 'd;
}

impl<Fut> FutureOr for Fut {
    fn or<'a, 'b, 'c, 'd, T>(self, f: impl Future<Output=<Self as Future>::Output> + Send + 'a)
    -> Pin<Box<dyn Future<Output=Option<T>> + Send + 'd>>
    where T: Send + 'b,
        Self: Future<Output=Option<T>> + Send + 'c,
        'a: 'd,
        'b: 'd,
        'c: 'd
    {
        Box::pin(self.then(|res| async move {
            match res {
                Some(res) => Some(res),
                None => f.await,
            }
        }))
    }
}

(Playground)