offset support in table

14 files changed, 204 insertions, 66 deletions

diff --git a/src/asm2wasm.h b/src/asm2wasm.h
index 868aa31d3..a5754f725 100644
--- a/src/asm2wasm.h
+++ b/src/asm2wasm.h
@@ -660,13 +660,17 @@ void Asm2WasmBuilder::processAsm(Ref ast) {
           // TODO: when not using aliasing function pointers, we could merge them by noticing that
           //       index 0 in each table is the null func, and each other index should only have one
           //       non-null func. However, that breaks down when function pointer casts are emulated.
-          functionTableStarts[name] = wasm.table.names.size(); // this table starts here
+          if (wasm.table.segments.size() == 0) {
+            wasm.table.segments.emplace_back(wasm.allocator.alloc<Const>()->set(Literal(uint32_t(0))));
+          }
+          auto& segment = wasm.table.segments[0];
+          functionTableStarts[name] = segment.data.size(); // this table starts here
           Ref contents = value[1];
           for (unsigned k = 0; k < contents->size(); k++) {
             IString curr = contents[k][1]->getIString();
-            wasm.table.names.push_back(curr);
+            segment.data.push_back(curr);
           }
-          wasm.table.initial = wasm.table.max = wasm.table.names.size();
+          wasm.table.initial = wasm.table.max = segment.data.size();
         } else {
           abort_on("invalid var element", pair);
         }
diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
index d7f866663..83c626eb1 100644
--- a/src/binaryen-c.cpp
+++ b/src/binaryen-c.cpp
@@ -729,10 +729,12 @@ void BinaryenSetFunctionTable(BinaryenModuleRef module, BinaryenFunctionRef* fun
   }
 
   auto* wasm = (Module*)module;
+  Table::Segment segment(wasm->allocator.alloc<Const>()->set(Literal(int32_t(0))));
   for (BinaryenIndex i = 0; i < numFuncs; i++) {
-    wasm->table.names.push_back(((Function*)funcs[i])->name);
+    segment.data.push_back(((Function*)funcs[i])->name);
   }
-  wasm->table.initial = wasm->table.max = wasm->table.names.size();
+  wasm->table.segments.push_back(segment);
+  wasm->table.initial = wasm->table.max = numFuncs;
 }
 
 // Memory. One per module
diff --git a/src/passes/DuplicateFunctionElimination.cpp b/src/passes/DuplicateFunctionElimination.cpp
index 961d26ba5..2b8e69b54 100644
--- a/src/passes/DuplicateFunctionElimination.cpp
+++ b/src/passes/DuplicateFunctionElimination.cpp
@@ -123,10 +123,12 @@ struct DuplicateFunctionElimination : public Pass {
         replacerRunner.add<FunctionReplacer>(&replacements);
         replacerRunner.run();
         // replace in table
-        for (auto& name : module->table.names) {
-          auto iter = replacements.find(name);
-          if (iter != replacements.end()) {
-            name = iter->second;
+        for (auto& segment : module->table.segments) {
+          for (auto& name : segment.data) {
+            auto iter = replacements.find(name);
+            if (iter != replacements.end()) {
+              name = iter->second;
+            }
           }
         }
         // replace in start
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index 711cc8a84..5eea38bdc 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -579,12 +579,15 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     if (curr->max && curr->max != Table::kMaxSize) o << ' ' << curr->max;
     o << " anyfunc)\n";
     doIndent(o, indent);
-    printOpening(o, "elem", true);
-    for (auto name : curr->names) {
-      o << ' ';
-      printName(name);
+    for (auto& segment : curr->segments) {
+      printOpening(o, "elem ", true);
+      visit(segment.offset);
+      for (auto name : segment.data) {
+        o << ' ';
+        printName(name);
+      }
+      o << ')';
     }
-    o << ')';
   }
   void visitModule(Module *curr) {
     currModule = curr;
@@ -652,7 +655,7 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
       visitGlobal(child.get());
       o << maybeNewLine;
     }
-    if (curr->table.names.size() > 0) {
+    if (curr->table.segments.size() > 0 || curr->table.initial > 0 || curr->table.max != Table::kMaxSize) {
       doIndent(o, indent);
       visitTable(&curr->table);
       o << maybeNewLine;
diff --git a/src/passes/RemoveUnusedFunctions.cpp b/src/passes/RemoveUnusedFunctions.cpp
index a2941aff6..78b0f0ffc 100644
--- a/src/passes/RemoveUnusedFunctions.cpp
+++ b/src/passes/RemoveUnusedFunctions.cpp
@@ -39,8 +39,10 @@ struct RemoveUnusedFunctions : public Pass {
       root.push_back(module->getFunction(curr->value));
     }
     // For now, all functions that can be called indirectly are marked as roots.
-    for (auto& curr : module->table.names) {
-      root.push_back(module->getFunction(curr));
+    for (auto& segment : module->table.segments) {
+      for (auto& curr : segment.data) {
+        root.push_back(module->getFunction(curr));
+      }
     }
     // Compute function reachability starting from the root set.
     DirectCallGraphAnalyzer analyzer(module, root);
diff --git a/src/passes/ReorderFunctions.cpp b/src/passes/ReorderFunctions.cpp
index 38ef98afb..679fedb61 100644
--- a/src/passes/ReorderFunctions.cpp
+++ b/src/passes/ReorderFunctions.cpp
@@ -38,8 +38,10 @@ struct ReorderFunctions : public WalkerPass<PostWalker<ReorderFunctions, Visitor
     for (auto& curr : module->exports) {
       counts[curr->value]++;
     }
-    for (auto& curr : module->table.names) {
-      counts[curr]++;
+    for (auto& segment : module->table.segments) {
+      for (auto& curr : segment.data) {
+        counts[curr]++;
+      }
     }
     std::sort(module->functions.begin(), module->functions.end(), [this](
       const std::unique_ptr<Function>& a,
diff --git a/src/shell-interface.h b/src/shell-interface.h
index 55f9b2a0b..f332307ad 100644
--- a/src/shell-interface.h
+++ b/src/shell-interface.h
@@ -86,6 +86,8 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
     }
   } memory;
 
+  std::vector<Name> table;
+
   ShellExternalInterface() : memory() {}
 
   void init(Module& wasm) override {
@@ -98,6 +100,15 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
         memory.set(offset + i, segment.data[i]);
       }
     }
+
+    table.resize(wasm.table.initial);
+    for (auto& segment : wasm.table.segments) {
+      Address offset = ConstantExpressionRunner().visit(segment.offset).value.geti32();
+      assert(offset + segment.data.size() <= wasm.table.initial);
+      for (size_t i = 0; i != segment.data.size(); ++i) {
+        table[offset + i] = segment.data[i];
+      }
+    }
   }
 
   Literal callImport(Import *import, LiteralList& arguments) override {
@@ -115,6 +126,19 @@ struct ShellExternalInterface : ModuleInstance::ExternalInterface {
     abort();
   }
 
+  Literal callTable(Index index, Name type, LiteralList& arguments, ModuleInstance& instance) override {
+    if (index >= table.size()) trap("callTable overflow");
+    auto* func = instance.wasm.getFunction(table[index]);
+    if (func->type.is() && func->type != type) trap("callIndirect: bad type");
+    if (func->params.size() != arguments.size()) trap("callIndirect: bad # of arguments");
+    for (size_t i = 0; i < func->params.size(); i++) {
+      if (func->params[i] != arguments[i].type) {
+        trap("callIndirect: bad argument type");
+      }
+    }
+    return instance.callFunctionInternal(func->name, arguments);
+  }
+
   Literal load(Load* load, Address addr) override {
     switch (load->type) {
       case i32: {
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index d405c91dc..04bebddc5 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -740,12 +740,19 @@ public:
   }
 
   void writeFunctionTable() {
-    if (wasm->table.names.size() == 0) return;
+    if (wasm->table.segments.size() == 0) return;
     if (debug) std::cerr << "== writeFunctionTable" << std::endl;
     auto start = startSection(BinaryConsts::Section::FunctionTable);
-    o << U32LEB(wasm->table.names.size());
-    for (auto name : wasm->table.names) {
-      o << U32LEB(getFunctionIndex(name));
+    o << U32LEB(wasm->table.initial);
+    o << U32LEB(wasm->table.max);
+    o << U32LEB(wasm->table.segments.size());
+    for (auto& segment : wasm->table.segments) {
+      writeExpression(segment.offset);
+      o << int8_t(BinaryConsts::End);
+      o << U32LEB(segment.data.size());
+      for (auto name : segment.data) {
+        o << U32LEB(getFunctionIndex(name));
+      }
     }
     finishSection(start);
   }
@@ -1644,11 +1651,13 @@ public:
       }
     }
 
-    for (size_t index : functionTable) {
-      assert(index < wasm.functions.size());
-      wasm.table.names.push_back(wasm.functions[index]->name);
+    for (auto& pair : functionTable) {
+      auto i = pair.first;
+      auto& indexes = pair.second;
+      for (auto j : indexes) {
+        wasm.table.segments[i].data.push_back(wasm.functions[j]->name);
+      }
     }
-    wasm.table.initial = wasm.table.max = wasm.table.names.size();
   }
 
   void readDataSegments() {
@@ -1667,14 +1676,20 @@ public:
     }
   }
 
-  std::vector<size_t> functionTable;
+  std::map<Index, std::vector<Index>> functionTable;
 
   void readFunctionTable() {
     if (debug) std::cerr << "== readFunctionTable" << std::endl;
+    wasm.table.initial = getU32LEB();
+    wasm.table.max = getU32LEB();
     auto num = getU32LEB();
     for (size_t i = 0; i < num; i++) {
-      auto index = getU32LEB();
-      functionTable.push_back(index);
+      wasm.table.segments.emplace_back(readExpression());
+      auto& temporary = functionTable[i];
+      auto size = getU32LEB();
+      for (Index j = 0; j < size; j++) {
+        temporary.push_back(getU32LEB());
+      }
     }
   }
 
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 0be36a677..f2deed039 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -533,6 +533,7 @@ public:
   struct ExternalInterface {
     virtual void init(Module& wasm) {}
     virtual Literal callImport(Import* import, LiteralList& arguments) = 0;
+    virtual Literal callTable(Index index, Name type, LiteralList& arguments, ModuleInstance& instance) = 0;
     virtual Literal load(Load* load, Address addr) = 0;
     virtual void store(Store* store, Address addr, Literal value) = 0;
     virtual void growMemory(Address oldSize, Address newSize) = 0;
@@ -591,8 +592,8 @@ private:
     return callFunctionInternal(name, arguments);
   }
 
-private:
-  // Internal function call.
+public:
+  // Internal function call. Must be public so that callTable implementations can use it (refactor?)
   Literal callFunctionInternal(IString name, LiteralList& arguments) {
 
     class FunctionScope {
@@ -672,18 +673,8 @@ private:
         LiteralList arguments;
         Flow flow = generateArguments(curr->operands, arguments);
         if (flow.breaking()) return flow;
-        size_t index = target.value.geti32();
-        if (index >= instance.wasm.table.names.size()) trap("callIndirect: overflow");
-        Name name = instance.wasm.table.names[index];
-        Function *func = instance.wasm.getFunction(name);
-        if (func->type.is() && func->type != curr->fullType) trap("callIndirect: bad type");
-        if (func->params.size() != arguments.size()) trap("callIndirect: bad # of arguments");
-        for (size_t i = 0; i < func->params.size(); i++) {
-          if (func->params[i] != arguments[i].type) {
-            trap("callIndirect: bad argument type");
-          }
-        }
-        return instance.callFunctionInternal(name, arguments);
+        Index index = target.value.geti32();
+        return instance.externalInterface->callTable(index, curr->fullType, arguments, instance);
       }
 
       Flow visitGetLocal(GetLocal *curr) {
@@ -802,6 +793,8 @@ private:
     return ret;
   }
 
+private:
+
   Address memorySize; // in pages
 
   template <class LS>
diff --git a/src/wasm-js.cpp b/src/wasm-js.cpp
index 2aacd94d4..83956e47e 100644
--- a/src/wasm-js.cpp
+++ b/src/wasm-js.cpp
@@ -195,10 +195,23 @@ extern "C" void EMSCRIPTEN_KEEPALIVE instantiate() {
           target.set(source, $0);
         }, ConstantExpressionRunner().visit(segment.offset).value.geti32(), &segment.data[0], segment.data.size());
       }
+      // Table support is in a JS array. If the entry is a number, it's a function pointer. If not, it's a JS method to be called directly
+      // TODO: make them all JS methods, wrapping a dynCall where necessary?
+      EM_ASM_({
+        Module['outside']['wasmTable'] = new Array($0);
+      }, wasm.table.initial);
+      for (auto segment : wasm.table.segments) {
+        Address offset = ConstantExpressionRunner().visit(segment.offset).value.geti32();
+        assert(offset + segment.data.size() <= wasm.table.initial);
+        for (size_t i = 0; i != segment.data.size(); ++i) {
+          EM_ASM_({
+            Module['outside']['wasmTable'][$0] = $1;
+          }, offset + i, wasm.getFunction(segment.data[i]));
+        }
+      }
     }
 
-    Literal callImport(Import *import, LiteralList& arguments) override {
-      if (wasmJSDebug) std::cout << "calling import " << import->name.str << '\n';
+    void prepareTempArgments(LiteralList& arguments) {
       EM_ASM({
         Module['tempArguments'] = [];
       });
@@ -213,6 +226,21 @@ extern "C" void EMSCRIPTEN_KEEPALIVE instantiate() {
           abort();
         }
       }
+    }
+
+    Literal getResultFromJS(double ret, WasmType type) {
+      switch (type) {
+        case none: return Literal(0);
+        case i32: return Literal((int32_t)ret);
+        case f32: return Literal((float)ret);
+        case f64: return Literal((double)ret);
+        default: abort();
+      }
+    }
+
+    Literal callImport(Import *import, LiteralList& arguments) override {
+      if (wasmJSDebug) std::cout << "calling import " << import->name.str << '\n';
+      prepareTempArgments(arguments);
       double ret = EM_ASM_DOUBLE({
         var mod = Pointer_stringify($0);
         var base = Pointer_stringify($1);
@@ -224,12 +252,36 @@ extern "C" void EMSCRIPTEN_KEEPALIVE instantiate() {
 
       if (wasmJSDebug) std::cout << "calling import returning " << ret << '\n';
 
-      switch (import->type->result) {
-        case none: return Literal(0);
-        case i32: return Literal((int32_t)ret);
-        case f32: return Literal((float)ret);
-        case f64: return Literal((double)ret);
-        default: abort();
+      return getResultFromJS(ret, import->type->result);
+    }
+
+    Literal callTable(Index index, Name type, LiteralList& arguments, ModuleInstance& instance) override {
+      void* ptr = (void*)EM_ASM_INT({
+        var value = Module['outside']['wasmTable'][$0];
+        return typeof value === "number" ? value : -1;
+      }, index);
+      if (ptr == nullptr) trap("callTable overflow");
+      if (ptr != (void*)-1) {
+        // a Function we can call
+        Function* func = (Function*)ptr;
+        if (func->type.is() && func->type != type) trap("callIndirect: bad type");
+        if (func->params.size() != arguments.size()) trap("callIndirect: bad # of arguments");
+        for (size_t i = 0; i < func->params.size(); i++) {
+          if (func->params[i] != arguments[i].type) {
+            trap("callIndirect: bad argument type");
+          }
+        }
+        return instance.callFunctionInternal(func->name, arguments);
+      } else {
+        // A JS function JS can call
+        prepareTempArgments(arguments);
+        double ret = EM_ASM_DOUBLE({
+          var func = Module['outside']['wasmTable'][$0];
+          var tempArguments = Module['tempArguments'];
+          Module['tempArguments'] = null;
+          return func.apply(null, tempArguments);
+        }, index);
+        return getResultFromJS(ret, instance.wasm.getFunctionType(type)->result);
       }
     }
 
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index 29f439bfb..2c9a4cd4d 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -132,7 +132,7 @@ void Linker::layout() {
 
   // Emit the pre-assigned function names in sorted order
   for (const auto& P : functionNames) {
-    out.wasm.table.names.push_back(P.second);
+    getTableSegment().data.push_back(P.second);
   }
 
   for (auto& relocation : out.relocations) {
@@ -206,9 +206,11 @@ void Linker::layout() {
   }
 
   // ensure an explicit function type for indirect call targets
-  for (auto& name : out.wasm.table.names) {
-    auto* func = out.wasm.getFunction(name);
-    func->type = ensureFunctionType(getSig(func), &out.wasm)->name;
+  for (auto& segment : out.wasm.table.segments) {
+    for (auto& name : segment.data) {
+      auto* func = out.wasm.getFunction(name);
+      func->type = ensureFunctionType(getSig(func), &out.wasm)->name;
+    }
   }
 
   // Export malloc and free whenever availble. JavsScript version of malloc has
@@ -225,7 +227,9 @@ void Linker::layout() {
   }
 
   // finalize function table
-  out.wasm.table.initial = out.wasm.table.max = out.wasm.table.names.size();
+  if (out.wasm.table.segments.size() > 0) {
+    out.wasm.table.initial = out.wasm.table.max = getTableSegment().data.size();
+  }
 }
 
 bool Linker::linkObject(S2WasmBuilder& builder) {
@@ -385,10 +389,19 @@ void Linker::emscriptenGlue(std::ostream& o) {
   o << " }\n";
 }
 
+Table::Segment& Linker::getTableSegment() {
+  if (out.wasm.table.segments.size() == 0) {
+    out.wasm.table.segments.emplace_back(out.wasm.allocator.alloc<Const>()->set(Literal(uint32_t(0))));
+  } else {
+    assert(out.wasm.table.segments.size() == 1);
+  }
+  return out.wasm.table.segments[0];
+}
+
 Index Linker::getFunctionIndex(Name name) {
   if (!functionIndexes.count(name)) {
-    functionIndexes[name] = out.wasm.table.names.size();
-    out.wasm.table.names.push_back(name);
+    functionIndexes[name] = getTableSegment().data.size();
+    getTableSegment().data.push_back(name);
     if (debug) {
       std::cerr << "function index: " << name << ": "
                 << functionIndexes[name] << '\n';
@@ -406,7 +419,6 @@ bool hasI64ResultOrParam(FunctionType* ft) {
 }
 
 void Linker::makeDummyFunction() {
-  assert(out.wasm.table.names.empty());
   bool create = false;
   // Check if there are address-taken functions
   for (auto& relocation : out.relocations) {
@@ -424,9 +436,10 @@ void Linker::makeDummyFunction() {
 }
 
 void Linker::makeDynCallThunks() {
+  if (out.wasm.table.segments.size() == 0) return;
   std::unordered_set<std::string> sigs;
   wasm::Builder wasmBuilder(out.wasm);
-  for (const auto& indirectFunc : out.wasm.table.names) {
+  for (const auto& indirectFunc : getTableSegment().data) {
     // Skip generating thunks for the dummy function
     if (indirectFunc == dummyFunction) continue;
     std::string sig(getSig(out.wasm.getFunction(indirectFunc)));
diff --git a/src/wasm-linker.h b/src/wasm-linker.h
index bf76a8f94..a6f5d319a 100644
--- a/src/wasm-linker.h
+++ b/src/wasm-linker.h
@@ -282,6 +282,10 @@ class Linker {
 
   void ensureImport(Name target, std::string signature);
 
+  // Makes sure the table has a single segment, with offset 0,
+  // to which we can add content.
+  Table::Segment& getTableSegment();
+
   // Retrieves (and assigns) an entry index in the indirect function table for
   // a given function.
   Index getFunctionIndex(Name name);
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 7f652da5a..b70f70b4e 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -1430,7 +1430,7 @@ private:
       if (s[1]->str() == ANYFUNC) {
         // (table type (elem ..))
         parseElem(*s[2]);
-        wasm.table.initial = wasm.table.max = wasm.table.names.size();
+        wasm.table.initial = wasm.table.max = wasm.table.segments[0].data.size();
         return;
       }
       // first element isn't dollared, and isn't anyfunc. this could be old syntax for (table 0 1) which means function 0 and 1, or it could be (table initial max? type), look for type
@@ -1443,13 +1443,23 @@ private:
     }
     // old notation (table func1 func2 ..)
     parseElem(s);
-    wasm.table.initial = wasm.table.max = wasm.table.names.size();
+    wasm.table.initial = wasm.table.max = wasm.table.segments[0].data.size();
   }
 
   void parseElem(Element& s) {
-    for (Index i = 1; i < s.size(); i++) {
-      wasm.table.names.push_back(getFunctionName(*s[i]));
+    Index i = 1;
+    Expression* offset;
+    if (s[i]->isList()) {
+      // there is an init expression
+      offset = parseExpression(s[i++]);
+    } else {
+      offset = allocator.alloc<Const>()->set(Literal(int32_t(0)));
+    }
+    Table::Segment segment(offset);
+    for (; i < s.size(); i++) {
+      segment.data.push_back(getFunctionName(*s[i]));
     }
+    wasm.table.segments.push_back(segment);
   }
 
   void parseType(Element& s) {
diff --git a/src/wasm.h b/src/wasm.h
index 82933bda0..68558033d 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1434,8 +1434,19 @@ class Table {
 public:
   static const Index kMaxSize = Index(-1);
 
+  struct Segment {
+    Expression* offset;
+    std::vector<Name> data;
+    Segment() {}
+    Segment(Expression* offset) : offset(offset) {
+    }
+    Segment(Expression* offset, std::vector<Name>& init) : offset(offset) {
+      data.swap(init);
+    }
+  };
+
   Address initial, max;
-  std::vector<Name> names;
+  std::vector<Segment> segments;
 
   Table() : initial(0), max(kMaxSize) {}
 };
@@ -1445,6 +1456,7 @@ public:
   static const Address::address_t kPageSize = 64 * 1024;
   static const Address::address_t kMaxSize = ~Address::address_t(0) / kPageSize;
   static const Address::address_t kPageMask = ~(kPageSize - 1);
+
   struct Segment {
     Expression* offset;
     std::vector<char> data; // TODO: optimize