Skip to content

Commit daf88e2

Browse files
BjornTheProgrammerBjornTheProgrammer
committed
feat(rust) make async work with native
1 parent a19207a commit daf88e2

13 files changed

Lines changed: 782 additions & 66 deletions

File tree

‎.gitignore‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,7 @@ crates/guest-rust/src/cabi_realloc.o
1111
wit_component
1212

1313
/wit-bindgen.sln
14+
15+
# Built by crates/rust/tests/native_e2e.rs
16+
crates/rust/tests/native-e2e/Cargo.lock
17+
crates/rust/tests/native-e2e/target

‎Cargo.lock‎

Lines changed: 11 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/guest-rust/src/rt/async_support.rs‎

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,33 +27,49 @@ macro_rules! rtdebug {
2727

2828
/// Helper macro to deduplicate foreign definitions of wasm functions.
2929
///
30-
/// This automatically imports when on wasm targets and then defines a dummy
31-
/// panicking shim for native targets to support native compilation but fail at
32-
/// runtime.
30+
/// On wasm targets this declares the canonical ABI built-ins as ordinary
31+
/// linker-resolved imports. On native targets each one instead becomes a shim
32+
/// that asks the host's import resolver for its implementation on first call
33+
/// (see `rt::native_imports`), identified by the same module and name, so the
34+
/// generated code is the same on both targets and only who satisfies the
35+
/// import differs.
3336
macro_rules! extern_wasm {
3437
(
35-
$(#[$extern_attr:meta])*
38+
#[link(wasm_import_module = $module:literal)]
3639
unsafe extern "C" {
3740
$(
38-
$(#[$func_attr:meta])*
39-
$vis:vis fn $func_name:ident ( $($args:tt)* ) $(-> $ret:ty)?;
41+
#[link_name = $name:literal]
42+
$vis:vis fn $func_name:ident ( $($arg:ident : $ty:ty),* $(,)? ) $(-> $ret:ty)?;
4043
)*
4144
}
4245
) => {
4346
$(
4447
#[cfg(not(target_family = "wasm"))]
45-
#[allow(unused, reason = "dummy shim for non-wasm compilation, never invoked")]
46-
$vis unsafe fn $func_name($($args)*) $(-> $ret)? {
47-
unreachable!();
48+
#[allow(dead_code, reason = "mirrors the wasm import set even if unused natively")]
49+
$vis unsafe fn $func_name($($arg: $ty),*) $(-> $ret)? {
50+
static CACHE: ::core::sync::atomic::AtomicPtr<()> =
51+
::core::sync::atomic::AtomicPtr::new(::core::ptr::null_mut());
52+
// Named so as not to shadow any parameter.
53+
let mut __impl = CACHE.load(::core::sync::atomic::Ordering::Acquire);
54+
if __impl.is_null() {
55+
__impl = crate::rt::resolve_import(
56+
crate::rt::native_imports::cstr(concat!($module, "\0")),
57+
crate::rt::native_imports::cstr(concat!($name, "\0")),
58+
);
59+
CACHE.store(__impl, ::core::sync::atomic::Ordering::Release);
60+
}
61+
let __func: unsafe extern "C" fn($($ty),*) $(-> $ret)? =
62+
unsafe { ::core::mem::transmute(__impl) };
63+
unsafe { __func($($arg),*) }
4864
}
4965
)*
5066

5167
#[cfg(target_family = "wasm")]
52-
$(#[$extern_attr])*
68+
#[link(wasm_import_module = $module)]
5369
unsafe extern "C" {
5470
$(
55-
$(#[$func_attr])*
56-
$vis fn $func_name($($args)*) $(-> $ret)?;
71+
#[link_name = $name]
72+
$vis fn $func_name($($arg: $ty),*) $(-> $ret)?;
5773
)*
5874
}
5975
};

‎crates/guest-rust/src/rt/async_support/cabi.rs‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -69,21 +69,42 @@
6969
7070
use core::ffi::c_void;
7171

72-
extern_wasm! {
73-
unsafe extern "C" {
74-
/// Sets the global task pointer to `ptr` provided. Returns the previous
75-
/// value.
76-
///
77-
/// This function acts as a dual getter and a setter. To get the
78-
/// current task pointer a dummy `ptr` can be provided (e.g. NULL) and then
79-
/// it's passed back when you're done working with it. When setting the
80-
/// current task pointer it's recommended to call this and then call it
81-
/// again with the previous value when the tasks's work is done.
82-
///
83-
/// For executors they need to ensure that the `ptr` passed in lives for
84-
/// the entire lifetime of the component model task.
85-
pub fn wasip3_task_set(ptr: *mut wasip3_task) -> *mut wasip3_task;
86-
}
72+
#[cfg(target_family = "wasm")]
73+
unsafe extern "C" {
74+
/// Sets the global task pointer to `ptr` provided. Returns the previous
75+
/// value.
76+
///
77+
/// This function acts as a dual getter and a setter. To get the
78+
/// current task pointer a dummy `ptr` can be provided (e.g. NULL) and then
79+
/// it's passed back when you're done working with it. When setting the
80+
/// current task pointer it's recommended to call this and then call it
81+
/// again with the previous value when the tasks's work is done.
82+
///
83+
/// For executors they need to ensure that the `ptr` passed in lives for
84+
/// the entire lifetime of the component model task.
85+
pub fn wasip3_task_set(ptr: *mut wasip3_task) -> *mut wasip3_task;
86+
}
87+
88+
/// Native counterpart of the C-defined `wasip3_task_set` above. Uses
89+
/// thread local so there is no possible panic on multi thread when polling
90+
/// two async exports at the same time.
91+
#[cfg(all(not(target_family = "wasm"), feature = "std"))]
92+
pub unsafe fn wasip3_task_set(ptr: *mut wasip3_task) -> *mut wasip3_task {
93+
use core::cell::Cell;
94+
std::thread_local!(
95+
static CURRENT: Cell<*mut wasip3_task> = const { Cell::new(core::ptr::null_mut()) }
96+
);
97+
CURRENT.with(|current| current.replace(ptr))
98+
}
99+
100+
/// Without `std` there are no thread-locals on stable Rust, so this falls
101+
/// back to one global slot, which requires the host to not poll two async
102+
/// exports at the same time.
103+
#[cfg(all(not(target_family = "wasm"), not(feature = "std")))]
104+
pub unsafe fn wasip3_task_set(ptr: *mut wasip3_task) -> *mut wasip3_task {
105+
use core::sync::atomic::{AtomicPtr, Ordering};
106+
static CURRENT: AtomicPtr<wasip3_task> = AtomicPtr::new(core::ptr::null_mut());
107+
CURRENT.swap(ptr, Ordering::AcqRel)
87108
}
88109

89110
/// The first version of `wasip3_task` which implies the existence of the

‎crates/guest-rust/src/rt/async_support/error_context.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,10 +74,10 @@ extern_wasm! {
7474
#[link(wasm_import_module = "$root")]
7575
unsafe extern "C" {
7676
#[link_name = "[error-context-new-utf8]"]
77-
fn new(_: *const u8, _: usize) -> u32;
77+
fn new(ptr: *const u8, len: usize) -> u32;
7878
#[link_name = "[error-context-drop]"]
79-
fn drop(_: u32);
79+
fn drop(handle: u32);
8080
#[link_name = "[error-context-debug-message-utf8]"]
81-
fn debug_message(_: u32, _: &mut RetPtr);
81+
fn debug_message(handle: u32, ret: &mut RetPtr);
8282
}
8383
}

‎crates/guest-rust/src/rt/async_support/waitable_set.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,8 +75,8 @@ extern_wasm! {
7575
#[link_name = "[waitable-join]"]
7676
fn join(waitable: u32, set: u32);
7777
#[link_name = "[waitable-set-wait]"]
78-
fn wait(_: u32, _: *mut [u32; 2]) -> u32;
78+
fn wait(set: u32, event: *mut [u32; 2]) -> u32;
7979
#[link_name = "[waitable-set-poll]"]
80-
fn poll(_: u32, _: *mut [u32; 2]) -> u32;
80+
fn poll(set: u32, event: *mut [u32; 2]) -> u32;
8181
}
8282
}

‎crates/guest-rust/src/rt/native_imports.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,15 @@ pub unsafe extern "C" fn __wit_bindgen_set_import_resolver(
4848
RESOLVER.store(resolver, Ordering::Release);
4949
}
5050

51+
/// Builds the `&CStr` for an import name from a `concat!(name, "\0")`
52+
/// literal, for the runtime's own intrinsic shims.
53+
pub(crate) const fn cstr(with_nul: &'static str) -> &'static CStr {
54+
match CStr::from_bytes_with_nul(with_nul.as_bytes()) {
55+
Ok(s) => s,
56+
Err(_) => panic!("import name contains an interior NUL"),
57+
}
58+
}
59+
5160
/// Called by generated import shims on their first invocation.
5261
///
5362
/// `module` and `name` are the strings described on [`ImportResolver`].

‎crates/rust/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ test-helpers = { path = '../test-helpers' }
4141
# For use with the custom attributes test
4242
serde_json = { workspace = true }
4343
bytes = "1"
44+
# For the native end-to-end test
45+
libloading = "0.8"
4446

4547
[features]
4648
serde = ['dep:serde', 'wit-bindgen-core/serde']

‎crates/rust/src/interface.rs‎

Lines changed: 67 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -725,26 +725,79 @@ macro_rules! {macro_name} {{
725725
}
726726
}
727727

728+
// Natively the intrinsics are resolved through the host's import
729+
// resolver, exactly like every other import.
730+
let rt = self.r#gen.runtime_path();
731+
let handle = || ("handle".to_string(), "u32");
732+
let mut extra: Vec<(String, &str)> = Vec::new();
733+
if let PayloadFor::Stream = payload_for {
734+
extra.push(("amt".to_string(), "usize"));
735+
}
736+
let mut native_intrinsics = String::new();
737+
for (rust_name, core_name, params, result) in [
738+
(
739+
"new",
740+
format!("[{import_prefix}-new-{index}]{func_name}"),
741+
vec![],
742+
Some("u64"),
743+
),
744+
(
745+
"cancel_write",
746+
format!("[{import_prefix}-cancel-write-{index}]{func_name}"),
747+
vec![handle()],
748+
Some("u32"),
749+
),
750+
(
751+
"cancel_read",
752+
format!("[{import_prefix}-cancel-read-{index}]{func_name}"),
753+
vec![handle()],
754+
Some("u32"),
755+
),
756+
(
757+
"drop_writable",
758+
format!("[{import_prefix}-drop-writable-{index}]{func_name}"),
759+
vec![handle()],
760+
None,
761+
),
762+
(
763+
"drop_readable",
764+
format!("[{import_prefix}-drop-readable-{index}]{func_name}"),
765+
vec![handle()],
766+
None,
767+
),
768+
(
769+
"start_read",
770+
format!("[async-lower][{import_prefix}-read-{index}]{func_name}"),
771+
[
772+
vec![handle(), ("ptr".to_string(), "*mut u8")],
773+
extra.clone(),
774+
]
775+
.concat(),
776+
Some("u32"),
777+
),
778+
(
779+
"start_write",
780+
format!("[async-lower][{import_prefix}-write-{index}]{func_name}"),
781+
[
782+
vec![handle(), ("ptr".to_string(), "*const u8")],
783+
extra.clone(),
784+
]
785+
.concat(),
786+
Some("u32"),
787+
),
788+
] {
789+
native_intrinsics.push_str(&crate::native_import_shim(
790+
&module, &core_name, rust_name, &params, result, rt,
791+
));
792+
}
793+
728794
let code = format!(
729795
r#"
730796
#[doc(hidden)]
731797
#[allow(unused_unsafe)]
732798
pub mod vtable{ordinal} {{
733799
734-
#[cfg(not(target_arch = "wasm32"))]
735-
unsafe extern "C" fn cancel_write(_: u32) -> u32 {{ unreachable!() }}
736-
#[cfg(not(target_arch = "wasm32"))]
737-
unsafe extern "C" fn cancel_read(_: u32) -> u32 {{ unreachable!() }}
738-
#[cfg(not(target_arch = "wasm32"))]
739-
unsafe extern "C" fn drop_writable(_: u32) {{ unreachable!() }}
740-
#[cfg(not(target_arch = "wasm32"))]
741-
unsafe extern "C" fn drop_readable(_: u32) {{ unreachable!() }}
742-
#[cfg(not(target_arch = "wasm32"))]
743-
unsafe extern "C" fn new() -> u64 {{ unreachable!() }}
744-
#[cfg(not(target_arch = "wasm32"))]
745-
unsafe extern "C" fn start_read(_: u32, _: *mut u8{start_extra}) -> u32 {{ unreachable!() }}
746-
#[cfg(not(target_arch = "wasm32"))]
747-
unsafe extern "C" fn start_write(_: u32, _: *const u8{start_extra}) -> u32 {{ unreachable!() }}
800+
{native_intrinsics}
748801
749802
// Work around a behavior of LLD where in a shared library when an address
750803
// is taken of an imported function that only shows up as a `GOT.func`

‎crates/rust/src/lib.rs‎

Lines changed: 54 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1906,8 +1906,7 @@ fn wasm_type(ty: WasmType) -> &'static str {
19061906

19071907
/// Declares the core import `wasm_import_module`/`wasm_import_name` as a
19081908
/// function named `rust_name`. On `wasm32` this is a plain linker-resolved
1909-
/// import; natively it's a shim that asks the host's resolver for the
1910-
/// implementation on first call (see `rt::resolve_import`) and caches it.
1909+
/// import; natively it's the shim from [`native_import_shim`].
19111910
fn declare_import(
19121911
wasm_import_module: &str,
19131912
wasm_import_name: &str,
@@ -1929,20 +1928,19 @@ fn declare_import(
19291928
sig.push_str(wasm_type(*result));
19301929
}
19311930

1932-
let named_params: Vec<String> = params
1931+
let named_params: Vec<(String, &str)> = params
19331932
.iter()
19341933
.enumerate()
1935-
.map(|(i, ty)| format!("arg{i}: {}", wasm_type(*ty)))
1934+
.map(|(i, ty)| (format!("arg{i}"), wasm_type(*ty)))
19361935
.collect();
1937-
let ret_sig = results
1938-
.first()
1939-
.map(|r| format!(" -> {}", wasm_type(*r)))
1940-
.unwrap_or_default();
1941-
let call_args = (0..params.len())
1942-
.map(|i| format!("arg{i}"))
1943-
.collect::<Vec<_>>()
1944-
.join(", ");
1945-
let named_params_str = named_params.join(", ");
1936+
let native = native_import_shim(
1937+
wasm_import_module,
1938+
wasm_import_name,
1939+
rust_name,
1940+
&named_params,
1941+
results.first().map(|r| wasm_type(*r)),
1942+
rt,
1943+
);
19461944

19471945
format!(
19481946
r#"
@@ -1952,18 +1950,54 @@ fn declare_import(
19521950
#[link_name = "{wasm_import_name}"]
19531951
fn {rust_name}{sig};
19541952
}}
1953+
{native}
1954+
"#,
1955+
)
1956+
}
1957+
1958+
/// Emits the native (non-`wasm32`) definition of the core import
1959+
/// `module`/`name`: an `unsafe extern "C" fn` named `rust_name` that asks the
1960+
/// host's import resolver for the implementation on first call (see
1961+
/// `rt::resolve_import`) and caches it. `params` are `(name, type)` pairs.
1962+
fn native_import_shim(
1963+
module: &str,
1964+
name: &str,
1965+
rust_name: &str,
1966+
params: &[(String, &str)],
1967+
result: Option<&str>,
1968+
rt: &str,
1969+
) -> String {
1970+
let named_params = params
1971+
.iter()
1972+
.map(|(arg, ty)| format!("{arg}: {ty}"))
1973+
.collect::<Vec<_>>()
1974+
.join(", ");
1975+
let param_types = params
1976+
.iter()
1977+
.map(|(_, ty)| *ty)
1978+
.collect::<Vec<_>>()
1979+
.join(", ");
1980+
let call_args = params
1981+
.iter()
1982+
.map(|(arg, _)| arg.as_str())
1983+
.collect::<Vec<_>>()
1984+
.join(", ");
1985+
let ret_sig = result.map(|r| format!(" -> {r}")).unwrap_or_default();
19551986

1987+
format!(
1988+
r#"
19561989
#[cfg(not(target_arch = "wasm32"))]
1957-
unsafe extern "C" fn {rust_name}({named_params_str}){ret_sig} {{
1990+
unsafe extern "C" fn {rust_name}({named_params}){ret_sig} {{
19581991
static CACHE: ::core::sync::atomic::AtomicPtr<()> =
19591992
::core::sync::atomic::AtomicPtr::new(::core::ptr::null_mut());
1960-
let mut ptr = CACHE.load(::core::sync::atomic::Ordering::Acquire);
1961-
if ptr.is_null() {{
1962-
ptr = {rt}::resolve_import(c"{wasm_import_module}", c"{wasm_import_name}");
1963-
CACHE.store(ptr, ::core::sync::atomic::Ordering::Release);
1993+
// Named so as not to shadow any parameter.
1994+
let mut __impl = CACHE.load(::core::sync::atomic::Ordering::Acquire);
1995+
if __impl.is_null() {{
1996+
__impl = {rt}::resolve_import(c"{module}", c"{name}");
1997+
CACHE.store(__impl, ::core::sync::atomic::Ordering::Release);
19641998
}}
1965-
let f: unsafe extern "C" fn{sig} = unsafe {{ ::core::mem::transmute(ptr) }};
1966-
unsafe {{ f({call_args}) }}
1999+
let __func: unsafe extern "C" fn({param_types}){ret_sig} = unsafe {{ ::core::mem::transmute(__impl) }};
2000+
unsafe {{ __func({call_args}) }}
19672001
}}
19682002
"#,
19692003
)

0 commit comments

Comments
 (0)