rust: store local wrapper addresses in async stream/future vtables

crates/rust/src/interface.rs

@@ -684,6 +684,18 @@ macro_rules! {macro_name} {{
             PayloadFor::Future => "",
             PayloadFor::Stream => ", _: usize",
         };
+        // On wasm the VTABLE stores the address of a local wrapper rather than
+        // the canonical import directly (see the wrapper functions below). The
+        // wrappers need *named* extra params to forward, where the extern block
+        // and native stubs only need the type.
+        let start_extra_named = match payload_for {
+            PayloadFor::Future => "",
+            PayloadFor::Stream => ", _extra: usize",
+        };
+        let start_extra_arg = match payload_for {
+            PayloadFor::Future => "",
+            PayloadFor::Stream => ", _extra",
+        };
         let mut lift_fn = format!("unsafe fn lift(ptr: *mut u8) -> {name} {{ {lift} }}");
         let mut lower_fn = format!("unsafe fn lower(value: {name}, ptr: *mut u8) {{ {lower} }}");
         let mut dealloc_lists_fn =
@@ -733,25 +745,49 @@ pub mod vtable{ordinal} {{
     #[cfg(not(target_arch = "wasm32"))]
     unsafe extern "C" fn start_write(_: u32, _: *const u8{start_extra}) -> u32 {{ unreachable!() }}
 
+    // Work around a behavior of LLD where in a shared library when an address
+    // is taken of an imported function that only shows up as a `GOT.func`
+    // import which means there's not actual import for wit-component to
+    // generate bindings for. This is worked around with local Rust functions
+    // that simply delegate to the imports. Seems to work for now even though
+    // it seems like optimizations should defeat this, but it's not all that
+    // much more indirection and in theory gets things working for now.
     #[cfg(target_arch = "wasm32")]
-    #[link(wasm_import_module = "{module}")]
-    unsafe extern "C" {{
-        #[link_name = "[{import_prefix}-new-{index}]{func_name}"]
-        fn new() -> u64;
-        #[link_name = "[{import_prefix}-cancel-write-{index}]{func_name}"]
-        fn cancel_write(_: u32) -> u32;
-        #[link_name = "[{import_prefix}-cancel-read-{index}]{func_name}"]
-        fn cancel_read(_: u32) -> u32;
-        #[link_name = "[{import_prefix}-drop-writable-{index}]{func_name}"]
-        fn drop_writable(_: u32);
-        #[link_name = "[{import_prefix}-drop-readable-{index}]{func_name}"]
-        fn drop_readable(_: u32);
-        #[link_name = "[async-lower][{import_prefix}-read-{index}]{func_name}"]
-        fn start_read(_: u32, _: *mut u8{start_extra}) -> u32;
-        #[link_name = "[async-lower][{import_prefix}-write-{index}]{func_name}"]
-        fn start_write(_: u32, _: *const u8{start_extra}) -> u32;
+    mod imports {{
+        #[link(wasm_import_module = "{module}")]
+        unsafe extern "C" {{
+            #[link_name = "[{import_prefix}-new-{index}]{func_name}"]
+            pub(super) fn new() -> u64;
+            #[link_name = "[{import_prefix}-cancel-write-{index}]{func_name}"]
+            pub(super) fn cancel_write(_: u32) -> u32;
+            #[link_name = "[{import_prefix}-cancel-read-{index}]{func_name}"]
+            pub(super) fn cancel_read(_: u32) -> u32;
+            #[link_name = "[{import_prefix}-drop-writable-{index}]{func_name}"]
+            pub(super) fn drop_writable(_: u32);
+            #[link_name = "[{import_prefix}-drop-readable-{index}]{func_name}"]
+            pub(super) fn drop_readable(_: u32);
+            #[link_name = "[async-lower][{import_prefix}-read-{index}]{func_name}"]
+            pub(super) fn start_read(_: u32, _: *mut u8{start_extra}) -> u32;
+            #[link_name = "[async-lower][{import_prefix}-write-{index}]{func_name}"]
+            pub(super) fn start_write(_: u32, _: *const u8{start_extra}) -> u32;
+        }}
     }}
 
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn new() -> u64 {{ unsafe {{ imports::new() }} }}
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn cancel_write(handle: u32) -> u32 {{ unsafe {{ imports::cancel_write(handle) }} }}
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn cancel_read(handle: u32) -> u32 {{ unsafe {{ imports::cancel_read(handle) }} }}
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn drop_writable(handle: u32) {{ unsafe {{ imports::drop_writable(handle) }} }}
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn drop_readable(handle: u32) {{ unsafe {{ imports::drop_readable(handle) }} }}
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn start_read(handle: u32, ptr: *mut u8{start_extra_named}) -> u32 {{ unsafe {{ imports::start_read(handle, ptr{start_extra_arg}) }} }}
+    #[cfg(target_arch = "wasm32")]
+    unsafe extern "C" fn start_write(handle: u32, ptr: *const u8{start_extra_named}) -> u32 {{ unsafe {{ imports::start_write(handle, ptr{start_extra_arg}) }} }}
+
     {lift_fn}
     {lower_fn}
     {dealloc_lists_fn}

crates/test/src/lib.rs

@@ -1087,6 +1087,29 @@ status: {}",
         bail!("{error}")
     }
 
+    /// Converts the list of dynamic libraries in `dylibs` into a component and places it
+    /// in the destination specified by `compile`.
+    ///
+    /// This is similar to `convert_p1_to_component` except usese a
+    /// `wit_component::Linker` instead of a `wit_component::ComponentEncoder`.
+    fn link_dylibs_to_component(&self, dylibs: &[PathBuf], compile: &Compile<'_>) -> Result<()> {
+        let mut linker = wit_component::Linker::default();
+        for dylib in dylibs {
+            let dylib_bytes =
+                fs::read(dylib).with_context(|| format!("failed to read dylib file {dylib:?}"))?;
+            let name = dylib
+                .file_name()
+                .and_then(|s| s.to_str())
+                .context("non-utf-8 dylib filename")?;
+            linker
+                .library(name, &dylib_bytes, false)
+                .with_context(|| format!("failed to register {name}"))?;
+        }
+        let component = linker.encode().context("failed to link")?;
+        write_if_different(compile.output, component)?;
+        Ok(())
+    }
+
     /// Converts the WASIp1 module at `p1` to a component using the information
     /// stored within `compile`.
     ///

crates/test/src/rust.rs

@@ -40,10 +40,22 @@ struct RustConfig {
     /// Space-separated list or array of compiler flags to pass.
     #[serde(default)]
     rustflags: StringList,
+
     /// List of path to rust files to build as external crates and link to the
     /// main crate.
     #[serde(default)]
     externs: Vec<String>,
+
+    /// Whether or not to link the main crate as a shared library.
+    ///
+    /// This is implied if `extern_dylibs` is specified.
+    #[serde(default)]
+    link_shared: bool,
+
+    /// Paths to Rust files to build as dynamic libraries. If non-empty the
+    /// main file is also built as a dynamic library.
+    #[serde(default)]
+    extern_dylibs: Vec<String>,
 }
 
 #[derive(Deserialize)]
@@ -208,6 +220,7 @@ path = 'lib.rs'
 
     fn compile(&self, runner: &Runner, compile: &Compile) -> Result<()> {
         let config = compile.component.deserialize_lang_config::<RustConfig>()?;
+        let link_shared = config.link_shared || !config.extern_dylibs.is_empty();
 
         // If this rust target doesn't natively produce a component then place
         // the compiler output in a temporary location which is componentized
@@ -216,8 +229,16 @@ path = 'lib.rs'
 
         // Compile all extern crates, if any
         let mut externs = Vec::new();
+        let mut dylibs = Vec::new();
         let manifest_dir = compile.component.path.parent().unwrap();
 
+        let wasi_sdk_path = if link_shared {
+            let path = runner.opts.c.wasi_sdk_path.as_ref();
+            Some(path.ok_or_else(|| anyhow::anyhow!("need a wasi-sdk-path"))?)
+        } else {
+            None
+        };
+
         let rustc = |path: &Path, output: &Path| {
             // Compile the main crate, passing `--extern` for all upstream crates.
             let mut cmd = runner.rustc(Edition::E2021);
@@ -228,9 +249,29 @@ path = 'lib.rs'
             for flag in Vec::from(config.rustflags.clone()) {
                 cmd.arg(flag);
             }
+            if link_shared {
+                cmd.arg("-Clink-arg=-shared");
+                cmd.arg("-Clink-self-contained=n");
+                cmd.arg(&format!(
+                    "-Clinker={}/bin/clang",
+                    wasi_sdk_path.unwrap().display()
+                ));
+                cmd.arg("-L").arg(&compile.artifacts_dir);
+            }
             cmd
         };
 
+        let compile_cdylib = |cmd: &mut Command| {
+            cmd.arg("--crate-type=cdylib");
+            if runner.produces_component() {
+                if link_shared {
+                    cmd.arg("-Clink-arg=-Wl,--skip-wit-component");
+                } else {
+                    cmd.arg("-Clink-arg=--skip-wit-component");
+                }
+            }
+        };
+
         for file in config.externs.iter() {
             let file = manifest_dir.join(file);
             let stem = file.file_stem().unwrap().to_str().unwrap();
@@ -239,6 +280,16 @@ path = 'lib.rs'
             externs.push((stem.to_string(), output));
         }
 
+        for file in config.extern_dylibs.iter() {
+            let file = manifest_dir.join(file);
+            let stem = file.file_stem().unwrap().to_str().unwrap();
+            let output = compile.artifacts_dir.join(format!("lib{stem}.so"));
+            let mut cmd = rustc(&file, &output);
+            compile_cdylib(&mut cmd);
+            runner.run_command(&mut cmd)?;
+            dylibs.push(output);
+        }
+
         // Compile the main crate, passing `--extern` for all upstream crates.
         let mut cmd = rustc(&compile.component.path, &output);
         cmd.env(
@@ -252,15 +303,27 @@ path = 'lib.rs'
             let arg = format!("--extern={name}={}", path.display());
             cmd.arg(arg);
         }
-        cmd.arg("--crate-type=cdylib");
-        if runner.produces_component() {
-            cmd.arg("-Clink-arg=--skip-wit-component");
-        }
+        compile_cdylib(&mut cmd);
         runner.run_command(&mut cmd)?;
 
-        runner
-            .convert_p1_to_component(&output, compile)
-            .with_context(|| format!("failed to convert {output:?}"))?;
+        if link_shared {
+            let libc_so = wasi_sdk_path.unwrap().join(&format!(
+                "share/wasi-sysroot/lib/{}/libc.so",
+                runner.opts.rust.rust_target,
+            ));
+            if !libc_so.is_file() {
+                anyhow::bail!("libc.so not found at {libc_so:?}");
+            }
+            dylibs.insert(0, libc_so);
+            dylibs.push(output.clone());
+            runner
+                .link_dylibs_to_component(&dylibs, compile)
+                .with_context(|| format!("failed to link {output:?}"))?;
+        } else {
+            runner
+                .convert_p1_to_component(&output, compile)
+                .with_context(|| format!("failed to convert {output:?}"))?;
+        }
 
         Ok(())
     }

tests/runtime/rust/dylibs-async/runner.rs

@@ -0,0 +1,18 @@
+//@ [lang]
+//@ link_shared = true
+
+include!(env!("BINDINGS"));
+
+use crate::my::inline::a;
+
+struct Component;
+
+export!(Component);
+
+impl Guest for Component {
+    async fn run() {
+        let (tx, rx) = wit_future::new(|| 3u32);
+        drop(tx);
+        assert_eq!(a::a(rx).await, 3);
+    }
+}

tests/runtime/rust/dylibs-async/test.rs

@@ -0,0 +1,16 @@
+//@ [lang]
+//@ link_shared = true
+
+include!(env!("BINDINGS"));
+
+use crate::exports::my::inline::a;
+use wit_bindgen::FutureReader;
+
+struct Component;
+export!(Component);
+
+impl a::Guest for Component {
+    async fn a(f: FutureReader<u32>) -> u32 {
+        f.await
+    }
+}

tests/runtime/rust/dylibs-async/test.wit

@@ -0,0 +1,15 @@
+package my:inline;
+
+interface a {
+  a: async func(f: future<u32>) -> u32;
+}
+
+world test {
+    export a;
+}
+
+world runner {
+    import a;
+
+    export run: async func();
+}

tests/runtime/rust/dylibs/dylib_dep.rs

@@ -0,0 +1,8 @@
+wit_bindgen::generate!("runner-dep" in "test.wit");
+
+use crate::my::inline::b;
+
+#[unsafe(no_mangle)]
+pub extern "C" fn dylib_dep() {
+    b::b();
+}

tests/runtime/rust/dylibs/runner.rs

@@ -0,0 +1,24 @@
+//@ [lang]
+//@ extern_dylibs = ["dylib_dep.rs"]
+
+include!(env!("BINDINGS"));
+
+use crate::my::inline::a;
+
+struct Component;
+
+export!(Component);
+
+#[link(name = "dylib_dep")]
+unsafe extern "C" {
+    fn dylib_dep();
+}
+
+impl Guest for Component {
+    fn run() {
+        a::a();
+        unsafe {
+            dylib_dep();
+        }
+    }
+}

tests/runtime/rust/dylibs/test.rs

@@ -0,0 +1,17 @@
+//@ [lang]
+//@ link_shared = true
+
+include!(env!("BINDINGS"));
+
+use crate::exports::my::inline::{a, b};
+
+struct Component;
+export!(Component);
+
+impl a::Guest for Component {
+    fn a() {}
+}
+
+impl b::Guest for Component {
+    fn b() {}
+}

tests/runtime/rust/dylibs/test.wit

@@ -0,0 +1,24 @@
+package my:inline;
+
+interface a {
+  a: func();
+}
+
+interface b {
+  b: func();
+}
+
+world test {
+    export a;
+    export b;
+}
+
+world runner-dep {
+    import b;
+}
+
+world runner {
+    import a;
+
+    export run: func();
+}