3 files changed, 27 insertions, 25 deletions

diff --git a/src/asm2wasm-main.cpp b/src/asm2wasm-main.cpp
index 1e9dc6680..aca9434bb 100644
--- a/src/asm2wasm-main.cpp
+++ b/src/asm2wasm-main.cpp
@@ -37,7 +37,7 @@ int main(int argc, char **argv) {
   if (debug) std::cerr << "wasming...\n";
   Module wasm;
   wasm.memory.initial = wasm.memory.max = 16*1024*1024; // we would normally receive this from the compiler
-  Asm2WasmBuilder asm2wasm(wasm);
+  Asm2WasmBuilder asm2wasm(wasm, pre.memoryGrowth);
   asm2wasm.processAsm(asmjs);
 
   if (debug) std::cerr << "optimizing...\n";
diff --git a/src/asm2wasm.h b/src/asm2wasm.h
index 187bfeee6..6cbb4e084 100644
--- a/src/asm2wasm.h
+++ b/src/asm2wasm.h
@@ -39,8 +39,9 @@ IString GLOBAL("global"), NAN_("NaN"), INFINITY_("Infinity"),
         FLOOR("floor"),
         SQRT("sqrt"),
         I32_TEMP("asm2wasm_i32_temp"),
-        DEBUGGER("debugger");
-
+        DEBUGGER("debugger"),
+        GROW_WASM_MEMORY("__growWasmMemory"),
+        NEW_SIZE("newSize");
 
 static void abort_on(std::string why) {
   std::cerr << why << '\n';
@@ -163,6 +164,8 @@ class Asm2WasmBuilder {
   std::map<IString, int> functionTableStarts; // each asm function table gets a range in the one wasm table, starting at a location
   std::map<CallIndirect*, IString> callIndirects; // track these, as we need to fix them after we know the functionTableStarts. this maps call => its function table
 
+  bool memoryGrowth;
+
 public:
   std::map<IString, MappedGlobal> mappedGlobals;
 
@@ -269,7 +272,7 @@ private:
   }
 
 public:
-  Asm2WasmBuilder(Module& wasm) : wasm(wasm), nextGlobal(8), maxGlobal(1000) {}
+  Asm2WasmBuilder(Module& wasm, bool memoryGrowth) : wasm(wasm), nextGlobal(8), maxGlobal(1000), memoryGrowth(memoryGrowth) {}
 
   void processAsm(Ref ast);
   void optimize();
@@ -740,6 +743,25 @@ void Asm2WasmBuilder::processAsm(Ref ast) {
     sub->type = WasmType::i32;
     call->target = sub;
   }
+
+  // apply memory growth, if relevant
+  if (memoryGrowth) {
+    // create and export a function that just calls memory growth
+    auto growWasmMemory = allocator.alloc<Function>();
+    growWasmMemory->name = GROW_WASM_MEMORY;
+    growWasmMemory->params.emplace_back(NEW_SIZE, i32); // the new size
+    auto get = allocator.alloc<GetLocal>();
+    get->name = NEW_SIZE;
+    auto grow = allocator.alloc<Host>();
+    grow->op = GrowMemory;
+    grow->operands.push_back(get);
+    growWasmMemory->body = grow;
+    wasm.addFunction(growWasmMemory);
+    auto export_ = allocator.alloc<Export>();
+    export_->name = export_->value = GROW_WASM_MEMORY;
+    wasm.addExport(export_);
+  }
+
 }
 
 Function* Asm2WasmBuilder::processFunction(Ref ast) {
diff --git a/src/wasm-js.cpp b/src/wasm-js.cpp
index fbd3031ee..d0951f671 100644
--- a/src/wasm-js.cpp
+++ b/src/wasm-js.cpp
@@ -14,9 +14,6 @@
 using namespace cashew;
 using namespace wasm;
 
-IString GROW_WASM_MEMORY("__growWasmMemory"),
-        NEW_SIZE("newSize");
-
 // global singletons
 Asm2WasmBuilder* asm2wasm = nullptr;
 ModuleInstance* instance = nullptr;
@@ -52,26 +49,9 @@ extern "C" void EMSCRIPTEN_KEEPALIVE load_asm(char *input) {
   module->memory.max = pre.memoryGrowth ? -1 : module->memory.initial;
 
   if (wasmJSDebug) std::cerr << "wasming...\n";
-  asm2wasm = new Asm2WasmBuilder(*module);
+  asm2wasm = new Asm2WasmBuilder(*module, pre.memoryGrowth);
   asm2wasm->processAsm(asmjs);
 
-  if (pre.memoryGrowth) {
-    // create and export a function that just calls memory growth
-    auto growWasmMemory = new Function();
-    growWasmMemory->name = GROW_WASM_MEMORY;
-    growWasmMemory->params.emplace_back(NEW_SIZE, i32); // the new size
-    auto get = new GetLocal();
-    get->name = NEW_SIZE;
-    auto grow = new Host();
-    grow->op = GrowMemory;
-    grow->operands.push_back(get);
-    growWasmMemory->body = grow;
-    module->addFunction(growWasmMemory);
-    auto export_ = new Export();
-    export_->name = export_->value = GROW_WASM_MEMORY;
-    module->addExport(export_);
-  }
-
   if (wasmJSDebug) std::cerr << "optimizing...\n";
   asm2wasm->optimize();