Minor wasm-split improvements (#3825)

- Support functions appearing more than once in the table. It turns out we
   were assuming and asserting that functions would appear at most once, but we
   weren't making use of that assumption in any way.

 - Make TableSlotManager::getSlot take a function Name rather than a RefFunc
   expression to avoid allocating and leaking unnecessary expressions.

 - Add and use a Builder interface for building TableElementSegments to make
   them more similar to other module-level items.

3 files changed, 49 insertions, 43 deletions

diff --git a/src/ir/module-splitting.cpp b/src/ir/module-splitting.cpp
index 7f12269f8..5f9ba0aad 100644
--- a/src/ir/module-splitting.cpp
+++ b/src/ir/module-splitting.cpp
@@ -54,8 +54,8 @@
 // ref.func instructions, they will have to be modified to use a similar layer
 // of indirection.
 //
-// The code as currently written makes a few assumptions about the module that
-// is being split:
+// The code as currently written makes a couple assumptions about the module
+// that is being split:
 //
 //   1. It assumes that mutable-globals is allowed. This could be worked around
 //      by introducing wrapper functions for globals and rewriting secondary
@@ -66,11 +66,6 @@
 //      is exactly one segment that may have a non-constant offset. It also
 //      assumes that all segments are active segments (although Binaryen does
 //      not yet support passive table segments anyway).
-//
-//   3. It assumes that each function appears in the table at most once. This
-//      isn't necessarily true in general or even for LLVM output after function
-//      deduplication. Relaxing this assumption would just require slightly more
-//      complex code, so it is a good candidate for a follow up PR.
 
 #include "ir/module-splitting.h"
 #include "ir/element-utils.h"
@@ -103,11 +98,6 @@ template<class F> void forEachElement(Module& module, F f) {
   });
 }
 
-static RefFunc* makeRefFunc(Module& wasm, Function* func) {
-  // FIXME: make the type NonNullable when we support it!
-  return Builder(wasm).makeRefFunc(func->name, func->sig);
-}
-
 struct TableSlotManager {
   struct Slot {
     Name tableName;
@@ -129,9 +119,10 @@ struct TableSlotManager {
   TableSlotManager(Module& module);
 
   Table* makeTable();
+  ElementSegment* makeElementSegment();
 
   // Returns the table index for `func`, allocating a new index if necessary.
-  Slot getSlot(RefFunc* entry);
+  Slot getSlot(Name func, Signature sig);
   void addSlot(Name func, Slot slot);
 };
 
@@ -148,14 +139,12 @@ Expression* TableSlotManager::Slot::makeExpr(Module& module) {
 }
 
 void TableSlotManager::addSlot(Name func, Slot slot) {
-  auto it = funcIndices.insert(std::make_pair(func, slot));
-  WASM_UNUSED(it);
-  assert(it.second && "Function already has multiple table slots");
+  // Ignore functions that already have slots.
+  funcIndices.insert({func, slot});
 }
 
 TableSlotManager::TableSlotManager(Module& module) : module(module) {
   // TODO: Reject or handle passive element segments
-
   auto it = std::find_if(module.tables.begin(),
                          module.tables.end(),
                          [&](std::unique_ptr<Table>& table) {
@@ -213,8 +202,15 @@ Table* TableSlotManager::makeTable() {
     Builder::makeTable(Names::getValidTableName(module, Name::fromInt(0))));
 }
 
-TableSlotManager::Slot TableSlotManager::getSlot(RefFunc* entry) {
-  auto slotIt = funcIndices.find(entry->func);
+ElementSegment* TableSlotManager::makeElementSegment() {
+  return module.addElementSegment(Builder::makeElementSegment(
+    Names::getValidElementSegmentName(module, Name::fromInt(0)),
+    activeTable->name,
+    Builder(module).makeConst(int32_t(0))));
+}
+
+TableSlotManager::Slot TableSlotManager::getSlot(Name func, Signature sig) {
+  auto slotIt = funcIndices.find(func);
   if (slotIt != funcIndices.end()) {
     return slotIt->second;
   }
@@ -226,27 +222,27 @@ TableSlotManager::Slot TableSlotManager::getSlot(RefFunc* entry) {
       activeBase = {activeTable->name, "", 0};
     }
 
+    // None of the existing segments should refer to the active table
     assert(std::all_of(module.elementSegments.begin(),
                        module.elementSegments.end(),
                        [&](std::unique_ptr<ElementSegment>& segment) {
                          return segment->table != activeTable->name;
                        }));
-    auto segment = std::make_unique<ElementSegment>(
-      activeTable->name, Builder(module).makeConst(int32_t(0)));
-    segment->setName(Name::fromInt(0), false);
-    activeSegment = segment.get();
-    module.addElementSegment(std::move(segment));
+
+    activeSegment = makeElementSegment();
   }
 
   Slot newSlot = {activeBase.tableName,
                   activeBase.global,
                   activeBase.index + Index(activeSegment->data.size())};
 
-  activeSegment->data.push_back(entry);
+  Builder builder(module);
+  activeSegment->data.push_back(builder.makeRefFunc(func, sig));
 
-  addSlot(entry->func, newSlot);
+  addSlot(func, newSlot);
   if (activeTable->initial <= newSlot.index) {
     activeTable->initial = newSlot.index + 1;
+    // TODO: handle the active table not being the dylink table (#3823)
     if (module.dylinkSection) {
       module.dylinkSection->tableSize = activeTable->initial;
     }
@@ -391,18 +387,15 @@ void ModuleSplitter::thunkExportedSecondaryFunctions() {
       // We've already created a thunk for this function
       continue;
     }
-    auto func = std::make_unique<Function>();
-    func->name = secondaryFunc;
-    func->sig = secondary.getFunction(secondaryFunc)->sig;
+    auto* func = primary.addFunction(Builder::makeFunction(
+      secondaryFunc, secondary.getFunction(secondaryFunc)->sig, {}));
     std::vector<Expression*> args;
     for (size_t i = 0, size = func->sig.params.size(); i < size; ++i) {
       args.push_back(builder.makeLocalGet(i, func->sig.params[i]));
     }
-
-    auto tableSlot = tableManager.getSlot(makeRefFunc(primary, func.get()));
+    auto tableSlot = tableManager.getSlot(secondaryFunc, func->sig);
     func->body = builder.makeCallIndirect(
       tableSlot.tableName, tableSlot.makeExpr(primary), args, func->sig);
-    primary.addFunction(std::move(func));
   }
 }
 
@@ -420,9 +413,8 @@ void ModuleSplitter::indirectCallsToSecondaryFunctions() {
       if (!parent.secondaryFuncs.count(curr->target)) {
         return;
       }
-      auto func = parent.secondary.getFunction(curr->target);
-      auto tableSlot =
-        parent.tableManager.getSlot(makeRefFunc(parent.primary, func));
+      auto* func = parent.secondary.getFunction(curr->target);
+      auto tableSlot = parent.tableManager.getSlot(curr->target, func->sig);
       replaceCurrent(
         builder.makeCallIndirect(tableSlot.tableName,
                                  tableSlot.makeExpr(parent.primary),
@@ -530,7 +522,9 @@ void ModuleSplitter::setupTablePatching() {
          ++i) {
       if (replacement->first == i) {
         // primarySeg->data[i] is a placeholder, so use the secondary function.
-        secondaryElems.push_back(makeRefFunc(secondary, replacement->second));
+        auto* func = replacement->second;
+        auto* ref = Builder(secondary).makeRefFunc(func->name, func->sig);
+        secondaryElems.push_back(ref);
         ++replacement;
       } else if (auto* get = primarySeg->data[i]->dynCast<RefFunc>()) {
         exportImportFunction(get->func);
@@ -556,8 +550,9 @@ void ModuleSplitter::setupTablePatching() {
     auto* offset = Builder(secondary).makeConst(int32_t(currBase));
     auto secondarySeg = std::make_unique<ElementSegment>(
       secondaryTable->name, offset, secondaryTable->type, currData);
-    secondarySeg->setName(Name::fromInt(secondary.elementSegments.size()),
-                          false);
+    Name name = Names::getValidElementSegmentName(
+      secondary, Name::fromInt(secondary.elementSegments.size()));
+    secondarySeg->setName(name, false);
     secondary.addElementSegment(std::move(secondarySeg));
   };
   for (auto curr = replacedElems.begin(); curr != replacedElems.end(); ++curr) {
@@ -566,7 +561,8 @@ void ModuleSplitter::setupTablePatching() {
       currBase = curr->first;
       currData.clear();
     }
-    currData.push_back(makeRefFunc(secondary, curr->second));
+    auto* func = curr->second;
+    currData.push_back(Builder(secondary).makeRefFunc(func->name, func->sig));
   }
   if (currData.size()) {
     finishSegment();
diff --git a/src/ir/module-splitting.h b/src/ir/module-splitting.h
index 09b4c0793..9a2597067 100644
--- a/src/ir/module-splitting.h
+++ b/src/ir/module-splitting.h
@@ -26,7 +26,7 @@
 // functions. The secondary module imports all of its dependencies from the
 // primary module.
 //
-// This code currently makes a few assumptions about the modules that will be
+// This code currently makes a couple assumptions about the modules that will be
 // split and will fail assertions if those assumptions are not true.
 //
 //   1) It assumes that mutable-globals are allowed.
@@ -34,8 +34,6 @@
 //   2) It assumes that either all segment offsets are constants or there is
 //      exactly one segment that may have a non-constant offset.
 //
-//   3) It assumes that each function appears in the table at most once.
-//
 // These requirements will be relaxed as necessary in the future, but for now
 // this code should be considered experimental and used with care.
 
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index 46ec734cc..bb739af89 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -88,10 +88,22 @@ public:
     table->type = type;
     table->initial = initial;
     table->max = max;
-
     return table;
   }
 
+  static std::unique_ptr<ElementSegment>
+  makeElementSegment(Name name,
+                     Name table,
+                     Expression* offset = nullptr,
+                     Type type = Type::funcref) {
+    auto seg = std::make_unique<ElementSegment>();
+    seg->name = name;
+    seg->table = table;
+    seg->offset = offset;
+    seg->type = type;
+    return seg;
+  }
+
   static std::unique_ptr<Export>
   makeExport(Name name, Name value, ExternalKind kind) {
     auto export_ = std::make_unique<Export>();