flesh out async support in wit-dylib

This does two things:

First, it adds a low-level API to `wit-dylib-ffi` for async imports and exports
as an alternative to the high-level API based on `wit-bindgen`'s `rt` library
and enabled with the `async` feature.  Thus enabling the `async` feature now
really just means "use the high-level API", while disabling it means "use the
low-level API".  The low-level API can make more sense for e.g. Python, where
the coroutine model is different enough from Rust's that it's easier to map it
directly to the component model ABI rather than try to plumb it through Rust's
async APIs.

Second, it populates the `wit_future` and `wit_stream` types with all the
relevant intrinsics, ABI size/align metadata, and generated payload lift and
lower functions.

I've tested this all end-to-end with `componentize-py`, but could use some
guidance on how to add tests here as well.

Signed-off-by: Joel Dice <[email protected]>

Author: dicej

crates/wit-dylib/ffi/src/ffi.rs

@@ -29,7 +29,7 @@ pub const WIT_TYPE_FUTURE: u32 = 25;
 pub const WIT_TYPE_STREAM: u32 = 26;
 pub const WIT_TYPE_ALIAS: u32 = 27;
 pub const WIT_TYPE_EMPTY: u32 = 255;
-pub const WIT_CURRENT_VERSION: u32 = 1;
+pub const WIT_CURRENT_VERSION: u32 = 2;
 pub type wit_type_t = u32;
 pub type wit_import_fn_t =
     ::std::option::Option<unsafe extern "C" fn(cx: *mut ::std::os::raw::c_void)>;
@@ -176,20 +176,74 @@ pub struct wit_fixed_size_list {
     pub ty: wit_type_t,
 }
 pub type wit_fixed_size_list_t = wit_fixed_size_list;
+pub type wit_lift_fn_t = ::std::option::Option<
+    unsafe extern "C" fn(cx: *mut ::std::os::raw::c_void, buffer: *const ::std::os::raw::c_void),
+>;
+pub type wit_lower_fn_t = ::std::option::Option<
+    unsafe extern "C" fn(cx: *mut ::std::os::raw::c_void, buffer: *mut ::std::os::raw::c_void),
+>;
+pub type wit_future_new_fn_t = ::std::option::Option<unsafe extern "C" fn() -> u64>;
+pub type wit_future_read_fn_t = ::std::option::Option<
+    unsafe extern "C" fn(future: u32, buffer: *mut ::std::os::raw::c_void) -> u32,
+>;
+pub type wit_future_write_fn_t = ::std::option::Option<
+    unsafe extern "C" fn(future: u32, buffer: *const ::std::os::raw::c_void) -> u32,
+>;
+pub type wit_future_cancel_read_fn_t =
+    ::std::option::Option<unsafe extern "C" fn(future: u32) -> u32>;
+pub type wit_future_cancel_write_fn_t =
+    ::std::option::Option<unsafe extern "C" fn(future: u32) -> u32>;
+pub type wit_future_drop_readable_fn_t = ::std::option::Option<unsafe extern "C" fn(future: u32)>;
+pub type wit_future_drop_writable_fn_t = ::std::option::Option<unsafe extern "C" fn(future: u32)>;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct wit_future {
     pub interface: *const ::std::os::raw::c_char,
     pub name: *const ::std::os::raw::c_char,
     pub ty: wit_type_t,
+    pub new: wit_future_new_fn_t,
+    pub read: wit_future_read_fn_t,
+    pub write: wit_future_write_fn_t,
+    pub cancel_read: wit_future_cancel_read_fn_t,
+    pub cancel_write: wit_future_cancel_write_fn_t,
+    pub drop_readable: wit_future_drop_readable_fn_t,
+    pub drop_writable: wit_future_drop_writable_fn_t,
+    pub lift: wit_lift_fn_t,
+    pub lower: wit_lower_fn_t,
+    pub abi_payload_size: usize,
+    pub abi_payload_align: usize,
 }
 pub type wit_future_t = wit_future;
+pub type wit_stream_new_fn_t = ::std::option::Option<unsafe extern "C" fn() -> u64>;
+pub type wit_stream_read_fn_t = ::std::option::Option<
+    unsafe extern "C" fn(stream: u32, buffer: *mut ::std::os::raw::c_void, count: usize) -> u32,
+>;
+pub type wit_stream_write_fn_t = ::std::option::Option<
+    unsafe extern "C" fn(stream: u32, buffer: *const ::std::os::raw::c_void, count: usize) -> u32,
+>;
+pub type wit_stream_cancel_read_fn_t =
+    ::std::option::Option<unsafe extern "C" fn(stream: u32) -> u32>;
+pub type wit_stream_cancel_write_fn_t =
+    ::std::option::Option<unsafe extern "C" fn(stream: u32) -> u32>;
+pub type wit_stream_drop_writable_fn_t = ::std::option::Option<unsafe extern "C" fn(stream: u32)>;
+pub type wit_stream_drop_readable_fn_t = ::std::option::Option<unsafe extern "C" fn(stream: u32)>;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct wit_stream {
     pub interface: *const ::std::os::raw::c_char,
     pub name: *const ::std::os::raw::c_char,
     pub ty: wit_type_t,
+    pub new: wit_stream_new_fn_t,
+    pub read: wit_stream_read_fn_t,
+    pub write: wit_stream_write_fn_t,
+    pub cancel_read: wit_stream_cancel_read_fn_t,
+    pub cancel_write: wit_stream_cancel_write_fn_t,
+    pub drop_readable: wit_stream_drop_readable_fn_t,
+    pub drop_writable: wit_stream_drop_writable_fn_t,
+    pub lift: wit_lift_fn_t,
+    pub lower: wit_lower_fn_t,
+    pub abi_payload_size: usize,
+    pub abi_payload_align: usize,
 }
 pub type wit_stream_t = wit_stream;
 #[repr(C)]

crates/wit-dylib/ffi/src/lib.rs

@@ -372,12 +372,19 @@ pub trait Interpreter: 'static {
 
     fn export_call(wit: Wit, func: ExportFunction, cx: &mut Self::CallCx<'_>);
 
+    #[cfg(feature = "async")]
     fn export_call_async(
         wit: Wit,
         func: ExportFunction,
         cx: Box<Self::CallCx<'static>>,
     ) -> impl std::future::Future<Output = ()>;
 
+    #[cfg(not(feature = "async"))]
+    fn export_async_start(wit: Wit, func: ExportFunction, cx: Box<Self::CallCx<'static>>) -> u32;
+
+    #[cfg(not(feature = "async"))]
+    fn export_async_callback(event0: u32, event1: u32, event2: u32) -> u32;
+
     fn export_finish(cx: Box<Self::CallCx<'_>>, func: ExportFunction) {
         let _ = func;
         let _ = cx;
@@ -519,20 +526,22 @@ pub trait RawInterpreter: Interpreter {
 
     unsafe fn raw_export_async_call(cx: *mut u8, which: usize) -> u32 {
         debug_println!("export_async_call({cx:?}, {which})");
-        #[cfg(feature = "async")]
         unsafe {
             let wit = Wit::from_raw(WIT_T);
             let func = wit.export_func(which);
-            wit_bindgen::rt::async_support::start_task(Self::export_call_async(
-                wit,
-                func,
-                Box::from_raw(cx.cast()),
-            ))
-            .cast_unsigned()
-        }
-        #[cfg(not(feature = "async"))]
-        unsafe {
-            panic!("async support disabled at compile time");
+            #[cfg(feature = "async")]
+            {
+                wit_bindgen::rt::async_support::start_task(Self::export_call_async(
+                    wit,
+                    func,
+                    Box::from_raw(cx.cast()),
+                ))
+                .cast_unsigned()
+            }
+            #[cfg(not(feature = "async"))]
+            {
+                Self::export_async_start(wit, func, Box::from_raw(cx.cast()))
+            }
         }
     }
 
@@ -543,7 +552,7 @@ pub trait RawInterpreter: Interpreter {
             wit_bindgen::rt::async_support::callback(a, b, c)
         }
         #[cfg(not(feature = "async"))]
-        panic!("async support disabled at compile time");
+        Self::export_async_callback(a, b, c)
     }
 
     unsafe fn raw_export_finish(cx: *mut u8, which: usize) {

crates/wit-dylib/ffi/src/types.rs

@@ -1,6 +1,6 @@
 use crate::Call;
 use crate::ffi;
-use std::ffi::{CStr, c_char};
+use std::ffi::{CStr, c_char, c_void};
 use std::fmt;
 use std::hash::{Hash, Hasher};
 
@@ -329,6 +329,11 @@ unsafe fn slice<T>(ptr: *const T, len: usize) -> &'static [T] {
     }
 }
 
+pub struct PendingAsyncImportCall {
+    pub subtask: u32,
+    pub buffer: *mut u8,
+}
+
 #[derive(Copy, Clone)]
 pub struct ImportFunction {
     wit: Wit,
@@ -358,6 +363,10 @@ impl ImportFunction {
         self.ptr.async_impl
     }
 
+    pub fn async_import_lift_impl(&self) -> ffi::wit_import_async_lift_fn_t {
+        self.ptr.async_lift_impl
+    }
+
     pub fn is_async(&self) -> bool {
         self.ptr.async_impl.is_some()
     }
@@ -384,6 +393,40 @@ impl ImportFunction {
         }
     }
 
+    #[cfg(not(feature = "async"))]
+    pub unsafe fn call_import_async(&self, cx: &mut impl Call) -> Option<PendingAsyncImportCall> {
+        use core::alloc::Layout;
+
+        const STATUS_STARTING: u32 = 0;
+        const STATUS_STARTED: u32 = 1;
+        const STATUS_RETURNED: u32 = 2;
+
+        let layout =
+            Layout::from_size_align(self.ptr.async_abi_area_size, self.ptr.async_abi_area_align)
+                .unwrap();
+        let (status, buffer) = unsafe {
+            let buffer = std::alloc::alloc(layout);
+            cx.defer_deallocate(buffer, layout);
+            let status = self.ptr.async_impl.unwrap()((&raw mut *cx).cast(), buffer.cast());
+            (status, buffer)
+        };
+        let subtask = status >> 4;
+        let status = status & 0xF;
+        match status {
+            STATUS_STARTING | STATUS_STARTED => Some(PendingAsyncImportCall { subtask, buffer }),
+            STATUS_RETURNED => unsafe {
+                self.lift_import_async_result(cx, buffer);
+                None
+            },
+            _ => panic!("unexpected status from async import call: {status}"),
+        }
+    }
+
+    #[cfg(not(feature = "async"))]
+    pub unsafe fn lift_import_async_result(&self, cx: &mut impl Call, buffer: *mut u8) {
+        unsafe { self.ptr.async_lift_impl.unwrap()((&raw mut *cx).cast(), buffer.cast()) };
+    }
+
     #[cfg(feature = "async")]
     pub async fn call_import_async(&self, cx: &mut impl Call) {
         use core::alloc::Layout;
@@ -467,6 +510,14 @@ impl ExportFunction {
         self.ptr.task_return
     }
 
+    pub fn call_task_return(&self, cx: &mut impl Call) {
+        let task_return = self.task_return().unwrap();
+        unsafe {
+            let cx: *mut _ = cx;
+            task_return(cx.cast());
+        }
+    }
+
     pub fn params(&self) -> impl ExactSizeIterator<Item = Type> + DoubleEndedIterator + Clone + '_ {
         self.raw_params()
             .iter()
@@ -1013,6 +1064,50 @@ impl Future {
     pub fn ty(&self) -> Option<Type> {
         Type::from_raw_opt(self.wit, self.ptr.ty)
     }
+
+    pub fn new(&self) -> unsafe extern "C" fn() -> u64 {
+        self.ptr.new.unwrap()
+    }
+
+    pub fn read(&self) -> unsafe extern "C" fn(u32, *mut c_void) -> u32 {
+        self.ptr.read.unwrap()
+    }
+
+    pub fn write(&self) -> unsafe extern "C" fn(u32, *const c_void) -> u32 {
+        self.ptr.write.unwrap()
+    }
+
+    pub fn cancel_read(&self) -> unsafe extern "C" fn(u32) -> u32 {
+        self.ptr.cancel_read.unwrap()
+    }
+
+    pub fn cancel_write(&self) -> unsafe extern "C" fn(u32) -> u32 {
+        self.ptr.cancel_write.unwrap()
+    }
+
+    pub fn drop_readable(&self) -> unsafe extern "C" fn(u32) {
+        self.ptr.drop_readable.unwrap()
+    }
+
+    pub fn drop_writable(&self) -> unsafe extern "C" fn(u32) {
+        self.ptr.drop_writable.unwrap()
+    }
+
+    pub unsafe fn lift(&self, cx: &mut impl Call, buffer: *mut u8) {
+        unsafe { self.ptr.lift.unwrap()((&raw mut *cx).cast(), buffer.cast()) };
+    }
+
+    pub unsafe fn lower(&self, cx: &mut impl Call, buffer: *mut u8) {
+        unsafe { self.ptr.lower.unwrap()((&raw mut *cx).cast(), buffer.cast()) };
+    }
+
+    pub fn abi_payload_size(&self) -> usize {
+        self.ptr.abi_payload_size
+    }
+
+    pub fn abi_payload_align(&self) -> usize {
+        self.ptr.abi_payload_align
+    }
 }
 
 impl fmt::Debug for Future {
@@ -1049,6 +1144,50 @@ impl Stream {
     pub fn ty(&self) -> Option<Type> {
         Type::from_raw_opt(self.wit, self.ptr.ty)
     }
+
+    pub fn new(&self) -> unsafe extern "C" fn() -> u64 {
+        self.ptr.new.unwrap()
+    }
+
+    pub fn read(&self) -> unsafe extern "C" fn(u32, *mut c_void, usize) -> u32 {
+        self.ptr.read.unwrap()
+    }
+
+    pub fn write(&self) -> unsafe extern "C" fn(u32, *const c_void, usize) -> u32 {
+        self.ptr.write.unwrap()
+    }
+
+    pub fn cancel_read(&self) -> unsafe extern "C" fn(u32) -> u32 {
+        self.ptr.cancel_read.unwrap()
+    }
+
+    pub fn cancel_write(&self) -> unsafe extern "C" fn(u32) -> u32 {
+        self.ptr.cancel_write.unwrap()
+    }
+
+    pub fn drop_readable(&self) -> unsafe extern "C" fn(u32) {
+        self.ptr.drop_readable.unwrap()
+    }
+
+    pub fn drop_writable(&self) -> unsafe extern "C" fn(u32) {
+        self.ptr.drop_writable.unwrap()
+    }
+
+    pub unsafe fn lift(&self, cx: &mut impl Call, buffer: *mut u8) {
+        unsafe { self.ptr.lift.unwrap()((&raw mut *cx).cast(), buffer.cast()) };
+    }
+
+    pub unsafe fn lower(&self, cx: &mut impl Call, buffer: *mut u8) {
+        unsafe { self.ptr.lower.unwrap()((&raw mut *cx).cast(), buffer.cast()) };
+    }
+
+    pub fn abi_payload_size(&self) -> usize {
+        self.ptr.abi_payload_size
+    }
+
+    pub fn abi_payload_align(&self) -> usize {
+        self.ptr.abi_payload_align
+    }
 }
 
 impl fmt::Debug for Stream {

crates/wit-dylib/src/bindgen.rs

@@ -425,6 +425,48 @@ pub fn task_return(
     )
 }
 
+pub fn lift_payload(adapter: &mut Adapter, resolve: &Resolve, ty: Type) -> Function {
+    let mut c = FunctionCompiler::new(adapter, resolve, 2);
+    c.ctx = Some(TempLocal::new(0, ValType::I32));
+    let frame = StackFrame {
+        size: 16,
+        abi_param_offset: None,
+        retptr_offset: None,
+        stack_temp_offset: Some(0),
+    };
+    c.allocate_stack_frame(&frame);
+    c.lift(
+        &AbiLoc::Memory(Memory {
+            addr: TempLocal::new(1, ValType::I32),
+            offset: 0,
+        }),
+        ty,
+    );
+    c.deallocate_stack_frame(&frame);
+    c.finish()
+}
+
+pub fn lower_payload(adapter: &mut Adapter, resolve: &Resolve, ty: Type) -> Function {
+    let mut c = FunctionCompiler::new(adapter, resolve, 2);
+    c.ctx = Some(TempLocal::new(0, ValType::I32));
+    let frame = StackFrame {
+        size: 16,
+        abi_param_offset: None,
+        retptr_offset: None,
+        stack_temp_offset: Some(0),
+    };
+    c.allocate_stack_frame(&frame);
+    c.lower(
+        ty,
+        &AbiLoc::Memory(Memory {
+            addr: TempLocal::new(1, ValType::I32),
+            offset: 0,
+        }),
+    );
+    c.deallocate_stack_frame(&frame);
+    c.finish()
+}
+
 struct FunctionCompiler<'a> {
     adapter: &'a mut Adapter,
     resolve: &'a Resolve,