Skip to content

Commit 3badd19

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 3badd19

1 file changed

Lines changed: 179 additions & 78 deletions

File tree

‎crates/core/src/abi.rs‎

Lines changed: 179 additions & 78 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,93 @@ 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+
(
1140+
AbiVariant::GuestImport | AbiVariant::GuestImportAsync,
1141+
_is_async @ true,
1142+
None,
1143+
) => {
1144+
unreachable!("async guest imports always return a result")
1145+
}
1146+
// Async guest imports return a i32 status code
1147+
(
1148+
AbiVariant::GuestImport | AbiVariant::GuestImportAsync,
1149+
_is_async @ true,
1150+
Some(ty),
1151+
) => {
1152+
// For async guest imports, we know whether we must lower results
1153+
// if there are no params (i.e. the usual out pointer wasn't even required)
1154+
// and we always know the return value will be a i32 status code
1155+
assert!(matches!(ty, Type::U32 | Type::S32));
1156+
(sig.params.is_empty(), Some(Some(vec![WasmType::I32])))
1157+
}
1158+
// All other async cases
1159+
(_, _is_async @ true, func_result) => {
1160+
let results = match &func_result {
1161+
Some(ty) => flat_types(self.resolve, ty, Some(max_flat_params)),
1162+
None => Some(Vec::new()),
1163+
};
1164+
(results.is_none(), Some(results))
1165+
}
1166+
// All other non-async cases
1167+
(_, _is_async @ false, _) => (sig.retptr, None),
11291168
};
11301169

11311170
// This was dynamically allocated by the caller (or async start
@@ -1147,62 +1186,119 @@ impl<'a, B: Bindgen> Generator<'a, B> {
11471186

11481187
self.realloc = Some(realloc);
11491188

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);
1189+
// Perform memory lowing of relevant results, including out pointers as well as traditional results
1190+
match (lower_to_memory, sig.retptr, variant) {
1191+
// Async guest imports with do no lowering cannot have ret pointers
1192+
// not having to do lowering implies that there was no return pointer provided
1193+
(_lower_to_memory @ false, _has_ret_ptr @ true, AbiVariant::GuestImport)
1194+
if async_ =>
1195+
{
1196+
unreachable!(
1197+
"async guest import cannot avoid lowering when a ret ptr is present"
1198+
)
11551199
}
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-
}
11711200

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);
1201+
// For sync calls, if no lowering to memory is required and there *is* a return pointer in use
1202+
// then we need to lower then simply lower the result(s) and return that directly from the function.
1203+
(_lower_to_memory @ false, _, _) => {
1204+
if let Some(ty) = &func.result {
1205+
self.lower(ty);
11871206
}
1207+
}
11881208

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

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]);
1265+
// Build and emit the appropriate return
1266+
match (variant, async_flat_results) {
1267+
// Async guest imports always return a i32 status code
1268+
(AbiVariant::GuestImport | AbiVariant::GuestImportAsync, None) if async_ => {
1269+
unreachable!("async guest imports must have a return")
1270+
}
11981271

1199-
self.emit(&Instruction::AsyncTaskReturn { name, params });
1200-
} else {
1201-
self.emit(&Instruction::Return {
1202-
func,
1203-
amt: sig.results.len(),
1204-
});
1272+
// Async guest imports with results return the status code, not a pointer to any results
1273+
(AbiVariant::GuestImport | AbiVariant::GuestImportAsync, Some(results))
1274+
if async_ =>
1275+
{
1276+
let name = &format!("[task-return]{}", func.name);
1277+
let params = results.as_deref().unwrap_or(&[WasmType::I32]);
1278+
self.emit(&Instruction::AsyncTaskReturn { name, params });
1279+
}
1280+
1281+
// All async/non-async cases with results that need to be returned are present here
1282+
//
1283+
// In practice, async imports should not end up here, as the returned result of an
1284+
// async import is *not* a pointer but instead a status code.
1285+
(_, Some(results)) => {
1286+
let name = &format!("[task-return]{}", func.name);
1287+
let params = results.as_deref().unwrap_or(&[WasmType::Pointer]);
1288+
self.emit(&Instruction::AsyncTaskReturn { name, params });
1289+
}
1290+
1291+
// All async/non-async cases with no results simply return
1292+
//
1293+
// In practice, an async import will never get here (it always has a result, the error code)
1294+
(_, None) => {
1295+
self.emit(&Instruction::Return {
1296+
func,
1297+
amt: sig.results.len(),
1298+
});
1299+
}
12051300
}
1301+
12061302
self.realloc = None;
12071303
}
12081304
}
@@ -1257,7 +1353,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
12571353
let mut operands = operands;
12581354
let mut operands_for_ty;
12591355
for ty in types {
1260-
let types = flat_types(self.resolve, ty).unwrap();
1356+
let types = flat_types(self.resolve, ty, None).unwrap();
12611357
(operands_for_ty, operands) = operands.split_at(types.len());
12621358
self.stack.extend_from_slice(operands_for_ty);
12631359
self.deallocate(ty, what);
@@ -1455,7 +1551,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
14551551
cases: impl IntoIterator<Item = Option<&'b Type>>,
14561552
) -> Vec<WasmType> {
14571553
use Instruction::*;
1458-
let results = flat_types(self.resolve, ty).unwrap();
1554+
let results = flat_types(self.resolve, ty, None).unwrap();
14591555
let mut casts = Vec::new();
14601556
for (i, ty) in cases.into_iter().enumerate() {
14611557
self.push_block();
@@ -1472,7 +1568,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
14721568
// Determine the types of all the wasm values we just
14731569
// pushed, and record how many. If we pushed too few
14741570
// then we'll need to push some zeros after this.
1475-
let temp = flat_types(self.resolve, ty).unwrap();
1571+
let temp = flat_types(self.resolve, ty, None).unwrap();
14761572
pushed += temp.len();
14771573

14781574
// For all the types pushed we may need to insert some
@@ -1638,13 +1734,13 @@ impl<'a, B: Bindgen> Generator<'a, B> {
16381734
types: impl Iterator<Item = &'b Type>,
16391735
mut iter: impl FnMut(&mut Self, &Type),
16401736
) {
1641-
let temp = flat_types(self.resolve, container).unwrap();
1737+
let temp = flat_types(self.resolve, container, None).unwrap();
16421738
let mut args = self
16431739
.stack
16441740
.drain(self.stack.len() - temp.len()..)
16451741
.collect::<Vec<_>>();
16461742
for ty in types {
1647-
let temp = flat_types(self.resolve, ty).unwrap();
1743+
let temp = flat_types(self.resolve, ty, None).unwrap();
16481744
self.stack.extend(args.drain(..temp.len()));
16491745
iter(self, ty);
16501746
}
@@ -1657,7 +1753,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
16571753
cases: impl IntoIterator<Item = Option<&'b Type>>,
16581754
mut iter: impl FnMut(&mut Self, &Type),
16591755
) {
1660-
let params = flat_types(self.resolve, ty).unwrap();
1756+
let params = flat_types(self.resolve, ty, None).unwrap();
16611757
let mut casts = Vec::new();
16621758
let block_inputs = self
16631759
.stack
@@ -1668,7 +1764,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
16681764
if let Some(ty) = ty {
16691765
// Push only the values we need for this variant onto
16701766
// the stack.
1671-
let temp = flat_types(self.resolve, ty).unwrap();
1767+
let temp = flat_types(self.resolve, ty, None).unwrap();
16721768
self.stack
16731769
.extend(block_inputs[..temp.len()].iter().cloned());
16741770

@@ -2399,9 +2495,14 @@ fn cast(from: WasmType, to: WasmType) -> Bitcast {
23992495
}
24002496
}
24012497

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);
2498+
/// Flatten types in a given type
2499+
///
2500+
/// It is sometimes necessary to restrict the number of max parameters dynamically,
2501+
/// for example during an async guest import call (flat params are limited to 4)
2502+
fn flat_types(resolve: &Resolve, ty: &Type, max_params: Option<usize>) -> Option<Vec<WasmType>> {
2503+
let mut storage =
2504+
iter::repeat_n(WasmType::I32, max_params.unwrap_or(MAX_FLAT_PARAMS)).collect::<Vec<_>>();
2505+
let mut flat = FlatTypes::new(storage.as_mut_slice());
24052506
if resolve.push_flat(ty, &mut flat) {
24062507
Some(flat.to_vec())
24072508
} else {

0 commit comments

Comments
 (0)