Generic closure returns that can capture arguments

I've been struggling recently to model a library function where the goal of the library is to provide host functionality to an internal runtime. Host-provided functions can be provided so long as arguments and return values implement a particular trait in my crate. What I would ideally like to write is a function in my library that takes a host function as a parameter. Another feature of my library is that users can configure a "context" that is received on each host function call which reports current runtime state to the closure. I would like the closure to be able to capture over the returned context state in the return value of the closure, so I've been trying to write a function signature like so:

fn register_host_fn(host_fn: impl Fn(&Context) -> (impl HostReturn + '_)) {
    // ...
}

This generates a compile error however:

error[E0562]: `impl Trait` only allowed in function and inherent method return types, not in `Fn` trait return
 --> src/lib.rs:5:52
  |
5 | fn register_host_fn(host_fn: impl Fn(&Context) -> (impl HostReturn + '_)) {
  |                                                    ^^^^^^^^^^^^^^^^^^^^

For more information about this error, try `rustc --explain E0562`.

I have tried a number of ways to implement this function signature and my current belief is that this probably isn't possible in today's stable Rust. I am wondering, though, if future features on the horizon such as impl-trait-in-return or associated types intended on having this sort of pattern be able to compile? Is there perhaps a nightly feature I can test out today to ensure that it works? My initial attempts so far have been unfruitful.

A full example of what I've been trying to get to compile which replaces Context with a dummy u8 parameter looks like:

trait HostReturn {}

impl HostReturn for u8 {}
impl<T: HostReturn> HostReturn for &'_ T {}

fn register_host_fn(host_fn: impl Fn(&u8) -> (impl HostReturn + '_)) {
    // ...
}

fn main() {
   register_host_fn(|borrow| borrow); 
}

which I have not been able to compile given any combination of nightly features or signatures of register_host_fn.

fn register_host_fn<'a, R: HostReturn, F>(host_fn: F) 
    where F: Fn(&'a u8) -> R, R: 'a {
    // ...
}

This works.

You have to define a helper trait that aliases Fn. And at the call site you have you use a helper function to get closure type inference to work correctly.

trait HostReturn {}

impl HostReturn for u8 {}
impl<T: HostReturn> HostReturn for &'_ T {}

fn register_host_fn<F: for<'a> HostFn<&'a u8>>(host_fn: F) {
    // ...
}

trait HostFn<I>: Fn(I) -> <Self as HostFn<I>>::Output {
    type Output: HostReturn;
}
impl<F, I, O: HostReturn> HostFn<I> for F
where
    F: Fn(I) -> O,
{
    type Output = O;
}

fn main() {
    fn funnel<F: Fn(&u8) -> &u8>(f: F) -> F { f }

   register_host_fn(funnel(|borrow| borrow));
}

playground

The call site can be make sweeter by using the library higher-order-closure:

register_host_fn(hrtb!(|borrow: &u8| -> &u8 { borrow }));
5 Likes

Ah sorry I should have expanded my example but this unfortunately doesn't work if the host function is called since the type signature has the caller picking 'a when the callee needs to be able to pick 'a (with some temporary lifetime)

Oh wow I had no idea about a trick like that! I will look more into this and try to understand what's going on here. Thanks for the help!

As a follow-up question, my actual use case is actually more complicated than I alluded to here. I'm working on a wasm runtime where host functions are given "runtime state" as well as "host state" but can only return a borrow to the "host state" since the runtime state is then later used to copy the returned value back into wasm. Code-wise what I ideally want is:

fn register_host_fn(host_fn: impl for<'a> Fn(&mut Runtime, &'a mut Host) -> (impl HostReturn + 'a)) {
    // ...
}

I don't fully understand how the lifetimes all work out above but it looks like the lifetime bounds (implicitly inserted by the compiler?) mean that the Output type can close over any of the input arguments, meaning my first attempt to do this is failing to compile. Is there a way to distinguish in the trait bounds specifically what the return value can capture?

I realize though that I'm getting quite far afield from my original use case and this is probably all too unergonomic or unwieldy to even use anyway. Nevertheless I'm still curious if this is possible to achieve one way or another.

You have to move the first parameter into the trait alias: Playground

3 Likes