Skip to content

Commit e0b816e

Browse files
committed
Adjust allocation strategy for TinyGo-derived modules so they don't immediately crash.
1 parent 6ed1788 commit e0b816e

File tree

2 files changed

+52
-3
lines changed

2 files changed

+52
-3
lines changed

crates/wasmparser/src/readers/core/custom.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ impl<'a> CustomSectionReader<'a> {
103103
/// Return value of [`CustomSectionReader::as_known`].
104104
///
105105
/// Note that this is `#[non_exhaustive]` because depending on crate features
106-
/// this enumeration will different entries.
106+
/// this enumeration will contain different entries.
107107
#[allow(missing_docs)]
108108
#[non_exhaustive]
109109
pub enum KnownCustom<'a> {

crates/wit-component/src/gc.rs

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,43 @@ impl<'a> Module<'a> {
500500
live_iter(&self.live_tables, self.tables.iter())
501501
}
502502

503+
/// Returns whether the module appears to have been generated by TinyGo. Absent
504+
/// any evidence, returns false.
505+
fn is_from_tiny_go(&self) -> bool {
506+
self.producers.as_ref().map_or(false, |producers| {
507+
producers.iter().any(|(field, values)| {
508+
field == "processed-by"
509+
&& values
510+
.iter()
511+
.any(|(processor, _version)| processor == "TinyGo")
512+
})
513+
})
514+
}
515+
516+
/// Returns a new Wasm function body which implements `cabi_realloc` to call a
517+
/// `malloc` that's exported from the main module.
518+
fn realloc_via_malloc(&self) -> Result<wasm_encoder::Function> {
519+
// Find `malloc`.
520+
let malloc_index = match self.exports.get("malloc") {
521+
None => bail!(
522+
"found no exported `malloc` function to use to allocate the adapter's state in a TinyGo-derived module"
523+
),
524+
Some(export) => {
525+
if export.kind != ExternalKind::Func {
526+
bail!("found a `malloc` export, but it was not a function")
527+
}
528+
export.index
529+
}
530+
};
531+
532+
let mut func = wasm_encoder::Function::new([]);
533+
534+
func.instructions().local_get(3); // desired new size
535+
func.instructions().call(malloc_index);
536+
537+
Ok(func)
538+
}
539+
503540
/// Encodes this `Module` to a new wasm module which is gc'd and only
504541
/// contains the items that are live as calculated by the `liveness` pass.
505542
fn encode(&mut self, main_module_realloc: Option<&str>) -> Result<Vec<u8>> {
@@ -733,10 +770,22 @@ impl<'a> Module<'a> {
733770

734771
if sp.is_some() && (realloc_index.is_none() || allocation_state.is_none()) {
735772
// Either the main module does _not_ export a realloc function, or it is not safe to use for stack
736-
// allocation because we have no way to short-circuit reentrance, so we'll use `memory.grow` instead.
773+
// allocation because we have no way to short-circuit reentrance, so we'll provide our own realloc
774+
// function instead.
737775
realloc_index = Some(num_func_imports + funcs.len());
738776
funcs.function(add_realloc_type(&mut types));
739-
code.function(&realloc_via_memory_grow());
777+
778+
let realloc_func = if self.is_from_tiny_go() {
779+
// If it appears the module was emitted by TinyGo, we delegate
780+
// to its `malloc()` function. (TinyGo assumes its GC has reign
781+
// over the whole memory and quickly overwrites the adapter's
782+
// State struct unless we inform its GC of the memory we use.)
783+
self.realloc_via_malloc()?
784+
} else {
785+
// If it's not TinyGo, use `memory.grow` instead.
786+
realloc_via_memory_grow()
787+
};
788+
code.function(&realloc_func);
740789
}
741790

742791
// Inject a start function to initialize the stack pointer which will be local to this module. This only

0 commit comments

Comments
 (0)