#async-trait #traits #async #variadic #future #middleware #request-response

async-variadic

Simple async variadic functions in Rust with trait bounds

3 stable releases

1.1.1 Jul 18, 2022
1.0.0 Jul 17, 2022

#631 in Asynchronous

MIT license

19KB
102 lines

async-variadic

Latest version Documentation License

Provides a way to pass along async functions with a trait that can support n-ary arguments. This code is inspired from the use of trait specialization in order to support variadics.

Examples

use async_variadic::AsyncFn;

async fn min() {}
async fn test2(_s: String) -> i32 { 3 }
async fn test3(a: i32, b: i32) -> i32 { a + b }

fn assert_impl_fn<T, O>(_: impl AsyncFn<T, O>) {}

assert_impl_fn(min);
assert_impl_fn(test2);
assert_impl_fn(test3);

Since the output is constrained on the trait, we can even add trait bounds to the input and the output like so:

struct Request {}
struct Response {}
struct Body {}

impl From<Request> for (Request,) {
    fn from(req: Request) -> Self {
        (req,)
    }
}

impl From<Request> for (Body,) {
    fn from(_req: Request) -> Self {
        (Body {},)
    }
}

impl From<&'static str> for Response {
    fn from(_s: &'static str) -> Self {
        Response {}
    }
}

fn assert_impl_output<T: From<Request>, O: Into<Response>>(_: impl AsyncFn<T, O>) {}

#[test]
fn test_trait_bounds() {
    async fn with_request_resp(_req: Request) -> &'static str {
        "hello"
    }

    async fn with_body_resp(_body: Body) -> &'static str {
        "hello"
    }

    assert_impl_output(with_request_resp);
    assert_impl_output(with_body_resp);
}

actix-web handlers + use of Into/From Traits

actix-web has an implementation of this used for handler functions that convert FromRequest and return any type that implements Responder. This is how you can write a web application that can bind in an arbitrary number of arguments at any position and return any type of data (so long as they implement the corresponding traits).

Background: Conflicting Implementations

In the background to this, you may try to implement a trait variation for any Function that has a given number of arguments. Looking at the example provided in the trait specialization article

trait VariadicFunction {
  fn call(&self, req: &[f64]) -> f64;
}

impl<Func> VariadicFunction for Func
where Func : Fn(f64)->f64 {
  fn eval(&self,args : &[f64])->f64 {
    (self)(args[0])
  }
}

impl<Func> VariadicFunction for Func
where Func : Fn(f64,f64)->f64 {
  fn eval(&self,args : &[f64])->f64 {
    (self)(args[0],args[1])
  } }

fn evaluate<F:VariadicFunction>(func : F, args: &[f64]) -> f64{
  func.eval(args)
}

The compiler will output something like:

error[E0119]: conflicting implementations of trait `VariadicFunction`:
  --> src/lib.rs:12:1
   |
5  | / impl<Func> VariadicFunction for Func
6  | | where Func : Fn(f64)->f64 {
7  | |   fn eval(&self,args : &[f64])->f64 {
8  | |     (self)(args[0])
9  | |   }
10 | | }
   | |_- first implementation here
11 |
12 | / impl<Func> VariadicFunction for Func
13 | | where Func : Fn(f64,f64)->f64 {
14 | |   fn eval(&self,args : &[f64])->f64 {
15 | |     (self)(args[0],args[1])
16 | |   }
17 | | }
   | |_^ conflicting implementation

error: aborting due to previous error

I've had similar issues when trying to create a Middleware trait that can handle a varying number of arguments. I wanted to do this specifically because I thought it might be nice to pass in closure that only specified one or more of the arguments.

#[async_trait]
pub trait Middleware<'a, 'b>: Send + Sync {
    #[must_use = "handle future must be used"]
    async fn handle(&self, request: &'a mut Request, response: &'b mut Response);
}

#[async_trait]
impl<'a, 'b, F, Fut, Res> Middleware<'a, 'b> for F
where
    F: Send + Sync + 'a + Fn(&'a mut Request, &'b mut Response) -> Fut,
    Fut: Future<Output = Res> + Send + 'b,
{
    async fn handle(&self, request: &'a mut Request, response: &'b mut Response) {
        (self)(request, response).await;
    }
}

Here I am bounding the F type to a Fn trait that takes two arguments. If I wanted to specify a different closure that only passed along the request that would look like this:

#[async_trait]
pub trait Middleware<'a, 'b>: Send + Sync {
    #[must_use = "handle future must be used"]
    async fn handle(&self, request: &'a mut Request, response: &'b mut Response);
}

#[async_trait]
impl<'a, 'b, F, Fut, Res> Middleware<'a, 'b> for F
where
    F: Send + Sync + 'a + Fn(&'a mut Request, &'b mut Response) -> Fut,
    Fut: Future<Output = Res> + Send + 'b,
{
    async fn handle(&self, request: &'a mut Request, response: &'b mut Response) {
        (self)(request, response).await;
    }
}

#[async_trait]
impl<'a, 'b, F, Fut, Res> Middleware<'a, 'b> for F
where
    F: Send + Sync + 'a + Fn(&'a mut Request) -> Fut,
    Fut: Future<Output = Res> + Send + 'b,
{
    async fn handle(&self, request: &'a mut Request, response: &'b mut Response) {
        (self)(request).await;
    }
}

Unfortunately I get the same compilation error as above. The main reason (as discussed in the article and in issue #60074) is that a closure could implement both Fn traits and create undefined behavior.

impl FnOnce<(u32,)> for Foo {
    type Output = u32;
    extern "rust-call" fn call_once(self, args: (u32,)) -> Self::Output {
        args.0
    }
}

impl FnOnce<(u32, u32)> for Foo {
    type Output = u32;
    extern "rust-call" fn call_once(self, args: (u32, u32)) -> Self::Output {
        args.0
    }
}

So that doesn't work for us. We need to find another way!

Trait Specialization

As mentioned in the article, the way around this is to create a trait specialization type that takes Args.

trait VariadicFunction<ArgList> {
  fn eval(&self, args: &[f64]) -> f64;
}

Then we can implement a Fn trait bound for each of the positional arguments while providing that type to the VariadicFunction.

impl<Func> VariadicFunction<f64> for Func
where Func : Fn(f64)->f64 {
  fn eval(&self,args : &[f64])->f64 {
    (self)(args[0])
  }
}

impl<Func> VariadicFunction<(f64,f64)> for Func
where Func : Fn(f64,f64)->f64 {
  fn eval(&self,args : &[f64])->f64 {
    (self)(args[0],args[1])
  }
}

fn evaluate<ArgList, F>(func : F, args: &[f64]) -> f64
where F: VariadicFunction<ArgList>{
  func.eval(args)
}

Notice that the VariadicFunction trait is actually 2 different traits once it becomes monomorphized by the compiler. We're implementing the entirely different traits for the functions that have the different arguments. Let's expand this for async!

Abstracted out to AsyncFn

An async function is nothing more than a function that returns a Future.

pub trait Future {
    type Output;

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

The associated output is the result of the Poll<Self::Output> returning Poll::Ready(output). We can define an AsyncFn trait that looks similar to the Future trait except it's main purpose is to be a trait bound for functions.

pub trait AsyncFn<Args> {
    type Output;
    type Future: Future<Output = Self::Output>;

    fn call(&self, args: Args) -> Self::Future;
}

Now all that's left is implementing the trait for different closures. The simplest case is going to be an empty async function.

impl<Func, Fut> AsyncFn<()> for Func
where
    Func: Fn() -> Fut,
    Fut: Future
{
    type Output = Fut::Output;
    type Future = Fut;

    fn call(&self, args: ()) -> Self::Future {
        (self)()
    }
}

The second case is going to be for a function that takes 1 argument. Notice that we don't have to specify the output since the Future::Output already specifies that for us.

impl<Func, Fut, A> AsyncFn<(A,)> for Func
where
    Func: Fn(A) -> Fut,
    Fut: Future
{
    type Output = Fut::Output;
    type Future = Fut;

    fn call(&self, args: (A,)) -> Self::Future {
        let (a,) = args;
        (self)(a)
    }
}

We can amend the call function here to destructure the arguments tuple into the local variable.

impl<Func, Fut, A> AsyncFn<(A,)> for Func
where
    Func: Fn(A) -> Fut,
    Fut: Future
{
    type Output = Fut::Output;
    type Future = Fut;

    fn call(&self, (A,): (A,)) -> Self::Future {
        (self)(A)
    }
}

This can be expanded upon to support more than one argument similar to the above implementation.

impl<Func, Fut, A, B> AsyncFn<(A,B)> for Func
where
    Func: Fn(A) -> Fut,
    Fut: Future
{
    type Output = Fut::Output;
    type Future = Fut;

    fn call(&self, (A,B): (A,B)) -> Self::Future {
        (self)(A, B)
    }
}

Automatically Implementing with a Macro

While it's simple enough to implement each of these variations manually, we can also write a simple macro to automatically build these for us.

macro_rules! ary ({ $($param:ident)* } => {
    impl<Func, Fut, $($param,)*> AsyncFn<($($param,)*)> for Func
    where
        Func: Fn($($param),*) -> Fut,
        Fut: Future
    {
        type Output = Fut::Output;
        type Future = Fut;

        #[inline]
        #[allow(non_snake_case)]
        fn call(&self, ($($param,)*): ($($param,)*)) -> Self::Future {
            (self)($($param,)*)
        }
    }
});

This macro simply takes the $(param:ident)* pattern and specifies it as both the generic argument to the impl<>, the trait specialization in the AsyncFn<T>, the trait bound on the Func arguments, and finally in the actual call function. Now all we need to do is call this macro with different identifiers and it's good to go.

ary! {}
ary! { A }
ary! { A B }
ary! { A B C }
ary! { A B C D }
ary! { A B C D E }
ary! { A B C D E F }
ary! { A B C D E F G }
ary! { A B C D E F G H }
ary! { A B C D E F G H I }
ary! { A B C D E F G H I J }
ary! { A B C D E F G H I J K }
ary! { A B C D E F G H I J K L }

This will let us specify up to 12 different parameters. Each new generic identifier will be included in an entirely new trait implementation.

Tests

The tests below show us how we can now specify any kind of async function that captures any type of variable (in any order) and returns any type (as indicated by the Future::Output).

#[cfg(test)]
mod tests {
    use super::*;

    fn assert_impl_fn<T>(_: impl AsyncFn<T>) {}

    #[test]
    fn test_args() {
        async fn min() {}
        async fn min_output() -> i32 {
            4
        }
        async fn with_req(_req: String) -> &'static str {
            "foo"
        }
        #[rustfmt::skip]
        #[allow(clippy::too_many_arguments, clippy::just_underscores_and_digits)]
        async fn max(
            _01: (), _02: (), _03: (), _04: (), _05: (), _06: (),
            _07: (), _08: (), _09: (), _10: (), _11: (), _12: (),
        ) {}

        assert_impl_fn(min);
        assert_impl_fn(with_req);
        assert_impl_fn(min_output);
        assert_impl_fn(max);
    }
}

The main use case here is now we can have a function like assert_impl_fn<T> that actually takes an async function or closure as an argument! This is especially useful for our handlers, middleware and web server implementations.

Trait Bounds

In order to allow for trait bounds on the output the trait will need to be modified to use async-trait. This works nicely for varying the input and output.

#[async_trait]
pub trait AsyncFn<Args, Output> {
    async fn call(&self, args: Args) -> Output;
}

With the macro now updated to require Send and Sync to support async-trait.

macro_rules! ary ({ $($param:ident)* } => {
    #[async_trait::async_trait]
    impl<Func, Fut, $($param:Send + 'static,)*> AsyncFn<($($param,)*), Fut::Output> for Func
    where
        Func: Send + Sync + Fn($($param),*) -> Fut,
        Fut: Future + Send
    {
        #[inline]
        #[allow(non_snake_case)]
        async fn call(&self, ($($param,)*): ($($param,)*)) -> Fut::Output {
            (self)($($param,)*).await
        }
    }
});

With this change, while it does box the output of the future, I believe is a net gain to the ergonomics as now we can do this:

struct Request {}
struct Response {}
struct Body {}

impl From<Request> for (Request,) {
    fn from(req: Request) -> Self {
        (req,)
    }
}

impl From<Request> for (Body,) {
    fn from(_req: Request) -> Self {
        (Body {},)
    }
}

impl From<&'static str> for Response {
    fn from(_s: &'static str) -> Self {
        Response {}
    }
}

fn assert_impl_output<T: From<Request>, O: Into<Response>>(_: impl AsyncFn<T, O>) {}

#[test]
fn test_trait_bounds() {
    async fn with_request_resp(_req: Request) -> &'static str {
        "hello"
    }

    async fn with_body_resp(_body: Body) -> &'static str {
        "hello"
    }

    assert_impl_output(with_request_resp);
    assert_impl_output(with_body_resp);
}

Dependencies

~315–770KB
~18K SLoC