Skip to content

Commit 8227b40

Browse files
vados-cosmonicvados-cosmonic
committed
feat: implement missing async arg result handling
This commit implements some (but not all) of the async argument and parameter handling that is emitted for function calls.
1 parent 24eebf7 commit 8227b40

1 file changed

Lines changed: 173 additions & 79 deletions

File tree

‎crates/core/src/abi.rs‎

Lines changed: 173 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
use std::fmt;
2+
use std::iter;
3+
24
pub use wit_parser::abi::{AbiVariant, FlatTypes, WasmSignature, WasmType};
35
use wit_parser::{
46
align_to_arch, Alignment, ArchitectureSize, ElementInfo, Enum, Flags, FlagsRepr, Function,
@@ -920,6 +922,7 @@ struct Generator<'a, B: Bindgen> {
920922
}
921923

922924
const MAX_FLAT_PARAMS: usize = 16;
925+
const MAX_FLAT_ASYNC_PARAMS: usize = 4;
923926

924927
impl<'a, B: Bindgen> Generator<'a, B> {
925928
fn new(resolve: &'a Resolve, bindgen: &'a mut B) -> Generator<'a, B> {
@@ -1075,57 +1078,85 @@ impl<'a, B: Bindgen> Generator<'a, B> {
10751078
amt: usize::from(func.result.is_some()),
10761079
});
10771080
}
1081+
10781082
LiftLower::LiftArgsLowerResults => {
1079-
if let (AbiVariant::GuestImport, true) = (variant, async_) {
1080-
todo!("implement host-side support for async lift/lower");
1081-
}
1083+
let max_flat_params = match (variant, async_) {
1084+
(AbiVariant::GuestImport | AbiVariant::GuestImportAsync, _is_async @ true) => {
1085+
MAX_FLAT_ASYNC_PARAMS
1086+
}
1087+
_ => MAX_FLAT_PARAMS,
1088+
};
10821089

1090+
// Read parameters from memory
10831091
let read_from_memory = |self_: &mut Self| {
10841092
let mut offset = ArchitectureSize::default();
1085-
let ptr = self_.stack.pop().unwrap();
1093+
let ptr = self_
1094+
.stack
1095+
.pop()
1096+
.expect("empty stack during read param from memory");
10861097
for (_, ty) in func.params.iter() {
10871098
offset = align_to_arch(offset, self_.bindgen.sizes().align(ty));
10881099
self_.read_from_memory(ty, ptr.clone(), offset);
10891100
offset += self_.bindgen.sizes().size(ty);
10901101
}
10911102
};
10921103

1093-
if !sig.indirect_params {
1094-
// If parameters are not passed indirectly then we lift each
1104+
// Resolve parameters
1105+
if sig.indirect_params {
1106+
// If parameters were passed indirectly, arguments must be
1107+
// read in succession from memory, with the pointer to the arguments
1108+
// being the first argument to the function.
1109+
self.emit(&Instruction::GetArg { nth: 0 });
1110+
read_from_memory(self);
1111+
} else {
1112+
// ... otherwise, if parameters were passed directly then we lift each
10951113
// argument in succession from the component wasm types that
10961114
// make-up the type.
10971115
let mut offset = 0;
10981116
for (_, ty) in func.params.iter() {
1099-
let types = flat_types(self.resolve, ty).unwrap();
1117+
let types = flat_types(self.resolve, ty, Some(max_flat_params))
1118+
.expect("direct parameter load failed to produce types during generation of fn call");
11001119
for _ in 0..types.len() {
11011120
self.emit(&Instruction::GetArg { nth: offset });
11021121
offset += 1;
11031122
}
11041123
self.lift(ty);
11051124
}
1106-
} else {
1107-
// ... otherwise argument is read in succession from memory
1108-
// where the pointer to the arguments is the first argument
1109-
// to the function.
1110-
self.emit(&Instruction::GetArg { nth: 0 });
1111-
read_from_memory(self);
11121125
}
11131126

11141127
// ... and that allows us to call the interface types function
11151128
self.emit(&Instruction::CallInterface { func, async_ });
11161129

1130+
// The return value of an async function is *not* the result of the function
1131+
// itself or a pointer but rather a status code.
1132+
//
11171133
// Asynchronous functions will call `task.return` after the
11181134
// interface function completes, so lowering is conditional
11191135
// based on slightly different logic for the `task.return`
11201136
// intrinsic.
1121-
let (lower_to_memory, async_flat_results) = if async_ {
1122-
let results = match &func.result {
1123-
Some(ty) => flat_types(self.resolve, ty),
1124-
None => Some(Vec::new()),
1125-
};
1126-
(results.is_none(), Some(results))
1127-
} else {
1128-
(sig.retptr, None)
1137+
let (lower_to_memory, async_flat_results) = match (variant, async_, &func.result) {
1138+
// Async guest imports return a i32 status code
1139+
(AbiVariant::GuestImport, _is_async @ true, None) => {
1140+
unreachable!("async guest imports always return a result")
1141+
}
1142+
// Async guest imports return a i32 status code
1143+
(AbiVariant::GuestImport, _is_async @ true, Some(ty)) => {
1144+
// For async guest imports, we know whether we must lower results
1145+
// if there are no params (i.e. the usual out pointer wasn't even required)
1146+
// and we always know the return value will be a i32 status code
1147+
assert!(matches!(ty, Type::U32 | Type::S32));
1148+
(sig.params.is_empty(), Some(Some(vec![WasmType::I32])))
1149+
}
1150+
// All other async cases
1151+
(_, _is_async @ true, func_result) => {
1152+
let results = match &func_result {
1153+
Some(ty) => flat_types(self.resolve, ty, Some(max_flat_params)),
1154+
None => Some(Vec::new()),
1155+
};
1156+
(results.is_none(), Some(results))
1157+
}
1158+
// All other non-async cases
1159+
(_, _is_async @ false, _) => (sig.retptr, None),
11291160
};
11301161

11311162
// This was dynamically allocated by the caller (or async start
@@ -1147,62 +1178,120 @@ impl<'a, B: Bindgen> Generator<'a, B> {
11471178

11481179
self.realloc = Some(realloc);
11491180

1150-
if !lower_to_memory {
1151-
// With no return pointer in use we simply lower the
1152-
// result(s) and return that directly from the function.
1153-
if let Some(ty) = &func.result {
1154-
self.lower(ty);
1181+
// Perform memory lowing of relevant results, including out pointers as well as traditional results
1182+
match (lower_to_memory, sig.retptr, variant) {
1183+
// If no lowering to memory is required, and there is no return pointer in use, we can do nothing
1184+
(_lower_to_memory @ false, _has_ret_ptr @ false, _) => {}
1185+
1186+
// Async guest imports with do no lowering cannot have ret pointers
1187+
// not having to do lowering implies that there was no return pointer provided
1188+
(_lower_to_memory @ false, _has_ret_ptr @ true, AbiVariant::GuestImport)
1189+
if async_ =>
1190+
{
1191+
unreachable!(
1192+
"async guest import cannot avoid lowering when a ret ptr is present"
1193+
)
11551194
}
1156-
} else {
1157-
match variant {
1158-
// When a function is imported to a guest this means
1159-
// it's a host providing the implementation of the
1160-
// import. The result is stored in the pointer
1161-
// specified in the last argument, so we get the
1162-
// pointer here and then write the return value into
1163-
// it.
1164-
AbiVariant::GuestImport => {
1165-
self.emit(&Instruction::GetArg {
1166-
nth: sig.params.len() - 1,
1167-
});
1168-
let ptr = self.stack.pop().unwrap();
1169-
self.write_params_to_memory(&func.result, ptr, Default::default());
1170-
}
11711195

1172-
// For a guest import this is a function defined in
1173-
// wasm, so we're returning a pointer where the
1174-
// value was stored at. Allocate some space here
1175-
// (statically) and then write the result into that
1176-
// memory, returning the pointer at the end.
1177-
AbiVariant::GuestExport | AbiVariant::GuestExportAsync => {
1178-
let ElementInfo { size, align } =
1179-
self.bindgen.sizes().params(&func.result);
1180-
let ptr = self.bindgen.return_pointer(size, align);
1181-
self.write_params_to_memory(
1182-
&func.result,
1183-
ptr.clone(),
1184-
Default::default(),
1185-
);
1186-
self.stack.push(ptr);
1187-
}
1196+
// For sync calls, if no lowering to memory is required and there *is* a return pointer in use
1197+
// then we need to lower then simply lower the result(s) and return that directly from the function.
1198+
(_lower_to_memory @ false, _has_ret_ptr @ true, _) => {
1199+
self.lower(&func.result.expect("return pointer must be present"));
1200+
}
11881201

1189-
AbiVariant::GuestImportAsync | AbiVariant::GuestExportAsyncStackful => {
1190-
unreachable!()
1191-
}
1202+
// We cannot lower to memory if the signature does not have a return pointer in
1203+
// either the params or the result
1204+
(_lower_to_memory @ true, _has_ret_ptr @ false, _) => unreachable!(
1205+
"lowering to memory cannot be performed without a return pointer"
1206+
),
1207+
1208+
// Lowering to memory for a guest import
1209+
//
1210+
// When a function is imported to a guest this means
1211+
// it's a host providing the implementation of the
1212+
// import. The result is stored in the pointer
1213+
// specified in the last argument, so we get the
1214+
// pointer here and then write the return value into
1215+
// it.
1216+
(
1217+
_lower_to_memory @ true,
1218+
_has_ret_ptr @ true,
1219+
AbiVariant::GuestImport | AbiVariant::GuestImportAsync,
1220+
) => {
1221+
self.emit(&Instruction::GetArg {
1222+
nth: sig.params.len() - 1,
1223+
});
1224+
let ptr = self
1225+
.stack
1226+
.pop()
1227+
.expect("empty stack during result lower to memory");
1228+
self.write_params_to_memory(&func.result, ptr, Default::default());
1229+
}
1230+
1231+
// Lowering to memory for a guest export
1232+
//
1233+
// For a guest import this is a function defined in
1234+
// wasm, so we're returning a pointer where the
1235+
// value was stored at. Allocate some space here
1236+
// (statically) and then write the result into that
1237+
// memory, returning the pointer at the end.
1238+
(
1239+
_lower_to_memory @ true,
1240+
_has_ret_ptr @ true,
1241+
AbiVariant::GuestExport | AbiVariant::GuestExportAsync,
1242+
) => {
1243+
let ElementInfo { size, align } = self.bindgen.sizes().params(&func.result);
1244+
let ptr = self.bindgen.return_pointer(size, align);
1245+
self.write_params_to_memory(&func.result, ptr.clone(), Default::default());
1246+
self.stack.push(ptr);
1247+
}
1248+
1249+
(
1250+
_lower_to_memory @ true,
1251+
_has_ret_ptr @ true,
1252+
AbiVariant::GuestExportAsyncStackful,
1253+
) => {
1254+
todo!("stackful async exports are not supported")
11921255
}
11931256
}
11941257

1195-
if let Some(results) = async_flat_results {
1196-
let name = &format!("[task-return]{}", func.name);
1197-
let params = results.as_deref().unwrap_or(&[WasmType::Pointer]);
1258+
// Build and emit the appropriate return
1259+
match (variant, async_flat_results) {
1260+
// Async guest imports always return a i32 status code
1261+
(AbiVariant::GuestImport | AbiVariant::GuestImportAsync, None) if async_ => {
1262+
unreachable!("async guest imports must have a return")
1263+
}
11981264

1199-
self.emit(&Instruction::AsyncTaskReturn { name, params });
1200-
} else {
1201-
self.emit(&Instruction::Return {
1202-
func,
1203-
amt: sig.results.len(),
1204-
});
1265+
// Async guest imports with results return the status code, not a pointer to any results
1266+
(AbiVariant::GuestImport | AbiVariant::GuestImportAsync, Some(results))
1267+
if async_ =>
1268+
{
1269+
let name = &format!("[task-return]{}", func.name);
1270+
let params = results.as_deref().unwrap_or(&[WasmType::I32]);
1271+
self.emit(&Instruction::AsyncTaskReturn { name, params });
1272+
}
1273+
1274+
// All async/non-async cases with results that need to be returned are present here
1275+
//
1276+
// In practice, async imports should not end up here, as the returned result of an
1277+
// async import is *not* a pointer but instead a status code.
1278+
(_, Some(results)) => {
1279+
let name = &format!("[task-return]{}", func.name);
1280+
let params = results.as_deref().unwrap_or(&[WasmType::Pointer]);
1281+
self.emit(&Instruction::AsyncTaskReturn { name, params });
1282+
}
1283+
1284+
// All async/non-async cases with no results simply return
1285+
//
1286+
// In practice, an async import will never get here (it always has a result, the error code)
1287+
(_, None) => {
1288+
self.emit(&Instruction::Return {
1289+
func,
1290+
amt: sig.results.len(),
1291+
});
1292+
}
12051293
}
1294+
12061295
self.realloc = None;
12071296
}
12081297
}
@@ -1257,7 +1346,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
12571346
let mut operands = operands;
12581347
let mut operands_for_ty;
12591348
for ty in types {
1260-
let types = flat_types(self.resolve, ty).unwrap();
1349+
let types = flat_types(self.resolve, ty, None).unwrap();
12611350
(operands_for_ty, operands) = operands.split_at(types.len());
12621351
self.stack.extend_from_slice(operands_for_ty);
12631352
self.deallocate(ty, what);
@@ -1455,7 +1544,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
14551544
cases: impl IntoIterator<Item = Option<&'b Type>>,
14561545
) -> Vec<WasmType> {
14571546
use Instruction::*;
1458-
let results = flat_types(self.resolve, ty).unwrap();
1547+
let results = flat_types(self.resolve, ty, None).unwrap();
14591548
let mut casts = Vec::new();
14601549
for (i, ty) in cases.into_iter().enumerate() {
14611550
self.push_block();
@@ -1472,7 +1561,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
14721561
// Determine the types of all the wasm values we just
14731562
// pushed, and record how many. If we pushed too few
14741563
// then we'll need to push some zeros after this.
1475-
let temp = flat_types(self.resolve, ty).unwrap();
1564+
let temp = flat_types(self.resolve, ty, None).unwrap();
14761565
pushed += temp.len();
14771566

14781567
// For all the types pushed we may need to insert some
@@ -1638,13 +1727,13 @@ impl<'a, B: Bindgen> Generator<'a, B> {
16381727
types: impl Iterator<Item = &'b Type>,
16391728
mut iter: impl FnMut(&mut Self, &Type),
16401729
) {
1641-
let temp = flat_types(self.resolve, container).unwrap();
1730+
let temp = flat_types(self.resolve, container, None).unwrap();
16421731
let mut args = self
16431732
.stack
16441733
.drain(self.stack.len() - temp.len()..)
16451734
.collect::<Vec<_>>();
16461735
for ty in types {
1647-
let temp = flat_types(self.resolve, ty).unwrap();
1736+
let temp = flat_types(self.resolve, ty, None).unwrap();
16481737
self.stack.extend(args.drain(..temp.len()));
16491738
iter(self, ty);
16501739
}
@@ -1657,7 +1746,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
16571746
cases: impl IntoIterator<Item = Option<&'b Type>>,
16581747
mut iter: impl FnMut(&mut Self, &Type),
16591748
) {
1660-
let params = flat_types(self.resolve, ty).unwrap();
1749+
let params = flat_types(self.resolve, ty, None).unwrap();
16611750
let mut casts = Vec::new();
16621751
let block_inputs = self
16631752
.stack
@@ -1668,7 +1757,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
16681757
if let Some(ty) = ty {
16691758
// Push only the values we need for this variant onto
16701759
// the stack.
1671-
let temp = flat_types(self.resolve, ty).unwrap();
1760+
let temp = flat_types(self.resolve, ty, None).unwrap();
16721761
self.stack
16731762
.extend(block_inputs[..temp.len()].iter().cloned());
16741763

@@ -2399,9 +2488,14 @@ fn cast(from: WasmType, to: WasmType) -> Bitcast {
23992488
}
24002489
}
24012490

2402-
fn flat_types(resolve: &Resolve, ty: &Type) -> Option<Vec<WasmType>> {
2403-
let mut storage = [WasmType::I32; MAX_FLAT_PARAMS];
2404-
let mut flat = FlatTypes::new(&mut storage);
2491+
/// Flatten types in a given type
2492+
///
2493+
/// It is sometimes necessary to restrict the number of max parameters dynamically,
2494+
/// for example during an async guest import call (flat params are limited to 4)
2495+
fn flat_types(resolve: &Resolve, ty: &Type, max_params: Option<usize>) -> Option<Vec<WasmType>> {
2496+
let mut storage =
2497+
iter::repeat_n(WasmType::I32, max_params.unwrap_or(MAX_FLAT_PARAMS)).collect::<Vec<_>>();
2498+
let mut flat = FlatTypes::new(storage.as_mut_slice());
24052499
if resolve.push_flat(ty, &mut flat) {
24062500
Some(flat.to_vec())
24072501
} else {

0 commit comments

Comments
 (0)