From a2d6b0342c32cfe8dbb6b62f3e5393e0a7449da7 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Tue, 4 Mar 2025 09:38:14 -0600 Subject: [PATCH] rust: Store raw intrinsics in vtables (#1187) Instead of generating a wrapper-per-intrinsic try to instead store raw functions inside of vtables to push as much code as possible into `mod async_support`, as it's generally easier to grok code in one location rather than generated code. Note that this does not update the `read` or `write` vtable functions as they're going to be trickier, if even possible, to expose the more raw version. This shouldn't have any functional change, it's just a reorganizing. The `cancel_*` intrinsics are all marked with a `TODO`, however, to come back and reconsider the return value which is otherwise discarded right now. --- crates/guest-rust/rt/src/async_support.rs | 4 +- .../rt/src/async_support/future_support.rs | 55 +++- .../rt/src/async_support/stream_support.rs | 54 +++- crates/rust/src/interface.rs | 289 +++++------------- crates/rust/src/lib.rs | 28 +- 5 files changed, 162 insertions(+), 268 deletions(-) diff --git a/crates/guest-rust/rt/src/async_support.rs b/crates/guest-rust/rt/src/async_support.rs index 3e9979d24..b3d52f7da 100644 --- a/crates/guest-rust/rt/src/async_support.rs +++ b/crates/guest-rust/rt/src/async_support.rs @@ -25,8 +25,8 @@ mod future_support; mod stream_support; pub use { - future_support::{FutureReader, FutureVtable, FutureWriter}, - stream_support::{StreamReader, StreamVtable, StreamWriter}, + future_support::{future_new, FutureReader, FutureVtable, FutureWriter}, + stream_support::{stream_new, StreamReader, StreamVtable, StreamWriter}, }; pub use futures; diff --git a/crates/guest-rust/rt/src/async_support/future_support.rs b/crates/guest-rust/rt/src/async_support/future_support.rs index c7a17836b..aa6e12f63 100644 --- a/crates/guest-rust/rt/src/async_support/future_support.rs +++ b/crates/guest-rust/rt/src/async_support/future_support.rs @@ -22,10 +22,29 @@ use { pub struct FutureVtable { pub write: fn(future: u32, value: T) -> Pin>>, pub read: fn(future: u32) -> Pin>>>>, - pub cancel_write: fn(future: u32), - pub cancel_read: fn(future: u32), - pub close_writable: fn(future: u32, err_ctx: u32), - pub close_readable: fn(future: u32), + pub cancel_write: unsafe extern "C" fn(future: u32) -> u32, + pub cancel_read: unsafe extern "C" fn(future: u32) -> u32, + pub close_writable: unsafe extern "C" fn(future: u32, err_ctx: u32), + pub close_readable: unsafe extern "C" fn(future: u32), + pub new: unsafe extern "C" fn() -> u32, +} + +/// Helper function to create a new read/write pair for a component model +/// future. +pub unsafe fn future_new( + vtable: &'static FutureVtable, +) -> (FutureWriter, FutureReader) { + let handle = unsafe { (vtable.new)() }; + super::with_entry(handle, |entry| match entry { + Entry::Vacant(entry) => { + entry.insert(Handle::LocalOpen); + } + Entry::Occupied(_) => unreachable!(), + }); + ( + FutureWriter::new(handle, vtable), + FutureReader::new(handle, vtable), + ) } /// Represents the writable end of a Component Model `future`. @@ -84,7 +103,11 @@ impl CancelableWrite { Handle::LocalReady(..) => { entry.insert(Handle::LocalOpen); } - Handle::Write => (writer.vtable.cancel_write)(writer.handle), + Handle::Write => unsafe { + // TODO: spec-wise this can return `BLOCKED` which seems + // bad? + (writer.vtable.cancel_write)(writer.handle); + }, }, }); writer @@ -179,17 +202,17 @@ impl Drop for FutureWriter { entry.insert(Handle::LocalClosed); } Handle::Read => unreachable!(), - Handle::Write | Handle::LocalClosed => { + Handle::Write | Handle::LocalClosed => unsafe { entry.remove(); (self.vtable.close_writable)(self.handle, 0); - } + }, Handle::WriteClosedErr(_) => match entry.remove() { - Handle::WriteClosedErr(None) => { + Handle::WriteClosedErr(None) => unsafe { (self.vtable.close_writable)(self.handle, 0); - } - Handle::WriteClosedErr(Some(err_ctx)) => { + }, + Handle::WriteClosedErr(Some(err_ctx)) => unsafe { (self.vtable.close_writable)(self.handle, err_ctx.handle()); - } + }, _ => unreachable!(), }, }, @@ -240,7 +263,11 @@ impl CancelableRead { Handle::LocalWaiting(_) => { entry.insert(Handle::LocalOpen); } - Handle::Read => (reader.vtable.cancel_read)(handle), + Handle::Read => unsafe { + // TODO: spec-wise this can return `BLOCKED` which seems + // bad? + (reader.vtable.cancel_read)(handle); + }, }, }); reader @@ -389,10 +416,10 @@ impl Drop for FutureReader { Handle::LocalOpen | Handle::LocalWaiting(_) => { entry.insert(Handle::LocalClosed); } - Handle::Read | Handle::LocalClosed => { + Handle::Read | Handle::LocalClosed => unsafe { entry.remove(); (self.vtable.close_readable)(handle); - } + }, Handle::Write | Handle::WriteClosedErr(_) => unreachable!(), }, }); diff --git a/crates/guest-rust/rt/src/async_support/stream_support.rs b/crates/guest-rust/rt/src/async_support/stream_support.rs index 4e4e78055..3a52bd73a 100644 --- a/crates/guest-rust/rt/src/async_support/stream_support.rs +++ b/crates/guest-rust/rt/src/async_support/stream_support.rs @@ -35,12 +35,30 @@ pub struct StreamVtable { future: u32, values: &mut [MaybeUninit], ) -> Pin>> + '_>>, - pub cancel_write: fn(future: u32), - pub cancel_read: fn(future: u32), - pub close_writable: fn(future: u32, err_ctx: u32), - pub close_readable: fn(future: u32), + pub cancel_write: unsafe extern "C" fn(future: u32) -> u32, + pub cancel_read: unsafe extern "C" fn(future: u32) -> u32, + pub close_writable: unsafe extern "C" fn(future: u32, err_ctx: u32), + pub close_readable: unsafe extern "C" fn(future: u32), + pub new: unsafe extern "C" fn() -> u32, } +/// Helper function to create a new read/write pair for a component model +/// stream. +pub unsafe fn stream_new( + vtable: &'static StreamVtable, +) -> (StreamWriter, StreamReader) { + let handle = unsafe { (vtable.new)() }; + super::with_entry(handle, |entry| match entry { + Entry::Vacant(entry) => { + entry.insert(Handle::LocalOpen); + } + Entry::Occupied(_) => unreachable!(), + }); + ( + StreamWriter::new(handle, vtable), + StreamReader::new(handle, vtable), + ) +} struct CancelWriteOnDrop { handle: Option, vtable: &'static StreamVtable, @@ -60,7 +78,11 @@ impl Drop for CancelWriteOnDrop { Handle::LocalReady(..) => { entry.insert(Handle::LocalOpen); } - Handle::Write => (self.vtable.cancel_write)(handle), + Handle::Write => unsafe { + // TODO: spec-wise this can return `BLOCKED` which seems + // bad? + (self.vtable.cancel_write)(handle); + }, }, }); } @@ -221,18 +243,20 @@ impl Drop for StreamWriter { entry.insert(Handle::LocalClosed); } Handle::Read => unreachable!(), - Handle::Write | Handle::LocalClosed => { + Handle::Write | Handle::LocalClosed => unsafe { entry.remove(); (self.vtable.close_writable)(self.handle, 0); - } + }, Handle::WriteClosedErr(_) => match entry.remove() { // Care is taken to avoid dropping the ErrorContext before close_writable is called. // If the error context is dropped prematurely, the component may garbage collect // the error context before it can be used/referenced by close_writable(). - Handle::WriteClosedErr(Some(e)) => { + Handle::WriteClosedErr(Some(e)) => unsafe { (self.vtable.close_writable)(self.handle, e.handle) - } - Handle::WriteClosedErr(None) => (self.vtable.close_writable)(self.handle, 0), + }, + Handle::WriteClosedErr(None) => unsafe { + (self.vtable.close_writable)(self.handle, 0) + }, _ => unreachable!(), }, }, @@ -259,7 +283,11 @@ impl Drop for CancelReadOnDrop { Handle::LocalWaiting(_) => { entry.insert(Handle::LocalOpen); } - Handle::Read => (self.vtable.cancel_read)(handle), + Handle::Read => unsafe { + // TODO: spec-wise this can return `BLOCKED` which seems + // bad? + (self.vtable.cancel_read)(handle); + }, }, }); } @@ -457,10 +485,10 @@ impl Drop for StreamReader { Handle::LocalOpen | Handle::LocalWaiting(_) => { entry.insert(Handle::LocalClosed); } - Handle::Read | Handle::LocalClosed => { + Handle::Read | Handle::LocalClosed => unsafe { entry.remove(); (self.vtable.close_readable)(handle); - } + }, Handle::Write | Handle::WriteClosedErr(_) => unreachable!(), }, }); diff --git a/crates/rust/src/interface.rs b/crates/rust/src/interface.rs index 7aab82c1c..6b0b43ba2 100644 --- a/crates/rust/src/interface.rs +++ b/crates/rust/src/interface.rs @@ -549,19 +549,7 @@ pub mod vtable{ordinal} {{ let address = buffer.0.as_mut_ptr() as *mut u8; {lower} - #[cfg(not(target_arch = "wasm32"))] - unsafe extern "C" fn wit_import(_: u32, _: *mut u8) -> u32 {{ - unreachable!() - }} - - #[cfg(target_arch = "wasm32")] - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[async][future-write-{index}]{func_name}"] - fn wit_import(_: u32, _: *mut u8) -> u32; - }} - - match unsafe {{ {async_support}::await_future_result(wit_import, future, address).await }} {{ + match unsafe {{ {async_support}::await_future_result(start_write, future, address).await }} {{ {async_support}::AsyncWaitResult::Values(_) => true, {async_support}::AsyncWaitResult::End => false, {async_support}::AsyncWaitResult::Error(_) => unreachable!("received error while performing write"), @@ -575,19 +563,7 @@ pub mod vtable{ordinal} {{ let mut buffer = Buffer([::core::mem::MaybeUninit::uninit(); {size}]); let address = buffer.0.as_mut_ptr() as *mut u8; - #[cfg(not(target_arch = "wasm32"))] - unsafe extern "C" fn wit_import(_: u32, _: *mut u8) -> u32 {{ - unreachable!() - }} - - #[cfg(target_arch = "wasm32")] - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[async][future-read-{index}]{func_name}"] - fn wit_import(_: u32, _: *mut u8) -> u32; - }} - - match unsafe {{ {async_support}::await_future_result(wit_import, future, address).await }} {{ + match unsafe {{ {async_support}::await_future_result(start_read, future, address).await }} {{ {async_support}::AsyncWaitResult::Values(v) => {{ {lift} Some(Ok(value)) @@ -600,95 +576,49 @@ pub mod vtable{ordinal} {{ }}) }} - fn cancel_write(writer: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[future-cancel-write-{index}]{func_name}"] - fn cancel(_: u32) -> u32; - }} - unsafe {{ cancel(writer) }}; - }} - }} - - fn cancel_read(reader: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[future-cancel-read-{index}]{func_name}"] - fn cancel(_: u32) -> u32; - }} - unsafe {{ cancel(reader) }}; - }} - }} - - fn close_writable(writer: u32, err_ctx: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn cancel_write(_: u32) -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn cancel_read(_: u32) -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn close_writable(_: u32, _: u32) {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn close_readable(_: u32) {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn new() -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn start_read(_: u32, _: *mut u8) -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn start_write(_: u32, _: *mut u8) -> u32 {{ unreachable!() }} - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[future-close-writable-{index}]{func_name}"] - fn drop(_: u32, _: u32); - }} - unsafe {{ drop(writer, err_ctx) }} - }} + #[cfg(target_arch = "wasm32")] + #[link(wasm_import_module = "{module}")] + extern "C" {{ + #[link_name = "[future-new-{index}]{func_name}"] + fn new() -> u32; + #[link_name = "[future-cancel-write-{index}]{func_name}"] + fn cancel_write(_: u32) -> u32; + #[link_name = "[future-cancel-read-{index}]{func_name}"] + fn cancel_read(_: u32) -> u32; + #[link_name = "[future-close-writable-{index}]{func_name}"] + fn close_writable(_: u32, _: u32); + #[link_name = "[future-close-readable-{index}]{func_name}"] + fn close_readable(_: u32); + #[link_name = "[future-new-{index}]{func_name}"] + fn new() -> u32; + #[link_name = "[async][future-read-{index}]{func_name}"] + fn start_read(_: u32, _: *mut u8) -> u32; + #[link_name = "[async][future-write-{index}]{func_name}"] + fn start_write(_: u32, _: *mut u8) -> u32; }} - fn close_readable(reader: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[future-close-readable-{index}]{func_name}"] - fn drop(_: u32); - }} - unsafe {{ drop(reader) }} - }} - }} pub static VTABLE: {async_support}::FutureVtable<{name}> = {async_support}::FutureVtable::<{name}> {{ - write, read, cancel_write, cancel_read, close_writable, close_readable + write, read, cancel_write, cancel_read, close_writable, close_readable, new }}; impl super::FuturePayload for {name} {{ - fn new() -> (u32, &'static {async_support}::FutureVtable) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[future-new-{index}]{func_name}"] - fn new() -> u32; - }} - (unsafe {{ new() }}, &VTABLE) - }} - }} + const VTABLE: &'static {async_support}::FutureVtable = &VTABLE; }} }} "#, @@ -785,24 +715,12 @@ pub mod vtable{ordinal} {{ {lower_address} {lower} - #[cfg(not(target_arch = "wasm32"))] - unsafe extern "C" fn wit_import(_: u32, _: *mut u8, _: u32) -> u32 {{ - unreachable!() - }} - - #[cfg(target_arch = "wasm32")] - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[async][stream-write-{index}]{func_name}"] - fn wit_import(_: u32, _: *mut u8, _: u32) -> u32; - }} - let mut total = 0; while total < values.len() {{ match unsafe {{ {async_support}::await_stream_result( - wit_import, + start_write, stream, address.add(total * {size}), u32::try_from(values.len() - (total * {size})).unwrap() @@ -824,21 +742,9 @@ pub mod vtable{ordinal} {{ {box_}::pin(async move {{ {lift_address} - #[cfg(not(target_arch = "wasm32"))] - unsafe extern "C" fn wit_import(_: u32, _: *mut u8, _: u32) -> u32 {{ - unreachable!() - }} - - #[cfg(target_arch = "wasm32")] - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[async][stream-read-{index}]{func_name}"] - fn wit_import(_: u32, _: *mut u8, _: u32) -> u32; - }} - match unsafe {{ {async_support}::await_stream_result( - wit_import, + start_read, stream, address, u32::try_from(values.len()).unwrap() @@ -854,95 +760,48 @@ pub mod vtable{ordinal} {{ }}) }} - fn cancel_write(writer: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[stream-cancel-write-{index}]{func_name}"] - fn cancel(_: u32) -> u32; - }} - unsafe {{ cancel(writer) }}; - }} - }} - - fn cancel_read(reader: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[stream-cancel-read-{index}]{func_name}"] - fn cancel(_: u32) -> u32; - }} - unsafe {{ cancel(reader) }}; - }} - }} - - fn close_writable(writer: u32, err_ctx: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[stream-close-writable-{index}]{func_name}"] - fn drop(_: u32, _: u32); - }} - unsafe {{ drop(writer, err_ctx) }} - }} - }} - - fn close_readable(reader: u32) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn cancel_write(_: u32) -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn cancel_read(_: u32) -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn close_writable(_: u32, _: u32) {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn close_readable(_: u32) {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn new() -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn start_read(_: u32, _: *mut u8, _: u32) -> u32 {{ unreachable!() }} + #[cfg(not(target_arch = "wasm32"))] + unsafe extern "C" fn start_write(_: u32, _: *mut u8, _: u32) -> u32 {{ unreachable!() }} - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[stream-close-readable-{index}]{func_name}"] - fn drop(_: u32); - }} - unsafe {{ drop(reader) }} - }} + #[cfg(target_arch = "wasm32")] + #[link(wasm_import_module = "{module}")] + extern "C" {{ + #[link_name = "[stream-new-{index}]{func_name}"] + fn new() -> u32; + #[link_name = "[stream-cancel-write-{index}]{func_name}"] + fn cancel_write(_: u32) -> u32; + #[link_name = "[stream-cancel-read-{index}]{func_name}"] + fn cancel_read(_: u32) -> u32; + #[link_name = "[stream-close-writable-{index}]{func_name}"] + fn close_writable(_: u32, _: u32); + #[link_name = "[stream-close-readable-{index}]{func_name}"] + fn close_readable(_: u32); + #[link_name = "[stream-new-{index}]{func_name}"] + fn new() -> u32; + #[link_name = "[async][stream-read-{index}]{func_name}"] + fn start_read(_: u32, _: *mut u8, _: u32) -> u32; + #[link_name = "[async][stream-write-{index}]{func_name}"] + fn start_write(_: u32, _: *mut u8, _: u32) -> u32; }} pub static VTABLE: {async_support}::StreamVtable<{name}> = {async_support}::StreamVtable::<{name}> {{ - write, read, cancel_write, cancel_read, close_writable, close_readable + write, read, cancel_write, cancel_read, close_writable, close_readable, new }}; impl super::StreamPayload for {name} {{ - fn new() -> (u32, &'static {async_support}::StreamVtable) {{ - #[cfg(not(target_arch = "wasm32"))] - {{ - unreachable!(); - }} - - #[cfg(target_arch = "wasm32")] - {{ - #[link(wasm_import_module = "{module}")] - extern "C" {{ - #[link_name = "[stream-new-{index}]{func_name}"] - fn new() -> u32; - }} - (unsafe {{ new() }}, &VTABLE) - }} - }} + const VTABLE: &'static {async_support}::StreamVtable = &VTABLE; }} }} "#, diff --git a/crates/rust/src/lib.rs b/crates/rust/src/lib.rs index 318e927a5..1ca3939d6 100644 --- a/crates/rust/src/lib.rs +++ b/crates/rust/src/lib.rs @@ -479,7 +479,7 @@ pub mod wit_future {{ #[doc(hidden)] pub trait FuturePayload: Unpin + Sized + 'static {{ - fn new() -> (u32, &'static {async_support}::FutureVtable); + const VTABLE: &'static {async_support}::FutureVtable; }}" )); for code in self.future_payloads.values() { @@ -489,17 +489,7 @@ pub mod wit_future {{ "\ /// Creates a new Component Model `future` with the specified payload type. pub fn new() -> ({async_support}::FutureWriter, {async_support}::FutureReader) {{ - let (handle, vtable) = T::new(); - {async_support}::with_entry(handle, |entry| match entry {{ - ::std::collections::hash_map::Entry::Vacant(entry) => {{ - entry.insert({async_support}::Handle::LocalOpen); - }} - ::std::collections::hash_map::Entry::Occupied(_) => unreachable!(), - }}); - ( - {async_support}::FutureWriter::new(handle, vtable), - {async_support}::FutureReader::new(handle, vtable), - ) + unsafe {{ {async_support}::future_new::(T::VTABLE) }} }} }} ", @@ -514,7 +504,7 @@ pub mod wit_stream {{ #![allow(dead_code, unused_variables, clippy::all)] pub trait StreamPayload: Unpin + Sized + 'static {{ - fn new() -> (u32, &'static {async_support}::StreamVtable); + const VTABLE: &'static {async_support}::StreamVtable; }}" )); for code in self.stream_payloads.values() { @@ -524,17 +514,7 @@ pub mod wit_stream {{ &format!("\ /// Creates a new Component Model `stream` with the specified payload type. pub fn new() -> ({async_support}::StreamWriter, {async_support}::StreamReader) {{ - let (handle, vtable) = T::new(); - {async_support}::with_entry(handle, |entry| match entry {{ - ::std::collections::hash_map::Entry::Vacant(entry) => {{ - entry.insert({async_support}::Handle::LocalOpen); - }} - ::std::collections::hash_map::Entry::Occupied(_) => unreachable!(), - }}); - ( - {async_support}::StreamWriter::new(handle, vtable), - {async_support}::StreamReader::new(handle, vtable), - ) + unsafe {{ {async_support}::stream_new::(T::VTABLE) }} }} }} "),