Search code examples
pointersasynchronousrustrust-tokiorust-futures

How to store a pointer to an async method in a container?


I have a struct that defines multiple async methods, and I'd like to store a pointer of each of them in a HashMap, so that I can call any method in one single line, knowing only a key that is given in parameter.

The aim here is to avoid as much as possible to have a huge match clause that would inflate more and more as I add new methods to my struct.

Methods all have the same signature:

async fn handle_xxx(&self, c: Command) -> Result<String, ()>

And I'd really like to call them the following way:

pub async fn execute_command(&mut self, command: Command) -> Result<String, ()> {
    let handler = self.command_callbacks.get(&command);
    let return_message: String = match handler {
        Some(f) => f(self, command).await.unwrap(), // The call is made here.
        None => return Err(()),
    };
    Ok(return_message)
}

However, obviously, in order to store something in a HashMap, you have to specify its type when declaring the HashMap, and that's when the trouble starts.

I tried the most obvious, which is declaring the wrapping function type:

type CommandExecutionNotWorking = fn(&CommandExecutor, Command) -> Future<Output = Result<String, ()>>;

Which does not work since Future is a trait, and not a type.

I tried to declare a generic type and specify it somewhere below:

type CommandExecutionNotWorkingEither<Fut> = fn(&CommandExecutor, Command) -> Fut;

But I encounter the same kind of issue since I need to specify the Future type, and have a HashMap declaration like following:

let mut command_callbacks: HashMap<
    Command,
    CommandExecutionFn<dyn Future<Output = Result<String, ()>>>,
> = HashMap::new();

impl Future obviously does not work since we're not in a function signature, Future either since it's not a type, and dyn Future creates a legitimate type mismatch.

Thus I tried to use Pin so that I can manipulate dyn Futures, and I ended up with the following signature:

type CommandExecutionStillNotWorking = fn(
    &CommandExecutor,
    Command,
) -> Pin<Box<dyn Future<Output = Result<String, ()>>>>;

But I need to manipulate functions that return Pin<Box<dyn Future<...>>> and not just Futures. So I tried to define a lambda that take an async function in parameter and returns a function that wraps the return value of my async method in a Pin<Box<...>>:

let wrap = |f| {
    |executor, command| Box::pin(f(&executor, command))
};

But the compiler is not happy since it expects me to define the type of f, which is what I tried to avoid here, so I'm back to square one.

Thus my question: Do you know if it's actually possible to write the type of an async function so that pointers on them can be easily manipulated like any variable or any other function pointer? Or should I go for another solution that may be less elegant, with a bit of duplication code or a huge match structure?


Solution

  • TL;DR: Yes, it is possible, but probably more complicated than you imagined.


    First, closures cannot be generic, and thus you'd need a function:

    fn wrap<Fut>(f: fn(&CommandExecutor, Command) -> Fut) -> CommandExecution
    where
        Fut: Future<Output = Result<String, ()>>
    {
        move |executor, command| Box::pin(f(executor, command))
    }
    

    But then you cannot turn the returned closure into a function pointer because it captures f.

    Now technically it should be possible since we only want to work with function items (that are non-capturing), and their type (unless converted into a function pointer) is zero-sized. So by the type alone we should be able to construct an instance. But doing that requires unsafe code:

    fn wrap<Fut, F>(_f: F) -> CommandExecution
    where
        Fut: Future<Output = Result<String, ()>>,
        F: Fn(&CommandExecutor, Command) -> Fut,
    {
        assert_eq!(std::mem::size_of::<F>(), 0, "expected a fn item");
        move |executor, command| {
            // SAFETY: `F` is a ZST (checked above), any (aligned!) pointer, even crafted
            // out of the thin air, is valid for it.
            let f: &F = unsafe { std::ptr::NonNull::dangling().as_ref() };
            Box::pin(f(executor, command))
        }
    }
    

    (We need _f as a parameter because we cannot specify the function type; let inference find it out itself.)

    However, are troubles don't end here. They just start.

    Now we get the following error:

    error[E0310]: the parameter type `Fut` may not live long enough
      --> src/lib.rs:28:9
       |
    18 | fn wrap<Fut, F>(_f: F) -> CommandExecution
       |         --- help: consider adding an explicit lifetime bound...: `Fut: 'static`
    ...
    28 |         Box::pin(f(executor, command))
       |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ...so that the type `Fut` will meet its required lifetime bounds
    

    Well, it suggests a solution. Let's try...

    It compiles! Successfully!!

    ...until we actually try to make use of it:

    let mut _command_callbacks: Vec<CommandExecution> = vec![
        wrap(CommandExecutor::handle_xxx),
        wrap(CommandExecutor::handle_xxx2),
    ];
    

    (a HashMap will have the same effect).

    error[E0308]: mismatched types
      --> src/lib.rs:34:9
       |
    34 |         wrap(CommandExecutor::handle_xxx),
       |         ^^^^ lifetime mismatch
       |
       = note: expected associated type `<for<'_> fn(&CommandExecutor, Command) -> impl Future<Output = Result<String, ()>> {CommandExecutor::handle_xxx} as FnOnce<(&CommandExecutor, Command)>>::Output`
                  found associated type `<for<'_> fn(&CommandExecutor, Command) -> impl Future<Output = Result<String, ()>> {CommandExecutor::handle_xxx} as FnOnce<(&CommandExecutor, Command)>>::Output`
       = note: the required lifetime does not necessarily outlive the static lifetime
    note: the lifetime requirement is introduced here
      --> src/lib.rs:21:41
       |
    21 |     F: Fn(&CommandExecutor, Command) -> Fut,
       |                                         ^^^
    
    error[E0308]: mismatched types
      --> src/lib.rs:35:9
       |
    35 |         wrap(CommandExecutor::handle_xxx2),
       |         ^^^^ lifetime mismatch
       |
       = note: expected associated type `<for<'_> fn(&CommandExecutor, Command) -> impl Future<Output = Result<String, ()>> {CommandExecutor::handle_xxx2} as FnOnce<(&CommandExecutor, Command)>>::Output`
                  found associated type `<for<'_> fn(&CommandExecutor, Command) -> impl Future<Output = Result<String, ()>> {CommandExecutor::handle_xxx2} as FnOnce<(&CommandExecutor, Command)>>::Output`
       = note: the required lifetime does not necessarily outlive the static lifetime
    note: the lifetime requirement is introduced here
      --> src/lib.rs:21:41
       |
    21 |     F: Fn(&CommandExecutor, Command) -> Fut,
       |                                         ^^^
    

    The problem is described in Lifetime of a reference passed to async callback. The solution is to use a trait to workaround the problem:

    type CommandExecution = for<'a> fn(
        &'a CommandExecutor,
        Command,
    ) -> Pin<Box<dyn Future<Output = Result<String, ()>> + 'a>>;
    
    trait CommandExecutionAsyncFn<CommandExecutor>:
        Fn(CommandExecutor, Command) -> <Self as CommandExecutionAsyncFn<CommandExecutor>>::Fut
    {
        type Fut: Future<Output = Result<String, ()>>;
    }
    
    impl<CommandExecutor, F, Fut> CommandExecutionAsyncFn<CommandExecutor> for F
    where
        F: Fn(CommandExecutor, Command) -> Fut,
        Fut: Future<Output = Result<String, ()>>,
    {
        type Fut = Fut;
    }
    
    fn wrap<F>(_f: F) -> CommandExecution
    where
        F: 'static + for<'a> CommandExecutionAsyncFn<&'a CommandExecutor>,
    {
        assert_eq!(std::mem::size_of::<F>(), 0, "expected a fn item");
        move |executor, command| {
            // SAFETY: `F` is a ZST (checked above), any (aligned!) pointer, even crafted
            // out of the thin air, is valid for it.
            let f: &F = unsafe { std::ptr::NonNull::dangling().as_ref() };
            Box::pin(f(executor, command))
        }
    }
    

    I will not expand on why this thing is needed or how it solves the problem. You can find an explanation in the linked question and the questions linked in it.

    And now our code works. Like, truly so.

    However, think carefully if you really want all of this stuff: it may be easier to just change the functions to return a boxed future.