wasm2js: support non-constant indexes for memory and table segments (#2055)

Mostly what we need for dynamic linking, at least on the binaryen side.

4 files changed, 134 insertions, 37 deletions

diff --git a/scripts/test/env.js b/scripts/test/env.js
index 44e6035de..af2cce1c7 100644
--- a/scripts/test/env.js
+++ b/scripts/test/env.js
@@ -8,3 +8,5 @@ export function getTempRet0() {
   return tempRet0;
 }
 
+export const memoryBase = 0;
+export const tableBase = 0;
diff --git a/src/wasm2js.h b/src/wasm2js.h
index fdfa77b84..6c97c5a6b 100644
--- a/src/wasm2js.h
+++ b/src/wasm2js.h
@@ -79,6 +79,10 @@ void sequenceAppend(Ref& ast, Ref extra) {
   ast = ValueBuilder::makeSeq(ast, extra);
 }
 
+IString stringToIString(std::string str) {
+  return IString(str.c_str(), false);
+}
+
 // Used when taking a wasm name and generating a JS identifier. Each scope here
 // is used to ensure that all names have a unique name but the same wasm name
 // within a scope always resolves to the same symbol.
@@ -89,16 +93,6 @@ enum class NameScope {
   Max,
 };
 
-template<typename T>
-static uint64_t constOffset(const T& segment) {
-  auto* c = segment.offset->template dynCast<Const>();
-  if (!c) {
-    Fatal() << "non-constant offsets aren't supported yet\n";
-    abort();
-  }
-  return c->value.getInteger();
-}
-
 //
 // Wasm2JSBuilder - converts a WebAssembly module's functions into JS
 //
@@ -196,7 +190,7 @@ public:
         out << "_" << i;
       }
       auto mangled = asmangle(out.str());
-      ret = IString(mangled.c_str(), false);
+      ret = stringToIString(mangled);
       if (!allMangledNames.count(ret)) {
         break;
       }
@@ -219,10 +213,6 @@ public:
     return ret;
   }
 
-  size_t getTableSize() {
-    return tableSize;
-  }
-
 private:
   Flags flags;
   PassOptions options;
@@ -237,8 +227,6 @@ private:
   std::unordered_map<const char*, IString> mangledNames[(int) NameScope::Max];
   std::unordered_set<IString> allMangledNames;
 
-  size_t tableSize;
-
   // If a function is callable from outside, we'll need to cast the inputs
   // and our return value. Otherwise, internally, casts are only needed
   // on operations.
@@ -331,17 +319,6 @@ Ref Wasm2JSBuilder::processWasm(Module* wasm, Name funcName) {
   ModuleUtils::iterImportedGlobals(*wasm, [&](Global* import) {
     addGlobalImport(asmFunc[3], import);
   });
-  // figure out the table size
-  tableSize = std::accumulate(wasm->table.segments.begin(),
-                              wasm->table.segments.end(),
-                              0, [&](size_t size, Table::Segment seg) -> size_t {
-                                return size + seg.data.size() + constOffset(seg);
-                              });
-  size_t pow2ed = 1;
-  while (pow2ed < tableSize) {
-    pow2ed <<= 1;
-  }
-  tableSize = pow2ed;
 
   // make sure exports get their expected names
   for (auto& e : wasm->exports) {
@@ -509,8 +486,7 @@ void Wasm2JSBuilder::addTable(Ref ast, Module* wasm) {
   // Emit a simple flat table as a JS array literal. Otherwise,
   // emit assignments separately for each index.
   FlatTable flat(wasm->table);
-  assert(flat.valid); // TODO: non-flat tables
-  if (!wasm->table.imported()) {
+  if (flat.valid && !wasm->table.imported()) {
     Ref theVar = ValueBuilder::makeVar();
     ast->push_back(theVar);
     Ref theArray = ValueBuilder::makeArray();
@@ -525,16 +501,33 @@ void Wasm2JSBuilder::addTable(Ref ast, Module* wasm) {
       ValueBuilder::appendToArray(theArray, ValueBuilder::makeName(name));
     }
   } else {
+    if (!wasm->table.imported()) {
+      Ref theVar = ValueBuilder::makeVar();
+      ast->push_back(theVar);
+      ValueBuilder::appendToVar(theVar, FUNCTION_TABLE, ValueBuilder::makeArray());
+    }
+
     // TODO: optimize for size
     for (auto& segment : wasm->table.segments) {
       auto offset = segment.offset;
-      Index start = offset->cast<Const>()->value.geti32();
       for (Index i = 0; i < segment.data.size(); i++) {
+        Ref index;
+        if (auto* c = offset->dynCast<Const>()) {
+          index = ValueBuilder::makeInt(c->value.geti32() + i);
+        } else if (auto* get = offset->dynCast<GetGlobal>()) {
+          index = ValueBuilder::makeBinary(
+            ValueBuilder::makeName(stringToIString(asmangle(get->name.str))),
+            PLUS,
+            ValueBuilder::makeNum(i)
+          );
+        } else {
+          WASM_UNREACHABLE();
+        }
         ast->push_back(ValueBuilder::makeStatement(
           ValueBuilder::makeBinary(
             ValueBuilder::makeSub(
               ValueBuilder::makeName(FUNCTION_TABLE),
-              ValueBuilder::makeInt(start + i)
+              index
             ),
             SET,
             ValueBuilder::makeName(fromName(segment.data[i], NameScope::Top))
@@ -1794,7 +1787,7 @@ private:
   void emitPostEmscripten();
   void emitPostES6();
 
-  void emitMemory(std::string buffer, std::string segmentWriter);
+  void emitMemory(std::string buffer, std::string segmentWriter, std::function<std::string (std::string)> accessGlobal);
   void emitScratchMemorySupport();
 };
 
@@ -1857,7 +1850,9 @@ void Wasm2JSGlue::emitPost() {
 }
 
 void Wasm2JSGlue::emitPostEmscripten() {
-  emitMemory("wasmMemory.buffer", "writeSegment");
+  emitMemory("wasmMemory.buffer", "writeSegment", [](std::string globalName) {
+    return std::string("asmLibraryArg['") + asmangle(globalName) + "']";
+  });
 
   out << "return asmFunc({\n"
       << "    'Int8Array': Int8Array,\n"
@@ -1895,7 +1890,8 @@ void Wasm2JSGlue::emitPostES6() {
   }
 
   emitMemory(std::string("mem") + moduleName.str,
-             std::string("assign") + moduleName.str);
+             std::string("assign") + moduleName.str,
+             [](std::string globalName) { return globalName; });
 
   // Actually invoke the `asmFunc` generated function, passing in all global
   // values followed by all imports
@@ -1958,7 +1954,7 @@ void Wasm2JSGlue::emitPostES6() {
   }
 }
 
-void Wasm2JSGlue::emitMemory(std::string buffer, std::string segmentWriter) {
+void Wasm2JSGlue::emitMemory(std::string buffer, std::string segmentWriter, std::function<std::string (std::string)> accessGlobal) {
   if (wasm.memory.segments.empty()) return;
 
   auto expr = R"(
@@ -1982,10 +1978,22 @@ void Wasm2JSGlue::emitMemory(std::string buffer, std::string segmentWriter) {
   out << "var " << segmentWriter
       << " = (" << expr << ")(" << buffer << ");\n";
 
+  auto globalOffset = [&](const Memory::Segment& segment) {
+    if (auto* c = segment.offset->template dynCast<Const>()) {;
+      return std::to_string(c->value.getInteger());
+    }
+    if (auto* get = segment.offset->template dynCast<GetGlobal>()) {
+      auto internalName = get->name;
+      auto importedName = wasm.getGlobal(internalName)->base;
+      return accessGlobal(asmangle(importedName.str));
+    }
+    Fatal() << "non-constant offsets aren't supported yet\n";
+  };
+
   for (auto& seg : wasm.memory.segments) {
     assert(!seg.isPassive && "passive segments not implemented yet");
     out << segmentWriter << "("
-      << constOffset(seg)
+      << globalOffset(seg)
       << ", \""
       << base64Encode(seg.data)
       << "\");\n";
diff --git a/test/wasm2js/dynamicLibrary.2asm.js b/test/wasm2js/dynamicLibrary.2asm.js
new file mode 100644
index 000000000..55c6d6500
--- /dev/null
+++ b/test/wasm2js/dynamicLibrary.2asm.js
@@ -0,0 +1,68 @@
+import { memoryBase } from 'env';
+import { tableBase } from 'env';
+
+function asmFunc(global, env, buffer) {
+ "almost asm";
+ var memory = env.memory;
+ var HEAP8 = new global.Int8Array(buffer);
+ var HEAP16 = new global.Int16Array(buffer);
+ var HEAP32 = new global.Int32Array(buffer);
+ var HEAPU8 = new global.Uint8Array(buffer);
+ var HEAPU16 = new global.Uint16Array(buffer);
+ var HEAPU32 = new global.Uint32Array(buffer);
+ var HEAPF32 = new global.Float32Array(buffer);
+ var HEAPF64 = new global.Float64Array(buffer);
+ var Math_imul = global.Math.imul;
+ var Math_fround = global.Math.fround;
+ var Math_abs = global.Math.abs;
+ var Math_clz32 = global.Math.clz32;
+ var Math_min = global.Math.min;
+ var Math_max = global.Math.max;
+ var Math_floor = global.Math.floor;
+ var Math_ceil = global.Math.ceil;
+ var Math_sqrt = global.Math.sqrt;
+ var abort = env.abort;
+ var nan = global.NaN;
+ var infinity = global.Infinity;
+ var import$memoryBase = env.memoryBase | 0;
+ var import$tableBase = env.tableBase | 0;
+ function foo() {
+  
+ }
+ 
+ function bar() {
+  
+ }
+ 
+ function baz() {
+  
+ }
+ 
+ var FUNCTION_TABLE = [];
+ FUNCTION_TABLE[import$tableBase + 0] = foo;
+ FUNCTION_TABLE[import$tableBase + 1] = bar;
+ return {
+  "baz": baz
+ };
+}
+
+var memasmFunc = new ArrayBuffer(16777216);
+var assignasmFunc = (
+    function(mem) {
+      var _mem = new Uint8Array(mem);
+      return function(offset, s) {
+        if (typeof Buffer === 'undefined') {
+          var bytes = atob(s);
+          for (var i = 0; i < bytes.length; i++)
+            _mem[offset + i] = bytes.charCodeAt(i);
+        } else {
+          var bytes = Buffer.from(s, 'base64');
+          for (var i = 0; i < bytes.length; i++)
+            _mem[offset + i] = bytes[i];
+        }
+      }
+    }
+  )(memasmFunc);
+assignasmFunc(memoryBase, "ZHluYW1pYyBkYXRh");
+var retasmFunc = asmFunc({Math,Int8Array,Uint8Array,Int16Array,Uint16Array,Int32Array,Uint32Array,Float32Array,Float64Array,NaN,Infinity}, {abort:function() { throw new Error('abort'); }},memasmFunc);
+export var baz = retasmFunc.baz;
diff --git a/test/wasm2js/dynamicLibrary.wast b/test/wasm2js/dynamicLibrary.wast
new file mode 100644
index 000000000..28175633c
--- /dev/null
+++ b/test/wasm2js/dynamicLibrary.wast
@@ -0,0 +1,19 @@
+(module
+ (type $FUNCSIG$vi (func (param i32)))
+ (type $FUNCSIG$ii (func (param i32) (result i32)))
+
+ (import "env" "memory" (memory $import$memory 256 256))
+ (import "env" "memoryBase" (global $import$memoryBase i32))
+ (data (global.get $import$memoryBase) "dynamic data")
+
+ (table 10 10 funcref)
+ (import "env" "tableBase" (global $import$tableBase i32))
+ (elem (global.get $import$tableBase) $foo $bar)
+
+ (export "baz" (func $baz))
+ (export "tab" (table 0))
+
+ (func $foo)
+ (func $bar)
+ (func $baz)
+)