Skip to content

Commit 88374b4

Browse files
authored
Store arguments for async params/results on the stack (#1185)
After #1176 there's no longer any need to store these values on the heap, so store them on the stack instead. This also updates to store params/results in an overlapping allocation which should be reasonable as it's local per-import and results are never written before parameters are read. (and params are never read after results are written).
1 parent 1d8cbb9 commit 88374b4

File tree

6 files changed

+51
-99
lines changed

6 files changed

+51
-99
lines changed

crates/core/src/abi.rs

Lines changed: 34 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -551,19 +551,8 @@ def_instruction! {
551551
blocks: usize,
552552
} : [1] => [0],
553553

554-
/// Allocate the parameter and/or return areas to use for an
555-
/// async-lowered import call.
556-
///
557-
/// This cannot be allocated on the (shadow-)stack since it needs to
558-
/// remain valid until the callee has finished using the buffers, which
559-
/// may be after we pop the current stack frame.
560-
AsyncMalloc { size: usize, align: usize } : [0] => [1],
561-
562554
/// Call an async-lowered import.
563-
///
564-
/// `size` and `align` are used to deallocate the parameter area
565-
/// allocated using `AsyncMalloc` after the callee task returns a value.
566-
AsyncCallWasm { name: &'a str, size: usize, align: usize } : [2] => [0],
555+
AsyncCallWasm { name: &'a str } : [2] => [0],
567556

568557
/// Generate code to run after `CallInterface` for an async-lifted export.
569558
///
@@ -913,18 +902,15 @@ impl<'a, B: Bindgen> Generator<'a, B> {
913902
self_.stack.push(ptr);
914903
};
915904

916-
let params_size_align = if self.async_ {
905+
if self.async_ {
917906
let ElementInfo { size, align } = self
918907
.bindgen
919908
.sizes()
920909
.record(func.params.iter().map(|(_, ty)| ty));
921-
self.emit(&Instruction::AsyncMalloc {
922-
size: size.size_wasm32(),
923-
align: align.align_wasm32(),
924-
});
925-
let ptr = self.stack.pop().unwrap();
910+
let ptr = self
911+
.bindgen
912+
.return_pointer(size.size_wasm32(), align.align_wasm32());
926913
lower_to_memory(self, ptr);
927-
Some((size, align))
928914
} else {
929915
if !sig.indirect_params {
930916
// If the parameters for this function aren't indirect
@@ -966,47 +952,39 @@ impl<'a, B: Bindgen> Generator<'a, B> {
966952
};
967953
lower_to_memory(self, ptr);
968954
}
969-
None
970-
};
955+
}
971956

972-
// If necessary we may need to prepare a return pointer for
973-
// this ABI.
974-
let dealloc_size_align =
975-
if let Some((params_size, params_align)) = params_size_align {
976-
let ElementInfo { size, align } =
977-
self.bindgen.sizes().record(func.result.iter());
978-
self.emit(&Instruction::AsyncMalloc {
979-
size: size.size_wasm32(),
980-
align: align.align_wasm32(),
981-
});
982-
let ptr = self.stack.pop().unwrap();
957+
if self.async_ {
958+
let ElementInfo { size, align } =
959+
self.bindgen.sizes().record(func.result.iter());
960+
let ptr = self
961+
.bindgen
962+
.return_pointer(size.size_wasm32(), align.align_wasm32());
963+
self.return_pointer = Some(ptr.clone());
964+
self.stack.push(ptr);
965+
966+
assert_eq!(self.stack.len(), 2);
967+
self.emit(&Instruction::AsyncCallWasm {
968+
name: &format!("[async-lower]{}", func.name),
969+
});
970+
} else {
971+
// If necessary we may need to prepare a return pointer for
972+
// this ABI.
973+
if self.variant == AbiVariant::GuestImport && sig.retptr {
974+
let info = self.bindgen.sizes().params(&func.result);
975+
let ptr = self
976+
.bindgen
977+
.return_pointer(info.size.size_wasm32(), info.align.align_wasm32());
983978
self.return_pointer = Some(ptr.clone());
984979
self.stack.push(ptr);
980+
}
985981

986-
assert_eq!(self.stack.len(), 2);
987-
self.emit(&Instruction::AsyncCallWasm {
988-
name: &format!("[async-lower]{}", func.name),
989-
size: params_size.size_wasm32(),
990-
align: params_align.align_wasm32(),
991-
});
992-
Some((size, align))
993-
} else {
994-
if self.variant == AbiVariant::GuestImport && sig.retptr {
995-
let info = self.bindgen.sizes().params(&func.result);
996-
let ptr = self
997-
.bindgen
998-
.return_pointer(info.size.size_wasm32(), info.align.align_wasm32());
999-
self.return_pointer = Some(ptr.clone());
1000-
self.stack.push(ptr);
1001-
}
1002-
1003-
assert_eq!(self.stack.len(), sig.params.len());
1004-
self.emit(&Instruction::CallWasm {
1005-
name: &func.name,
1006-
sig: &sig,
1007-
});
1008-
None
1009-
};
982+
assert_eq!(self.stack.len(), sig.params.len());
983+
self.emit(&Instruction::CallWasm {
984+
name: &func.name,
985+
sig: &sig,
986+
});
987+
}
1010988

1011989
if !(sig.retptr || self.async_) {
1012990
// With no return pointer in use we can simply lift the
@@ -1043,14 +1021,6 @@ impl<'a, B: Bindgen> Generator<'a, B> {
10431021
self.emit(&Instruction::Flush {
10441022
amt: usize::from(func.result.is_some()),
10451023
});
1046-
1047-
if let Some((size, align)) = dealloc_size_align {
1048-
self.stack.push(ptr);
1049-
self.emit(&Instruction::GuestDeallocate {
1050-
size: size.size_wasm32(),
1051-
align: align.align_wasm32(),
1052-
});
1053-
}
10541024
}
10551025

10561026
self.emit(&Instruction::Return {

crates/csharp/src/function.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1253,8 +1253,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
12531253
results.extend(operands.iter().take(*amt).map(|v| v.clone()));
12541254
}
12551255

1256-
Instruction::AsyncMalloc { .. }
1257-
| Instruction::AsyncPostCallInterface { .. }
1256+
Instruction::AsyncPostCallInterface { .. }
12581257
| Instruction::AsyncCallReturn { .. }
12591258
| Instruction::FutureLower { .. }
12601259
| Instruction::FutureLift { .. }

crates/guest-rust/rt/src/async_support.rs

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
extern crate std;
55
use core::sync::atomic::{AtomicBool, Ordering};
6-
use std::alloc::{self, Layout};
76
use std::any::Any;
87
use std::boxed::Box;
98
use std::collections::{hash_map, HashMap};
@@ -157,7 +156,6 @@ pub fn first_poll<T: 'static>(
157156
#[doc(hidden)]
158157
pub async unsafe fn await_result(
159158
import: unsafe extern "C" fn(*mut u8, *mut u8) -> i32,
160-
params_layout: Layout,
161159
params: *mut u8,
162160
results: *mut u8,
163161
) {
@@ -182,17 +180,13 @@ pub async unsafe fn await_result(
182180
let (tx, rx) = oneshot::channel();
183181
CALLS.insert(call, tx);
184182
rx.await.unwrap();
185-
alloc::dealloc(params, params_layout);
186183
}
187184
STATUS_STARTED => {
188-
alloc::dealloc(params, params_layout);
189185
let (tx, rx) = oneshot::channel();
190186
CALLS.insert(call, tx);
191187
rx.await.unwrap();
192188
}
193-
STATUS_RETURNED | STATUS_DONE => {
194-
alloc::dealloc(params, params_layout);
195-
}
189+
STATUS_RETURNED | STATUS_DONE => {}
196190
_ => unreachable!("unrecognized async call status"),
197191
}
198192

crates/moonbit/src/lib.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2665,8 +2665,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
26652665
results.extend(operands.iter().take(*amt).map(|v| v.clone()));
26662666
}
26672667

2668-
Instruction::AsyncMalloc { .. }
2669-
| Instruction::AsyncPostCallInterface { .. }
2668+
Instruction::AsyncPostCallInterface { .. }
26702669
| Instruction::AsyncCallReturn { .. }
26712670
| Instruction::FutureLower { .. }
26722671
| Instruction::FutureLift { .. }

crates/rust/src/bindgen.rs

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,14 @@ impl Bindgen for FunctionBindgen<'_, '_> {
278278
// stack whereas exports use a per-module return area to cut down on
279279
// stack usage. Note that for imports this also facilitates "adapter
280280
// modules" for components to not have data segments.
281-
if self.gen.in_import {
281+
if size == 0 {
282+
// If the size requested is 0 then we know it won't be written to so
283+
// hand out a null pointer. This can happen with async for example
284+
// when the params or results are zero-sized.
285+
uwrite!(self.src, "let ptr{tmp} = core::ptr::null_mut::<u8>();");
286+
} else if self.gen.in_import {
287+
// Import return areas are stored on the stack since this stack
288+
// frame will be around for the entire function call.
282289
self.import_return_pointer_area_size = self.import_return_pointer_area_size.max(size);
283290
self.import_return_pointer_area_align =
284291
self.import_return_pointer_area_align.max(align);
@@ -287,6 +294,9 @@ impl Bindgen for FunctionBindgen<'_, '_> {
287294
"let ptr{tmp} = ret_area.0.as_mut_ptr().cast::<u8>();"
288295
);
289296
} else {
297+
// Export return areas are stored in `static` memory as they need to
298+
// persist beyond the function call itself (and are cleaned-up in
299+
// `post-return`).
290300
self.gen.return_pointer_area_size = self.gen.return_pointer_area_size.max(size);
291301
self.gen.return_pointer_area_align = self.gen.return_pointer_area_align.max(align);
292302
uwriteln!(
@@ -870,20 +880,14 @@ impl Bindgen for FunctionBindgen<'_, '_> {
870880
self.push_str(");\n");
871881
}
872882

873-
Instruction::AsyncCallWasm { name, size, align } => {
883+
Instruction::AsyncCallWasm { name, .. } => {
874884
let func = self.declare_import(name, &[WasmType::Pointer; 2], &[WasmType::I32]);
875885

876886
let async_support = self.gen.gen.async_support_path();
877-
let tmp = self.tmp();
878-
let layout = format!("layout{tmp}");
879-
let alloc = self.gen.path_to_std_alloc_module();
880-
self.push_str(&format!(
881-
"let {layout} = {alloc}::Layout::from_size_align_unchecked({size}, {align});\n",
882-
));
883887
let operands = operands.join(", ");
884888
uwriteln!(
885889
self.src,
886-
"{async_support}::await_result({func}, {layout}, {operands}).await;"
890+
"{async_support}::await_result({func}, {operands}).await;"
887891
);
888892
}
889893

@@ -951,19 +955,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
951955
self.push_str(";\n");
952956
}
953957

954-
Instruction::AsyncMalloc { size, align } => {
955-
let alloc = self.gen.path_to_std_alloc_module();
956-
let tmp = self.tmp();
957-
let ptr = format!("ptr{tmp}");
958-
let layout = format!("layout{tmp}");
959-
uwriteln!(
960-
self.src,
961-
"let {layout} = {alloc}::Layout::from_size_align_unchecked({size}, {align});
962-
let {ptr} = {alloc}::alloc({layout});"
963-
);
964-
results.push(ptr);
965-
}
966-
967958
Instruction::AsyncPostCallInterface { func } => {
968959
let result = &operands[0];
969960
self.async_result_name = Some(result.clone());

crates/teavm-java/src/lib.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1999,8 +1999,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
19991999
results.extend(operands.iter().take(*amt).map(|v| v.clone()));
20002000
}
20012001

2002-
Instruction::AsyncMalloc { .. }
2003-
| Instruction::AsyncPostCallInterface { .. }
2002+
Instruction::AsyncPostCallInterface { .. }
20042003
| Instruction::AsyncCallReturn { .. }
20052004
| Instruction::FutureLower { .. }
20062005
| Instruction::FutureLift { .. }

0 commit comments

Comments
 (0)