From 66025636fd6bef4b3f717a75640f6c2bfca9eade Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Tue, 9 Aug 2022 12:59:53 -0500
Subject: [PATCH] Remove a layer of recursion in adapter compilation (#4657)

In #4640 a feature was added to adapter modules that whenever
translation goes through memory it instead goes through a helper
function as opposed to inlining it directly. The generation of the
helper function happened recursively at compile time, however, and sure
enough oss-fuzz has found an input which blows the host stack at compile
time.

This commit removes the compile-time recursion from the adapter compiler
when translating these helper functions by deferring the translation to
a worklist which is processed after the original function is translated.
This makes the stack-based recursion instead heap-based, removing the
stack overflow.
---
 crates/environ/src/fact.rs            | 29 +++++++++++++++++++++++++++
 crates/environ/src/fact/trampoline.rs | 26 ++++++------------------
 2 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/crates/environ/src/fact.rs b/crates/environ/src/fact.rs
index 19f8a899be..895db3ea5c 100644
--- a/crates/environ/src/fact.rs
+++ b/crates/environ/src/fact.rs
@@ -66,6 +66,7 @@ pub struct Module<'a> {
 
     funcs: PrimaryMap<FunctionId, Function>,
     translate_mem_funcs: HashMap<(InterfaceType, InterfaceType, Options, Options), FunctionId>,
+    translate_mem_worklist: Vec<(FunctionId, InterfaceType, InterfaceType, Options, Options)>,
 }
 
 struct AdapterData {
@@ -138,6 +139,7 @@ impl<'a> Module<'a> {
             imported_globals: PrimaryMap::new(),
             funcs: PrimaryMap::new(),
             translate_mem_funcs: HashMap::new(),
+            translate_mem_worklist: Vec::new(),
         }
     }
 
@@ -185,6 +187,10 @@ impl<'a> Module<'a> {
                 called_as_export: true,
             },
         );
+
+        while let Some((result, src, dst, src_opts, dst_opts)) = self.translate_mem_worklist.pop() {
+            trampoline::compile_translate_mem(self, result, src, &src_opts, dst, &dst_opts);
+        }
     }
 
     fn import_options(&mut self, ty: TypeFuncIndex, options: &AdapterOptionsDfg) -> AdapterOptions {
@@ -315,6 +321,29 @@ impl<'a> Module<'a> {
             })
     }
 
+    fn translate_mem(
+        &mut self,
+        src: InterfaceType,
+        src_opts: &Options,
+        dst: InterfaceType,
+        dst_opts: &Options,
+    ) -> FunctionId {
+        *self
+            .translate_mem_funcs
+            .entry((src, dst, *src_opts, *dst_opts))
+            .or_insert_with(|| {
+                // Generate a fresh `Function` with a unique id for what we're about to
+                // generate.
+                let ty = self
+                    .core_types
+                    .function(&[src_opts.ptr(), dst_opts.ptr()], &[]);
+                let id = self.funcs.push(Function::new(None, ty));
+                self.translate_mem_worklist
+                    .push((id, src, dst, *src_opts, *dst_opts));
+                id
+            })
+    }
+
     /// Encodes this module into a WebAssembly binary.
     pub fn encode(&mut self) -> Vec<u8> {
         // Build the function/export sections of the wasm module in a first pass
diff --git a/crates/environ/src/fact/trampoline.rs b/crates/environ/src/fact/trampoline.rs
index c00337f450..6e3fd62c1a 100644
--- a/crates/environ/src/fact/trampoline.rs
+++ b/crates/environ/src/fact/trampoline.rs
@@ -88,28 +88,14 @@ pub(super) fn compile(module: &mut Module<'_>, adapter: &AdapterData) {
 /// The generated function takes two arguments: the source pointer and
 /// destination pointer. The conversion operation is configured by the
 /// `src_opts` and `dst_opts` specified as well.
-fn compile_translate_mem(
+pub(super) fn compile_translate_mem(
     module: &mut Module<'_>,
+    result: FunctionId,
     src: InterfaceType,
     src_opts: &Options,
     dst: InterfaceType,
     dst_opts: &Options,
-) -> FunctionId {
-    // If a helper for this translation has already been generated then reuse
-    // that. Note that this is key to this function where by doing this it
-    // prevents an exponentially sized output given any particular input type.
-    let key = (src, dst, *src_opts, *dst_opts);
-    if module.translate_mem_funcs.contains_key(&key) {
-        return module.translate_mem_funcs[&key];
-    }
-
-    // Generate a fresh `Function` with a unique id for what we're about to
-    // generate.
-    let ty = module
-        .core_types
-        .function(&[src_opts.ptr(), dst_opts.ptr()], &[]);
-    let result = module.funcs.push(Function::new(None, ty));
-    module.translate_mem_funcs.insert(key, result);
+) {
     let mut compiler = Compiler {
         types: module.types,
         module,
@@ -138,7 +124,6 @@ fn compile_translate_mem(
         }),
     );
     compiler.finish();
-    result
 }
 
 /// Possible ways that a interface value is represented in the core wasm
@@ -486,8 +471,9 @@ impl Compiler<'_, '_> {
                 // because we don't know the final index of the generated
                 // function yet. It's filled in at the end of adapter module
                 // translation.
-                let helper =
-                    compile_translate_mem(self.module, *src_ty, src.opts, *dst_ty, dst.opts);
+                let helper = self
+                    .module
+                    .translate_mem(*src_ty, src.opts, *dst_ty, dst.opts);
 
                 // TODO: overflow checks?
                 self.instruction(LocalGet(src.addr.idx));