Adjust allocation strategy for TinyGo-derived modules so they don't immediately crash.

Author: erikrose

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

@@ -103,7 +103,7 @@ impl<'a> CustomSectionReader<'a> {
 /// Return value of [`CustomSectionReader::as_known`].
 ///
 /// Note that this is `#[non_exhaustive]` because depending on crate features
-/// this enumeration will different entries.
+/// this enumeration will contain different entries.
 #[allow(missing_docs)]
 #[non_exhaustive]
 pub enum KnownCustom<'a> {

crates/wit-component/src/gc.rs

@@ -500,6 +500,48 @@ impl<'a> Module<'a> {
         live_iter(&self.live_tables, self.tables.iter())
     }
 
+    /// Returns whether the module appears to have been generated by TinyGo. Absent
+    /// any evidence, returns false.
+    fn is_from_tiny_go(&self) -> bool {
+        self.producers.as_ref().map_or(false, |producers| {
+            producers.iter().any(|(field, values)| {
+                field == "processed-by"
+                    && values
+                        .iter()
+                        .any(|(processor, _version)| processor == "TinyGo")
+            })
+        })
+    }
+
+    /// Returns a new Wasm function body which implements `cabi_realloc` to call a
+    /// `malloc` that's exported from the main module.
+    fn realloc_via_malloc(&self) -> Result<wasm_encoder::Function> {
+        //  (func $cabi_realloc (;someIndex;) (type someType) (param i32 i32 i32 i32) (result i32)
+        //    local.get 3
+        //    call $malloc
+        //  )
+
+        // Find `malloc`.
+        let malloc_index = match self.exports.get("malloc") {
+            None => bail!(
+                "found no exported `malloc` function to use to allocate the adapter's state in a TinyGo-derived module"
+            ),
+            Some(export) => {
+                if export.kind != ExternalKind::Func {
+                    bail!("found a `malloc` export, but it was not a function")
+                }
+                export.index
+            }
+        };
+
+        let mut func = wasm_encoder::Function::new([]);
+
+        func.instructions().local_get(3); // desired new size
+        func.instructions().call(malloc_index);
+
+        Ok(func)
+    }
+
     /// Encodes this `Module` to a new wasm module which is gc'd and only
     /// contains the items that are live as calculated by the `liveness` pass.
     fn encode(&mut self, main_module_realloc: Option<&str>) -> Result<Vec<u8>> {
@@ -733,10 +775,20 @@ impl<'a> Module<'a> {
 
         if sp.is_some() && (realloc_index.is_none() || allocation_state.is_none()) {
             // Either the main module does _not_ export a realloc function, or it is not safe to use for stack
-            // allocation because we have no way to short-circuit reentrance, so we'll use `memory.grow` instead.
+            // allocation because we have no way to short-circuit reentrance, so we'll provide our own realloc
+            // function instead.
             realloc_index = Some(num_func_imports + funcs.len());
             funcs.function(add_realloc_type(&mut types));
-            code.function(&realloc_via_memory_grow());
+
+            // If it appears the module was emitted by TinyGo, we delegate to its `malloc()` function. (TinyGo
+            // assumes its GC has reign over the whole memory and quickly overwrites the adapter's State struct
+            // unless we inform its GC of the memory we use.)
+            if self.is_from_tiny_go() {
+                code.function(&self.realloc_via_malloc()?);
+            } else {
+                // If it's not TinyGo, use `memory.grow` instead.
+                code.function(&realloc_via_memory_grow());
+            }
         }
 
         // Inject a start function to initialize the stack pointer which will be local to this module. This only