support callback-less (AKA stackful) async lifts

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

Author: dicej

crates/wasmparser/src/validator/component.rs

@@ -3,7 +3,7 @@
 use super::{
     check_max,
     component_types::{
-        AliasableResourceId, ComponentAnyTypeId, ComponentCoreInstanceTypeId,
+        Abi, AliasableResourceId, ComponentAnyTypeId, ComponentCoreInstanceTypeId,
         ComponentCoreModuleTypeId, ComponentCoreTypeId, ComponentDefinedType,
         ComponentDefinedTypeId, ComponentEntityType, ComponentFuncType, ComponentFuncTypeId,
         ComponentInstanceType, ComponentInstanceTypeId, ComponentType, ComponentTypeId,
@@ -963,7 +963,21 @@ impl ComponentState {
 
         // Lifting a function is for an export, so match the expected canonical ABI
         // export signature
-        let info = ty.lower(types, false, options.contains(&CanonicalOption::Async));
+        let info = ty.lower(
+            types,
+            if options.contains(&CanonicalOption::Async) {
+                if options
+                    .iter()
+                    .any(|v| matches!(v, CanonicalOption::Callback(_)))
+                {
+                    Abi::LiftAsync
+                } else {
+                    Abi::LiftAsyncStackful
+                }
+            } else {
+                Abi::LiftSync
+            },
+        );
         self.check_options(
             Some(core_ty),
             &info,
@@ -1012,7 +1026,14 @@ impl ComponentState {
 
         // Lowering a function is for an import, so use a function type that matches
         // the expected canonical ABI import signature.
-        let info = ty.lower(types, true, options.contains(&CanonicalOption::Async));
+        let info = ty.lower(
+            types,
+            if options.contains(&CanonicalOption::Async) {
+                Abi::LowerAsync
+            } else {
+                Abi::LowerSync
+            },
+        );
 
         self.check_options(None, &info, &options, types, offset, features, true)?;
 

crates/wasmparser/src/validator/component_types.rs

@@ -891,21 +891,34 @@ impl TypeData for ComponentFuncType {
     }
 }
 
+#[derive(Copy, Clone, Debug)]
+pub(crate) enum Abi {
+    LowerSync,
+    LowerAsync,
+    LiftSync,
+    LiftAsync,
+    LiftAsyncStackful,
+}
+
 impl ComponentFuncType {
     /// Lowers the component function type to core parameter and result types for the
     /// canonical ABI.
-    pub(crate) fn lower(&self, types: &TypeList, is_lower: bool, async_: bool) -> LoweringInfo {
+    pub(crate) fn lower(&self, types: &TypeList, abi: Abi) -> LoweringInfo {
         let mut info = LoweringInfo::default();
 
-        if async_ && is_lower {
-            for _ in 0..2 {
-                info.params.push(ValType::I32);
+        let is_lower = match abi {
+            Abi::LowerAsync => {
+                for _ in 0..2 {
+                    info.params.push(ValType::I32);
+                }
+                info.results.push(ValType::I32);
+                info.requires_memory = true;
+                info.requires_realloc = self.results.iter().any(|(_, ty)| ty.contains_ptr(types));
+                return info;
             }
-            info.results.push(ValType::I32);
-            info.requires_memory = true;
-            info.requires_realloc = self.results.iter().any(|(_, ty)| ty.contains_ptr(types));
-            return info;
-        }
+            Abi::LowerSync => true,
+            Abi::LiftSync | Abi::LiftAsync | Abi::LiftAsyncStackful => false,
+        };
 
         for (_, ty) in self.params.iter() {
             // Check to see if `ty` has a pointer somewhere in it, needed for
@@ -940,32 +953,37 @@ impl ComponentFuncType {
             }
         }
 
-        if async_ {
-            info.results.push(ValType::I32);
-        } else {
-            for (_, ty) in self.results.iter() {
-                // Results of lowered functions that contains pointers must be
-                // allocated by the callee meaning that realloc is required.
-                // Results of lifted function are allocated by the guest which
-                // means that no realloc option is necessary.
-                if is_lower && !info.requires_realloc {
-                    info.requires_realloc = ty.contains_ptr(types);
-                }
+        match abi {
+            Abi::LowerAsync => unreachable!(),
+            Abi::LowerSync | Abi::LiftSync => {
+                for (_, ty) in self.results.iter() {
+                    // Results of lowered functions that contains pointers must be
+                    // allocated by the callee meaning that realloc is required.
+                    // Results of lifted function are allocated by the guest which
+                    // means that no realloc option is necessary.
+                    if is_lower && !info.requires_realloc {
+                        info.requires_realloc = ty.contains_ptr(types);
+                    }
 
-                if !ty.push_wasm_types(types, &mut info.results) {
-                    // Too many results to return directly, either a retptr parameter will be used (import)
-                    // or a single pointer will be returned (export)
-                    info.results.clear();
-                    if is_lower {
-                        info.params.max = MAX_LOWERED_TYPES;
-                        assert!(info.params.push(ValType::I32));
-                    } else {
-                        assert!(info.results.push(ValType::I32));
+                    if !ty.push_wasm_types(types, &mut info.results) {
+                        // Too many results to return directly, either a retptr parameter will be used (import)
+                        // or a single pointer will be returned (export)
+                        info.results.clear();
+                        if is_lower {
+                            info.params.max = MAX_LOWERED_TYPES;
+                            assert!(info.params.push(ValType::I32));
+                        } else {
+                            assert!(info.results.push(ValType::I32));
+                        }
+                        info.requires_memory = true;
+                        break;
                     }
-                    info.requires_memory = true;
-                    break;
                 }
             }
+            Abi::LiftAsync => {
+                info.results.push(ValType::I32);
+            }
+            Abi::LiftAsyncStackful => {}
         }
 
         // Memory is always required when realloc is required

crates/wit-component/src/encoding.rs

@@ -169,7 +169,7 @@ impl RequiredOptions {
                 ret |= RequiredOptions::REALLOC;
             }
         }
-        if abi == AbiVariant::GuestExportAsync {
+        if let AbiVariant::GuestExportAsync | AbiVariant::GuestExportAsyncStackful = abi {
             ret |= RequiredOptions::ASYNC;
         }
         ret

crates/wit-component/src/validation.rs

@@ -816,6 +816,8 @@ impl ExportMap {
         let full_name = name;
         let (abi, name) = if let Some(name) = names.async_name(name) {
             (AbiVariant::GuestExportAsync, name)
+        } else if let Some(name) = names.async_stackful_name(name) {
+            (AbiVariant::GuestExportAsyncStackful, name)
         } else {
             (AbiVariant::GuestExport, name)
         };
@@ -1081,6 +1083,7 @@ trait NameMangling {
     fn subtask_drop(&self) -> Option<&str>;
     fn callback_name<'a>(&self, s: &'a str) -> Option<&'a str>;
     fn async_name<'a>(&self, s: &'a str) -> Option<&'a str>;
+    fn async_stackful_name<'a>(&self, s: &'a str) -> Option<&'a str>;
     fn error_context_new(&self, s: &str) -> Option<StringEncoding>;
     fn error_context_debug_message<'a>(&self, s: &'a str) -> Option<(StringEncoding, &'a str)>;
     fn error_context_drop(&self) -> Option<&str>;
@@ -1176,6 +1179,10 @@ impl NameMangling for Standard {
         _ = s;
         None
     }
+    fn async_stackful_name<'a>(&self, s: &'a str) -> Option<&'a str> {
+        _ = s;
+        None
+    }
     fn error_context_new(&self, s: &str) -> Option<StringEncoding> {
         _ = s;
         None
@@ -1357,6 +1364,9 @@ impl NameMangling for Legacy {
     fn async_name<'a>(&self, s: &'a str) -> Option<&'a str> {
         s.strip_prefix("[async]")
     }
+    fn async_stackful_name<'a>(&self, s: &'a str) -> Option<&'a str> {
+        s.strip_prefix("[async-stackful]")
+    }
     fn error_context_new(&self, s: &str) -> Option<StringEncoding> {
         parse_encoding(
             s.strip_prefix("[error-context-new;encoding=")?

crates/wit-parser/src/abi.rs

@@ -129,6 +129,7 @@ pub enum AbiVariant {
     GuestExport,
     GuestImportAsync,
     GuestExportAsync,
+    GuestExportAsyncStackful,
 }
 
 impl Resolve {
@@ -164,7 +165,9 @@ impl Resolve {
                 (&func.kind, variant),
                 (
                     crate::FunctionKind::Method(_),
-                    AbiVariant::GuestExport | AbiVariant::GuestExportAsync
+                    AbiVariant::GuestExport
+                        | AbiVariant::GuestExportAsync
+                        | AbiVariant::GuestExportAsyncStackful
                 )
             ) {
                 // Guest exported methods always receive resource rep as first argument
@@ -178,13 +181,24 @@ impl Resolve {
             }
         }
 
-        if let AbiVariant::GuestExportAsync = variant {
-            return WasmSignature {
-                params,
-                indirect_params,
-                results: vec![WasmType::Pointer],
-                retptr: false,
-            };
+        match variant {
+            AbiVariant::GuestExportAsync => {
+                return WasmSignature {
+                    params,
+                    indirect_params,
+                    results: vec![WasmType::Pointer],
+                    retptr: false,
+                };
+            }
+            AbiVariant::GuestExportAsyncStackful => {
+                return WasmSignature {
+                    params,
+                    indirect_params,
+                    results: Vec::new(),
+                    retptr: false,
+                };
+            }
+            _ => {}
         }
 
         let mut results = Vec::new();