Changing the return type of add_to_linker closures

[lann]

From this Zulip discussion:

I've been experimenting with ways to provide reusable bits of host functionality and have run into a stumbling block with wasmtime's generated add_to_linker functions, which take getclosures with types (roughly) like:

get: impl Fn(&mut T) -> &mut U + Send + Sync + Copy + 'static
    ...
    where U: Host,

These closures are usually trivial field accesses like |data| &mut data.some_host_impl, but this is somewhat limiting because you can only return &muts borrowing from T. As a specific example, the wasmtime-wasi WasiView::table needs to be shared with wasmtime-wasi-http's WasiHttpView::table, but this currently requires both *View implementations to use the same concrete type within T, making the two host implementations uncomposable.

There are some unsafe approaches to deal with this but I'm wondering if it would be possible to change the add_to_linker closure signature to permit returning non-ref types, e.g. something that impls AsMut<U>. I've been trying to find a way to do this backward-compatibility but I'm running into the limits of my understanding of closure lifetime

[lann]

@alexcrichton:
I can definitely say that the &mut T -> &mut U closures I've never felt great about. I've always preferred to do &mut T -> U with trait bounds on U, which I think is what you're getting at as well. I've also tried to stretch my Rust knowledge of closures to achieve this and while it's possible it gets really nasty really quickly.

The next-closest approximation I've thought of is to allow returning &mut dyn Trait but I haven't gotten around to pursuing that much.

I was experimenting with this to clean up wasmtime-wasi APIs a bit ago but I seem to have misplace the branch

I asked about this on the Rust user forums a long time ago too

[alexcrichton]

One other point I'd add is that I wouldn't take strict backwards compatibility as a constraint here. Everything that works today should continue to work, but beyond that it's fine if callers need a slightly different function call or annotations or such.

[lann]

It turns out the closure lifetime problems were due to a type inference limitation: rust-lang/rust#58052

This seems to work as a rough approach to improve on:

use std::borrow::BorrowMut;

trait BorrowHost<T, U>: Send + Sync + 'static {
    fn borrow_host<'a>(&self, data: &'a mut T) -> impl BorrowMut<U>
    where
        U: 'a;
}

impl<T, U, F> BorrowHost<T, U> for F
where
    F: Fn(&mut T) -> &mut U + Send + Sync + Copy + 'static,
{
    fn borrow_host<'a>(&self, data: &'a mut T) -> impl BorrowMut<U> + 'a
    where
        U: 'a,
    {
        self(data)
    }
}

pub fn add_to_linker<T, U>(
    linker: &mut wasmtime::component::Linker<T>,
    get: impl Fn(&mut T) -> &mut U + Send + Sync + Copy + 'static,
) -> wasmtime::Result<()>
where
    U: Host,
{
    add_to_linker_borrow_host(linker, get)
}

fn add_to_linker_borrow_host<T, U: Host>(
    linker: &mut wasmtime::component::Linker<T>,
    getter: impl BorrowHost<T, U>,
) -> wasmtime::Result<()> {
    let mut inst = linker.instance("foo")?;
    inst.func_wrap(
        "bar",
        move |mut caller: wasmtime::StoreContextMut<'_, T>, (): ()| {
            let mut host = getter.borrow_host(caller.data_mut());
            host.borrow_mut().host_func();
            Ok(())
        },
    )
}

// add_to_linker(linker, |data| &mut data.host)?;

[alexcrichton]

Nice that looks reasonable to me 👍

[lann]

After some more iteration here I've come up with something that solves my problems. I'm not really happy with it but its the best compromise I've been able to make with the borrow checker: https://gist.github.com/lann/8d85f69ef334051206ee1c7c0897f0f4

If this looks acceptable I can PR it, though I'd also be happy if someone was able to refine the design further.

[alexcrichton]

Some minor cleanups on that:

• Testing locally I think you can drop the HostBorrowFn struct entirely in favor of impl<T, U, F> GetHost<T> for F directly.
• impl<T: Host> Host for &mut T { I think should be auto-synthesized by bindgen! and I don't think there's any issue with that.

One point I think is pretty important is that we'll need GetHost per-Host-trait which is a lot of them. That would be untenable if HostBorrowFn were required because you'd need a different-type-per-closure, but with being able to implement GetHost<T> for F directly that means that a single closure should be able to work with all Host traits just fine.

One other minor possibility might be to have:

    pub trait GetHost: Send + Sync + 'static {
        type T;
        type Host<'a>: Host;
        fn get_host<'a>(&self, data: &'a mut Self::T) -> Self::Host<'a>;
    }

As I write that out though I don't think I like that, so probably scratch that.

[alexcrichton]

Oh also I'd generalize impl<T: Host> Host for &mut T to impl<T: Host + ?Sized> Host for &mut T for the most-general-blanket-forward-impl

[lann]

Testing locally I think you can drop the HostBorrowFn struct entirely in favor of impl<T, U, F> GetHost<T> for F directly.

I added HostBorrowFn to work around some coherence issue but that may have been with a different signature on add_to_linker_get_host; it does seem to work without now.

[alexcrichton]

Heh yeah when I've tried to golf this stuff historically I've found that it's a very fine line to walk between rustc and usability here...

[lann]

I'll start on the PR. I also updated the gist with your suggestions in case someone is looking here before the PR is ready.

[lann]

One other minor possibility might be to have: (GAT stuff)

I tried many variations with GATs but ended up playing whack-a-mole with various lifetime constraints. It may be possible to express this with a single generic GetHost but I don't seem to be capable of figuring it out. 😬

[lann]

This is what I was (vaguely) worried about: https://github.com/bytecodealliance/wasmtime/actions/runs/8805989303/job/24169783936?pr=8448#step:17:392

e.g.

error[E0119]: conflicting implementations of trait `async_io::wasi::sockets::ip_name_lookup::Host` for type `&mut _`
   --> crates/wasi/src/ip_name_lookup.rs:23:1
    |
23  |   impl<T: WasiView> Host for T {
    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ conflicting implementation for `&mut _`

[lann]

From an out-of-band discussion, we're going to try:

• Adding a add_to_linker_with_closure: false that allows opting-out of the generated add_to_linker, impl Host for &mut Host, and impl GetHost for Fn(&mut T) -> &mut U
• Update wasi(-http) to set the above opt-out flag and change impls as necessary (e.g. impl ... for dyn WasiView)