wasm2js: support memory imports properly, updating their buffer too (#2013)

* Emit an import statement for the memory.
* Update the imported memory's buffer when we grow.

3 files changed, 70 insertions, 5 deletions

diff --git a/src/wasm2js.h b/src/wasm2js.h
index 50968046f..d563409f9 100644
--- a/src/wasm2js.h
+++ b/src/wasm2js.h
@@ -279,7 +279,7 @@ private:
   void addExports(Ref ast, Module* wasm);
   void addGlobal(Ref ast, Global* global);
   void setNeedsAlmostASM(const char *reason);
-  void addMemoryGrowthFuncs(Ref ast);
+  void addMemoryGrowthFuncs(Ref ast, Module* wasm);
 
   Wasm2JSBuilder() = delete;
   Wasm2JSBuilder(const Wasm2JSBuilder &) = delete;
@@ -321,6 +321,18 @@ Ref Wasm2JSBuilder::processWasm(Module* wasm, Name funcName) {
   ValueBuilder::appendArgumentToFunction(asmFunc, ENV);
   ValueBuilder::appendArgumentToFunction(asmFunc, BUFFER);
   asmFunc[3]->push_back(ValueBuilder::makeStatement(ValueBuilder::makeString(USE_ASM)));
+  // add memory import
+  if (wasm->memory.exists && wasm->memory.imported()) {
+    Ref theVar = ValueBuilder::makeVar();
+    asmFunc[3]->push_back(theVar);
+    ValueBuilder::appendToVar(theVar,
+      "memory",
+      ValueBuilder::makeDot(
+        ValueBuilder::makeName(ENV),
+        ValueBuilder::makeName("memory")
+      )
+    );
+  }
   // create heaps, etc
   addBasics(asmFunc[3]);
   ModuleUtils::iterImportedFunctions(*wasm, [&](Function* import) {
@@ -589,7 +601,7 @@ void Wasm2JSBuilder::addExports(Ref ast, Module* wasm) {
   if (almostASM) {
     // replace "use asm"
     ast[0] = ValueBuilder::makeStatement(ValueBuilder::makeString(ALMOST_ASM));
-    addMemoryGrowthFuncs(ast);
+    addMemoryGrowthFuncs(ast, wasm);
   }
   ast->push_back(ValueBuilder::makeStatement(ValueBuilder::makeReturn(exports)));
 }
@@ -1849,7 +1861,7 @@ void Wasm2JSBuilder::setNeedsAlmostASM(const char *reason) {
   }
 }
 
-void Wasm2JSBuilder::addMemoryGrowthFuncs(Ref ast) {
+void Wasm2JSBuilder::addMemoryGrowthFuncs(Ref ast, Module* wasm) {
   Ref growMemoryFunc = ValueBuilder::makeFunction(WASM_GROW_MEMORY);
   ValueBuilder::appendArgumentToFunction(growMemoryFunc, IString("pagesToAdd"));
 
@@ -1978,6 +1990,20 @@ void Wasm2JSBuilder::addMemoryGrowthFuncs(Ref ast) {
     )
   );
 
+  // apply the changes to the memory import
+  if (wasm->memory.imported()) {
+    ValueBuilder::appendToBlock(block,
+      ValueBuilder::makeBinary(
+        ValueBuilder::makeDot(
+          ValueBuilder::makeName("memory"),
+          ValueBuilder::makeName(BUFFER)
+        ),
+        SET,
+        ValueBuilder::makeName(IString("newBuffer"))
+      )
+    );
+  }
+
   growMemoryFunc[3]->push_back(
     ValueBuilder::makeReturn(
       ValueBuilder::makeName(IString("oldPages"))));
diff --git a/test/wasm2js/emscripten.2asm.js b/test/wasm2js/emscripten.2asm.js
index 31b87cd44..23eca023e 100644
--- a/test/wasm2js/emscripten.2asm.js
+++ b/test/wasm2js/emscripten.2asm.js
@@ -1,7 +1,8 @@
 function instantiate(asmLibraryArg, wasmMemory, FUNCTION_TABLE) {
 
 function asmFunc(global, env, buffer) {
- "use asm";
+ "almost asm";
+ var memory = env.memory;
  var HEAP8 = new global.Int8Array(buffer);
  var HEAP16 = new global.Int16Array(buffer);
  var HEAP32 = new global.Int32Array(buffer);
@@ -47,12 +48,44 @@ function asmFunc(global, env, buffer) {
   
  }
  
+ function __growWasmMemory($0) {
+  $0 = $0 | 0;
+  return __wasm_grow_memory($0 | 0) | 0;
+ }
+ 
  // EMSCRIPTEN_END_FUNCS;
  FUNCTION_TABLE[1] = foo;
  FUNCTION_TABLE[2] = bar;
+ function __wasm_grow_memory(pagesToAdd) {
+  pagesToAdd = pagesToAdd | 0;
+  var oldPages = __wasm_current_memory() | 0;
+  var newPages = oldPages + pagesToAdd | 0;
+  if ((oldPages < newPages) && (newPages < 65536)) {
+   var newBuffer = new ArrayBuffer(Math_imul(newPages, 65536));
+   var newHEAP8 = new global.Int8Array(newBuffer);
+   newHEAP8.set(HEAP8);
+   HEAP8 = newHEAP8;
+   HEAP16 = new global.Int16Array(newBuffer);
+   HEAP32 = new global.Int32Array(newBuffer);
+   HEAPU8 = new global.Uint8Array(newBuffer);
+   HEAPU16 = new global.Uint16Array(newBuffer);
+   HEAPU32 = new global.Uint32Array(newBuffer);
+   HEAPF32 = new global.Float32Array(newBuffer);
+   HEAPF64 = new global.Float64Array(newBuffer);
+   buffer = newBuffer;
+   memory.buffer = newBuffer;
+  }
+  return oldPages;
+ }
+ 
+ function __wasm_current_memory() {
+  return buffer.byteLength / 65536 | 0;
+ }
+ 
  return {
   main: main, 
-  other: other
+  other: other, 
+  __growWasmMemory: __growWasmMemory
  };
 }
 
diff --git a/test/wasm2js/emscripten.wast b/test/wasm2js/emscripten.wast
index 82bf73548..dd11ee4b2 100644
--- a/test/wasm2js/emscripten.wast
+++ b/test/wasm2js/emscripten.wast
@@ -10,6 +10,7 @@
  (global $global$0 (mut i32) (i32.const 5243904))
  (export "main" (func $main))
  (export "other" (func $other))
+ (export "__growWasmMemory" (func $__growWasmMemory))
  (func $main
   (drop (call $syscall$6 (i32.const 1) (i32.const 2)))
   (drop (call $syscall$54 (i32.const 3) (i32.const 4)))
@@ -28,5 +29,10 @@
  (func $bar
   (nop)
  )
+ (func $__growWasmMemory (param $0 i32) (result i32)
+  (grow_memory
+   (local.get $0)
+  )
+ )
 )