/*
 * Copyright 2020 WebAssembly Community Group participants
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// The process of module splitting involves these steps:
//
//   1. Create the new secondary modules.
//
//   2. Move the deferred functions from the primary to each of the secondary
//      modules.
//
//   3. For any secondary function exported from the primary module, export in
//      its place a trampoline function that makes an indirect call to its
//      placeholder function (and eventually to the original secondary
//      function), allocating a new table slot for the placeholder if necessary.
//
//   4. Export globals, tags, tables, and memories from the primary module and
//      import them in the secondary modules. If possible, move those module
//      items instead to the secondary modules.
//
//   5. Replace all references to each secondary module's functions in the
//      primary module's and each other secondary module's table segments with
//      references to imported placeholder functions.
//
//   6. Rewrite direct calls from primary functions to secondary functions to be
//      indirect calls to their placeholder functions (and eventually to their
//      original secondary functions), allocating new table slots for the
//      placeholders if necessary.
//
//   7. For each primary function directly called from a secondary function,
//      export the primary function if it is not already exported and import it
//      into each secondary module using it.
//
//   8. For each secondary module, create new dispatch table segments in the
//      module that will replace all the placeholder function references in the
//      table with references to their corresponding secondary functions upon
//      instantiation.
//
// Functions can be used or referenced three ways in a WebAssembly module: they
// can be exported, called, or referenced with ref.func. The above procedure
// introduces a layer of indirection to each of those mechanisms that removes
// all references to secondary functions from the primary module but restores
// the original program's semantics once the secondary modules are instantiated.
//
// 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
//      code that accesses them, but now that mutable-globals is shipped on all
//      browsers, hopefully that extra complexity won't be necessary.
//
//   2. It assumes that either all table segment offsets are constants or there
//      is exactly one segment that may have a non-constant offset. It also
//      assumes that all segments are dispatch segments.
//
//   3. It assumes that if exact function references are required for validity
//      (because they are stored in a local with an exact function type, for
//      example), then custom descriptors are allowed so primary functions can
//      be imported exactly. This could be worked around by removing exactness
//      from the IR before splitting.
//
#include "ir/module-splitting.h"
#include "ir/effects.h"
#include "ir/find_all.h"
#include "ir/module-utils.h"
#include "ir/names.h"
#include "support/small_vector.h"
#include "support/stdckdint.h"
#include "wasm-builder.h"
#include "wasm.h"

namespace wasm::ModuleSplitting {

namespace {

template<class F> void forEachElement(Module& module, F f) {
  ModuleUtils::iterActiveElementSegments(module, [&](ElementSegment* segment) {
    Name base;
    Index offset = 0;
    if (auto* c = segment->offset->dynCast<Const>()) {
      offset = c->value.getInteger();
    } else if (auto* g = segment->offset->dynCast<GlobalGet>()) {
      base = g->name;
    }
    for (Index i = 0; i < segment->data.size(); ++i) {
      f(segment->table, base, offset + i, segment->data[i]);
    }
  });
}

struct TableSlotManager {
  struct Slot {
    Name tableName;

    // If `global` is empty, then this slot is at a statically known index.
    Name global;
    Index index = 0;

    // Generate code to compute the index of this table slot
    Expression* makeExpr(Module& module);
  };
  Module& module;
  const std::vector<std::unique_ptr<Module>>& secondaries;
  Table* dispatchTable = nullptr;
  ElementSegment* dispatchSegment = nullptr;
  Slot dispatchBase;
  std::map<Name, Slot> funcIndices;
  std::vector<ElementSegment*> dispatchTableSegments;

  TableSlotManager(Module& module,
                   const std::vector<std::unique_ptr<Module>>& secondaries);

  Table* makeTable();
  ElementSegment* makeElementSegment();

  // Returns the table index for `func`, allocating a new index if necessary.
  Slot getSlot(Name func, HeapType type);
  void addSlot(Name func, Slot slot);
};

Expression* TableSlotManager::Slot::makeExpr(Module& module) {
  Builder builder(module);
  auto* table = module.getTable(tableName);
  auto makeIndex = [&]() {
    return builder.makeConst(Literal::makeFromInt32(index, table->addressType));
  };
  if (global) {
    Expression* getBase = builder.makeGlobalGet(global, table->addressType);
    auto addOp = table->is64() ? AddInt64 : AddInt32;
    return index == 0 ? getBase
                      : builder.makeBinary(addOp, getBase, makeIndex());
  } else {
    return makeIndex();
  }
}

void TableSlotManager::addSlot(Name func, Slot slot) {
  // Ignore functions that already have slots.
  funcIndices.insert({func, slot});
}

TableSlotManager::TableSlotManager(
  Module& module, const std::vector<std::unique_ptr<Module>>& secondaries)
  : module(module), secondaries(secondaries) {
  // If possible, just create a new table to manage all primary-to-secondary
  // calls lazily. Do not re-use slots for functions that will already be in
  // existing tables, since that is not correct in the face of table mutations.
  // However, do not do this for emscripten; its loader code (and dynamic
  // loading in particular) do not support this yet.
  // TODO: Reduce overhead by creating a separate table for each function type
  // if WasmGC is enabled.
  Export* emscriptenTableExport =
    module.getExportOrNull("__indirect_function_table");
  Table* singletonTable =
    module.tables.size() == 1 ? module.tables[0].get() : nullptr;
  bool emscriptenTableImport =
    singletonTable && singletonTable->imported() &&
    singletonTable->module == "env" &&
    singletonTable->base == "__indirect_function_table";

  if (module.features.hasReferenceTypes() && !emscriptenTableExport &&
      !emscriptenTableImport) {
    return;
  }

  // TODO: Reject or handle passive element segments
  auto funcref = Type(HeapType::func, Nullable);
  auto it = std::find_if(
    module.tables.begin(),
    module.tables.end(),
    [&](std::unique_ptr<Table>& table) { return table->type == funcref; });
  if (it == module.tables.end()) {
    // There is no indirect function table, so we will create one lazily.
    return;
  }

  dispatchTable = it->get();
  ModuleUtils::iterTableSegments(
    module, dispatchTable->name, [&](ElementSegment* segment) {
      dispatchTableSegments.push_back(segment);
    });

  if (dispatchTableSegments.empty()) {
    // There are no dispatch segments, so we will lazily create one and start
    // filling it at index 0.
    dispatchBase = {dispatchTable->name, Name(), 0};
  } else if (dispatchTableSegments.size() == 1 &&
             dispatchTableSegments[0]->type == funcref &&
             !dispatchTableSegments[0]->offset->is<Const>()) {
    // If there is exactly one table segment and that segment has a non-constant
    // offset, append new items to the end of that segment. In all other cases,
    // append new items at constant offsets after all existing items at constant
    // offsets.
    assert(dispatchTableSegments[0]->offset->is<GlobalGet>() &&
           "Unexpected initializer instruction");
    dispatchSegment = dispatchTableSegments[0];
    dispatchBase = {dispatchTable->name,
                    dispatchTableSegments[0]->offset->cast<GlobalGet>()->name,
                    0};
  } else {
    // Finds the segment with the highest occupied table slot so that new items
    // can be inserted contiguously at the end of it without accidentally
    // overwriting any other items. TODO: be more clever about filling gaps in
    // the table, if that is ever useful.
    Index maxIndex = 0;
    for (auto& segment : dispatchTableSegments) {
      assert(segment->offset->is<Const>() &&
             "Unexpected non-const segment offset with multiple segments");
      Index segmentBase = segment->offset->cast<Const>()->value.getInteger();
      if (segmentBase + segment->data.size() >= maxIndex) {
        maxIndex = segmentBase + segment->data.size();
        dispatchSegment = segment;
        dispatchBase = {dispatchTable->name, Name(), segmentBase};
      }
    }
  }

  // Initialize funcIndices with the functions already in the table.
  forEachElement(module,
                 [&](Name table, Name base, Index offset, Expression* elem) {
                   if (auto* func = elem->dynCast<RefFunc>()) {
                     addSlot(func->func, {table, base, offset});
                   }
                 });
}

Table* TableSlotManager::makeTable() {
  // Because the dispatch table will be imported in secondary modules, its name
  // should not collide with any existing tables in primary and secondary
  // modules.
  std::unordered_set<Name> secondaryTableNames;
  for (auto& secondary : secondaries) {
    for (auto& table : secondary->tables) {
      secondaryTableNames.insert(table->name);
    }
  }
  Name name = Names::getValidName("0", [&](Name test) {
    if (module.getTableOrNull(test)) {
      return false;
    }
    if (secondaryTableNames.contains(test)) {
      return false;
    }
    return true;
  });
  return module.addTable(Builder::makeTable(name));
}

ElementSegment* TableSlotManager::makeElementSegment() {
  Builder builder(module);
  Expression* offset =
    builder.makeConst(Literal::makeFromInt32(0, dispatchTable->addressType));
  return module.addElementSegment(Builder::makeElementSegment(
    Names::getValidElementSegmentName(module, Name::fromInt(0)),
    dispatchTable->name,
    offset));
}

TableSlotManager::Slot TableSlotManager::getSlot(Name func, HeapType type) {
  auto slotIt = funcIndices.find(func);
  if (slotIt != funcIndices.end()) {
    return slotIt->second;
  }

  // If there are no segments yet, allocate one.
  if (dispatchSegment == nullptr) {
    if (dispatchTable == nullptr) {
      dispatchTable = makeTable();
      dispatchBase = {dispatchTable->name, Name(), 0};
    }

    // None of the existing segments should refer to the dispatch table
    assert(std::all_of(module.elementSegments.begin(),
                       module.elementSegments.end(),
                       [&](std::unique_ptr<ElementSegment>& segment) {
                         return segment->table != dispatchTable->name;
                       }));

    dispatchSegment = makeElementSegment();
  }

  Slot newSlot = {dispatchBase.tableName,
                  dispatchBase.global,
                  dispatchBase.index + Index(dispatchSegment->data.size())};

  Builder builder(module);
  auto funcType = Type(type, NonNullable, Inexact);
  dispatchSegment->data.push_back(builder.makeRefFunc(func, funcType));

  addSlot(func, newSlot);
  if (dispatchTable->initial <= newSlot.index) {
    dispatchTable->initial = newSlot.index + 1;
    // TODO: handle the dispatch table not being the dylink table (#3823)
    if (module.dylinkSection) {
      module.dylinkSection->tableSize = dispatchTable->initial;
    }
  }
  if (dispatchTable->max <= newSlot.index) {
    dispatchTable->max = newSlot.index + 1;
  }
  return newSlot;
}

// Module items ownership tracking

// Struct containing sets of used module elements of a single module
struct UsedNames {
  std::unordered_set<Name> globals;
  std::unordered_set<Name> memories;
  std::unordered_set<Name> tables;
  std::unordered_set<Name> tags;
  std::unordered_set<Name> dataSegments;
  std::unordered_set<Name> elementSegments;
};

// A tracker that, given a module element, tracks which module is its owner,
// i.e., where the element should be placed, and the list of secondary modules
// using this element.
struct OwnershipTracker {
  UsedNames primaryUsed;
  std::vector<UsedNames> secondaryUsed;

  struct ItemInfo {
    UsedNames* owner = nullptr;
    SmallVector<Module*, 2> usingSecondaries;
  };

  std::unordered_map<Name, ItemInfo> tables;
  std::unordered_map<Name, ItemInfo> memories;
  std::unordered_map<Name, ItemInfo> globals;
  std::unordered_map<Name, ItemInfo> tags;
  std::unordered_map<Name, ItemInfo> dataSegments;
  std::unordered_map<Name, ItemInfo> elementSegments;

  const std::vector<std::unique_ptr<Module>>* secondaries = nullptr;

  using FieldType = std::unordered_set<Name> UsedNames::*;
  using MapType = std::unordered_map<Name, ItemInfo> OwnershipTracker::*;

  template<typename T> void insert(Name name, UsedNames* owner) {
#define INSERT_ITEM(ItemType, field)                                           \
  if constexpr (std::is_same_v<T, ItemType>) {                                 \
    insertImpl(name, owner, &OwnershipTracker::field, &UsedNames::field);      \
  }
    INSERT_ITEM(Table, tables)
    INSERT_ITEM(Memory, memories)
    INSERT_ITEM(Global, globals)
    INSERT_ITEM(Tag, tags)
    INSERT_ITEM(DataSegment, dataSegments)
    INSERT_ITEM(ElementSegment, elementSegments)
#undef INSERT_ITEM
  }

  // 'mapField' points to one of OwnershipTracker's maps, such as
  //   std::unordered_map<Name, ItemInfo> globals;
  // 'field' points to one of UsedName's sets, such as
  //   std::unordered_set<Name> globals;
  void
  insertImpl(Name name, UsedNames* owner, MapType mapField, FieldType field) {
    (owner->*field).insert(name);
    // Figure out which module is the 'owner' of this item. If it is used by a
    // single secondary module, that secondary module is the owner. If it is
    // used by the primary module or multiple secondary modules, the primary
    // module is the owner.
    auto [it, inserted] = (this->*mapField).insert({name, ItemInfo{owner, {}}});
    Module* secondary = nullptr;
    if (owner != &primaryUsed) {
      size_t index = owner - secondaryUsed.data();
      secondary = (*secondaries)[index].get();
    }
    if (inserted) {
      if (secondary) {
        it->second.usingSecondaries.push_back(secondary);
      }
    } else {
      if (it->second.owner != owner) {
        it->second.owner = &primaryUsed;
        (primaryUsed.*field).insert(name);
      }
      if (secondary) {
        auto& vec = it->second.usingSecondaries;
        if (std::find(vec.begin(), vec.end(), secondary) == vec.end()) {
          vec.push_back(secondary);
        }
      }
    }
  }

  UsedNames* getOwner(Name name,
                      const std::unordered_map<Name, ItemInfo>& map) {
    auto it = map.find(name);
    if (it != map.end()) {
      return it->second.owner;
    }
    return nullptr;
  }

  const SmallVector<Module*, 2>&
  getUsingSecondaries(Name name,
                      const std::unordered_map<Name, ItemInfo>& map) {
    auto it = map.find(name);
    if (it != map.end()) {
      return it->second.usingSecondaries;
    }
    static SmallVector<Module*, 2> empty;
    return empty;
  }

  bool isUnused(Name name, const std::unordered_map<Name, ItemInfo>& map) {
    return getOwner(name, map) == nullptr;
  }

  bool usedBySingleSecondary(Name name,
                             const std::unordered_map<Name, ItemInfo>& map) {
    auto* owner = getOwner(name, map);
    return owner != nullptr && owner != &primaryUsed;
  }
};

struct ModuleSplitter {
  const Config& config;
  std::vector<std::unique_ptr<Module>> secondaries;

  Module& primary;

  std::unordered_set<Name> primaryFuncs;
  std::unordered_set<Name> allSecondaryFuncs;
  std::unordered_map<Name, Index> funcToSecondaryIndex;

  TableSlotManager tableManager;

  Names::MinifiedNameGenerator minified;

  // Map from internal function names to (one of) their corresponding export
  // names.
  std::unordered_map<Name, Name> exportedPrimaryFuncs;

  // Map from <Kind, InternalName> to their corresponding export names for
  // non-function items.
  std::unordered_map<std::pair<ExternalKind, Name>, Name> exportedPrimaryItems;

  // For each table, map placeholder indices to the names of the functions they
  // replace.
  std::unordered_map<Name, std::map<size_t, Name>> placeholderMap;

  // Map from original secondary function name to its trampoline
  std::unordered_map<Name, Name> trampolineMap;

  void shareDispatchTable(Module* secondary);

  // Initialization helpers
  static std::unique_ptr<Module> initSecondary(const Module& primary);
  static std::unordered_map<Name, Name>
  initExportedPrimaryFuncs(const Module& primary);
  static std::unordered_map<std::pair<ExternalKind, Name>, Name>
  initExportedPrimaryItems(const Module& primary);

  // Other helpers
  void exportImportFunction(Name func, const std::set<Module*>& modules);
  void makeImportExport(Importable& primaryItem,
                        Importable& secondaryItem,
                        const std::string& genericExportName,
                        ExternalKind kind);
  Name getTrampoline(Name funcName);

  OwnershipTracker tracker;

  void computeUsedNames();

  // Main splitting steps
  void classifyFunctions();
  void moveSecondaryFunctions();
  void thunkExportedSecondaryFunctions();
  void shareImportableItems();
  void indirectReferencesToSecondaryFunctions();
  void indirectCallsToSecondaryFunctions();
  void exportImportCalledPrimaryFunctions();
  void setupTablePatching();

  ModuleSplitter(Module& primary, const Config& config)
    : config(config), primary(primary), tableManager(primary, secondaries),
      exportedPrimaryFuncs(initExportedPrimaryFuncs(primary)),
      exportedPrimaryItems(initExportedPrimaryItems(primary)) {
    classifyFunctions();
    moveSecondaryFunctions();
    thunkExportedSecondaryFunctions();
    shareImportableItems();
    indirectReferencesToSecondaryFunctions();
    indirectCallsToSecondaryFunctions();
    exportImportCalledPrimaryFunctions();
    setupTablePatching();
  }
};

std::unique_ptr<Module> ModuleSplitter::initSecondary(const Module& primary) {
  // Create the secondary module and copy trivial properties.
  auto secondary = std::make_unique<Module>();
  secondary->features = primary.features;
  secondary->hasFeaturesSection = primary.hasFeaturesSection;
  return secondary;
}

void ModuleSplitter::classifyFunctions() {
  // Find functions that refer to data or element segments. These functions must
  // remain in the primary module because segments cannot be exported to be
  // accessed from the secondary module.
  //
  // TODO: Investigate other options, such as moving the segments to the
  // secondary module or replacing the segment-using instructions in the
  // secondary module with calls to imports.
  ModuleUtils::ParallelFunctionAnalysis<std::vector<Name>>
    segmentReferrerCollector(
      primary, [&](Function* func, std::vector<Name>& segmentReferrers) {
        if (func->imported()) {
          return;
        }

        struct SegmentReferrerCollector
          : PostWalker<SegmentReferrerCollector,
                       UnifiedExpressionVisitor<SegmentReferrerCollector>> {
          bool hasSegmentReference = false;

          void visitExpression(Expression* curr) {

#define DELEGATE_ID curr->_id

#define DELEGATE_START(id) [[maybe_unused]] auto* cast = curr->cast<id>();
#define DELEGATE_GET_FIELD(id, field) cast->field
#define DELEGATE_FIELD_TYPE(id, field)
#define DELEGATE_FIELD_HEAPTYPE(id, field)
#define DELEGATE_FIELD_CHILD(id, field)
#define DELEGATE_FIELD_OPTIONAL_CHILD(id, field)
#define DELEGATE_FIELD_INT(id, field)
#define DELEGATE_FIELD_LITERAL(id, field)
#define DELEGATE_FIELD_NAME(id, field)
#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)
#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)
#define DELEGATE_FIELD_ADDRESS(id, field)

#define DELEGATE_FIELD_NAME_KIND(id, field, kind)                              \
  if (kind == ModuleItemKind::DataSegment ||                                   \
      kind == ModuleItemKind::ElementSegment) {                                \
    hasSegmentReference = true;                                                \
  }

#include "wasm-delegations-fields.def"
          }
        };
        SegmentReferrerCollector collector;
        collector.walkFunction(func);
        if (collector.hasSegmentReference) {
          segmentReferrers.push_back(func->name);
        }
      });

  std::unordered_set<Name> segmentReferrers;
  for (auto& [_, referrers] : segmentReferrerCollector.map) {
    segmentReferrers.insert(referrers.begin(), referrers.end());
  }

  std::unordered_set<Name> configSecondaryFuncs;
  for (auto& funcs : config.secondaryFuncs) {
    configSecondaryFuncs.insert(funcs.begin(), funcs.end());
  }
  for (auto& func : primary.functions) {
    if (func->imported() || !configSecondaryFuncs.contains(func->name) ||
        segmentReferrers.contains(func->name)) {
      primaryFuncs.insert(func->name);
    } else {
      assert(func->name != primary.start && "The start function must be kept");
      allSecondaryFuncs.insert(func->name);
    }
  }
}

std::unordered_map<Name, Name>
ModuleSplitter::initExportedPrimaryFuncs(const Module& primary) {
  std::unordered_map<Name, Name> functionExportNames;
  for (auto& ex : primary.exports) {
    if (ex->kind == ExternalKind::Function) {
      functionExportNames[*ex->getInternalName()] = ex->name;
    }
  }
  return functionExportNames;
}

std::unordered_map<std::pair<ExternalKind, Name>, Name>
ModuleSplitter::initExportedPrimaryItems(const Module& primary) {
  std::unordered_map<std::pair<ExternalKind, Name>, Name> exports;
  for (auto& ex : primary.exports) {
    if (ex->kind != ExternalKind::Function) {
      if (auto* name = ex->getInternalName()) {
        exports[std::make_pair(ex->kind, *name)] = ex->name;
      }
    }
  }
  return exports;
}

void ModuleSplitter::makeImportExport(Importable& primaryItem,
                                      Importable& secondaryItem,
                                      const std::string& genericExportName,
                                      ExternalKind kind) {
  secondaryItem.name = primaryItem.name;
  secondaryItem.hasExplicitName = primaryItem.hasExplicitName;
  secondaryItem.module = config.importNamespace;
  auto exportIt = exportedPrimaryItems.find({kind, primaryItem.name});
  if (exportIt != exportedPrimaryItems.end()) {
    secondaryItem.base = exportIt->second;
  } else {
    std::string baseName =
      config.newExportPrefix +
      (config.minimizeNewExportNames ? minified.getName() : genericExportName);
    Name exportName = Names::getValidExportName(primary, baseName);
    primary.addExport(
      std::make_unique<Export>(exportName, kind, primaryItem.name));
    secondaryItem.base = exportName;
    exportedPrimaryItems[{kind, primaryItem.name}] = exportName;
  }
}

void ModuleSplitter::exportImportFunction(Name funcName,
                                          const std::set<Module*>& modules) {
  Name exportName;
  // If the function is already exported, use the existing export name.
  // Otherwise, create a new export for it.
  auto exportIt = exportedPrimaryFuncs.find(funcName);
  if (exportIt != exportedPrimaryFuncs.end()) {
    exportName = exportIt->second;
  } else {
    if (config.minimizeNewExportNames) {
      do {
        exportName = config.newExportPrefix + minified.getName();
      } while (primary.getExportOrNull(exportName) != nullptr);
    } else {
      exportName = Names::getValidExportName(
        primary, config.newExportPrefix + funcName.toString());
    }
    primary.addExport(
      Builder::makeExport(exportName, funcName, ExternalKind::Function));
    exportedPrimaryFuncs[funcName] = exportName;
  }
  // Import the function if it is not already imported into the secondary
  // module.
  for (auto* secondary : modules) {
    if (secondary->getFunctionOrNull(funcName) == nullptr) {
      auto primaryFunc = primary.getFunction(funcName);
      auto func = Builder::makeFunction(funcName, primaryFunc->type, {});
      func->hasExplicitName = primaryFunc->hasExplicitName;
      func->module = config.importNamespace;
      func->base = exportName;
      func->type = func->type.withInexactIfNoCustomDescs(secondary->features);
      secondary->addFunction(std::move(func));
    }
  }
}

void ModuleSplitter::moveSecondaryFunctions() {
  // Move the specified functions from the primary to the secondary modules.
  for (auto& funcNames : config.secondaryFuncs) {
    auto secondary = initSecondary(primary);
    for (auto funcName : funcNames) {
      if (allSecondaryFuncs.contains(funcName)) {
        auto* func = primary.getFunction(funcName);
        ModuleUtils::copyFunction(func, *secondary);
        primary.removeFunction(funcName);
        funcToSecondaryIndex[funcName] = secondaries.size();
      }
    }
    secondaries.push_back(std::move(secondary));
  }
}

Name ModuleSplitter::getTrampoline(Name funcName) {
  auto [it, inserted] = trampolineMap.insert({funcName, Name()});
  if (!inserted) {
    return it->second;
  }

  Builder builder(primary);
  Module& secondary = *secondaries.at(funcToSecondaryIndex.at(funcName));
  auto* oldFunc = secondary.getFunction(funcName);
  auto trampoline = Names::getValidFunctionName(
    primary, std::string("trampoline_") + funcName.toString());
  it->second = trampoline;

  // Generate the call and the function. We generate a direct call here, but
  // this will be converted to a call_indirect in
  // indirectCallsToSecondaryFunctions.
  std::vector<Expression*> args;
  for (Index i = 0; i < oldFunc->getNumParams(); i++) {
    args.push_back(builder.makeLocalGet(i, oldFunc->getLocalType(i)));
  }
  auto* call = builder.makeCall(funcName, args, oldFunc->getResults());

  auto func = builder.makeFunction(trampoline, oldFunc->type, {}, call);
  func->hasExplicitName = oldFunc->hasExplicitName;
  primary.addFunction(std::move(func));
  primaryFuncs.insert(trampoline);
  return trampoline;
}

void ModuleSplitter::shareDispatchTable(Module* secondary) {
  assert(tableManager.dispatchTable);
  auto secondaryTable =
    secondary->getTableOrNull(tableManager.dispatchTable->name);
  if (secondaryTable) {
    // In case it's already in the secondary module, sync the initial/max
    secondaryTable->initial = tableManager.dispatchTable->initial;
    secondaryTable->max = tableManager.dispatchTable->max;
  } else {
    secondaryTable =
      ModuleUtils::copyTable(tableManager.dispatchTable, *secondary);
    makeImportExport(*tableManager.dispatchTable,
                     *secondaryTable,
                     "table",
                     ExternalKind::Table);
  }
  if (tableManager.dispatchBase.global) {
    auto* primaryGlobal = primary.getGlobal(tableManager.dispatchBase.global);
    auto* secondaryGlobal =
      secondary->getGlobalOrNull(tableManager.dispatchBase.global);
    if (!secondaryGlobal) {
      secondaryGlobal = ModuleUtils::copyGlobal(primaryGlobal, *secondary);
      makeImportExport(
        *primaryGlobal, *secondaryGlobal, "global", ExternalKind::Global);
    }
  }
}

void ModuleSplitter::thunkExportedSecondaryFunctions() {
  // Update exports of secondary functions in the primary module to export
  // wrapper functions that indirectly call the secondary functions. We are
  // adding secondary function names to the primary table here, but they will be
  // replaced with placeholder functions later along with any references to
  // secondary functions that were already in the table.
  Builder builder(primary);
  for (auto& ex : primary.exports) {
    if (ex->kind != ExternalKind::Function ||
        !allSecondaryFuncs.contains(*ex->getInternalName())) {
      continue;
    }
    Name trampoline = getTrampoline(*ex->getInternalName());
    ex->setInternalName(trampoline);
  }
}

void ModuleSplitter::computeUsedNames() {
  tracker.secondaries = &secondaries;
  tracker.secondaryUsed.resize(secondaries.size());
  UsedNames& primaryUsed = tracker.primaryUsed;
  std::vector<UsedNames>& secondaryUsed = tracker.secondaryUsed;

  struct NameCollector
    : public PostWalker<NameCollector,
                        UnifiedExpressionVisitor<NameCollector>> {
    UsedNames& used;
    OwnershipTracker& tracker;

    NameCollector(UsedNames& used, OwnershipTracker& tracker)
      : used(used), tracker(tracker) {}

    void visitExpression(Expression* curr) {
#define DELEGATE_ID curr->_id
#define DELEGATE_START(id) [[maybe_unused]] auto* cast = curr->cast<id>();
#define DELEGATE_GET_FIELD(id, field) cast->field
#define DELEGATE_FIELD_TYPE(id, field)
#define DELEGATE_FIELD_HEAPTYPE(id, field)
#define DELEGATE_FIELD_CHILD(id, field)
#define DELEGATE_FIELD_INT(id, field)
#define DELEGATE_FIELD_LITERAL(id, field)
#define DELEGATE_FIELD_NAME(id, field)
#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)
#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)
#define DELEGATE_FIELD_ADDRESS(id, field)

#define DELEGATE_FIELD_NAME_KIND(id, field, kind)                              \
  if (cast->field.is()) {                                                      \
    switch (kind) {                                                            \
      case ModuleItemKind::Table:                                              \
        tracker.insert<Table>(cast->field, &used);                             \
        break;                                                                 \
      case ModuleItemKind::Memory:                                             \
        tracker.insert<Memory>(cast->field, &used);                            \
        break;                                                                 \
      case ModuleItemKind::Global:                                             \
        tracker.insert<Global>(cast->field, &used);                            \
        break;                                                                 \
      case ModuleItemKind::Tag:                                                \
        tracker.insert<Tag>(cast->field, &used);                               \
        break;                                                                 \
      case ModuleItemKind::DataSegment:                                        \
        tracker.insert<DataSegment>(cast->field, &used);                       \
        break;                                                                 \
      case ModuleItemKind::ElementSegment:                                     \
        tracker.insert<ElementSegment>(cast->field, &used);                    \
        break;                                                                 \
      case ModuleItemKind::Function:                                           \
      case ModuleItemKind::Invalid:                                            \
        break;                                                                 \
    }                                                                          \
  }

#include "wasm-delegations-fields.def"
    }
  };

  // Given a module, collect names used in the module
  auto scanModule = [&](Module& module, UsedNames& used) {
    NameCollector collector(used, tracker);
    for (auto& func : module.functions) {
      if (!func->imported()) {
        collector.walk(func->body);
      }
    }
  };

  scanModule(primary, primaryUsed);
  for (size_t i = 0; i < secondaries.size(); ++i) {
    scanModule(*secondaries[i], secondaryUsed[i]);
  }

  // If primary module has exports, they are "used" in it. Secondary modules
  // don't have exports, so this only applies to the primary module.
  for (auto& ex : primary.exports) {
    switch (ex->kind) {
      case ExternalKind::Global:
        tracker.insert<Global>(*ex->getInternalName(), &primaryUsed);
        break;
      case ExternalKind::Memory:
        tracker.insert<Memory>(*ex->getInternalName(), &primaryUsed);
        break;
      case ExternalKind::Table:
        tracker.insert<Table>(*ex->getInternalName(), &primaryUsed);
        break;
      case ExternalKind::Tag:
        tracker.insert<Tag>(*ex->getInternalName(), &primaryUsed);
        break;
      default:
        break;
    }
  }

  // We need to assume the dispatch table and its base global are used in the
  // primary module, because we will create segments there later.
  if (tableManager.dispatchTable) {
    tracker.insert<Table>(tableManager.dispatchTable->name, &primaryUsed);
  }
  if (tableManager.dispatchBase.global) {
    tracker.insert<Global>(tableManager.dispatchBase.global, &primaryUsed);
  }

  // If custom-descirptors is enabled, global and table initializers can trap.
  // Trapping globals should stay in the primary module to preserve the trapping
  // behavior upon instantiation.
  if (primary.features.hasCustomDescriptors()) {
    for (auto& global : primary.globals) {
      if (global->init &&
          EffectAnalyzer(config.passOptions, primary, global->init)
            .hasUnremovableSideEffects()) {
        tracker.insert<Global>(global->name, &primaryUsed);
      }
    }
    for (auto& table : primary.tables) {
      if (table->init &&
          EffectAnalyzer(config.passOptions, primary, table->init)
            .hasUnremovableSideEffects()) {
        tracker.insert<Table>(table->name, &primaryUsed);
      }
    }
  }

  // Scan table initializers into their owning modules. If a table is used by a
  // single secondary module, its initializer dependencies are marked as "used"
  // in that secondary module. Otherwise, they are marked as used in the primary
  // module.
  if (primary.features.hasGC()) {
    for (auto& table : primary.tables) {
      if (!table->init) {
        continue;
      }
      if (UsedNames* owner = tracker.getOwner(table->name, tracker.tables)) {
        NameCollector(*owner, tracker).walk(table->init);
      }
    }
  }

  auto mayTrap = [&](auto* segment) {
    if constexpr (std::is_same_v<decltype(segment), ElementSegment*>) {
      if (primary.features.hasCustomDescriptors()) {
        for (auto* item : segment->data) {
          if (EffectAnalyzer(config.passOptions, primary, item)
                .hasUnremovableSideEffects()) {
            return true;
          }
        }
      }
    }

    // Check for out-of-bounds offset. This is adapted from maybeRootSegment
    // function in RemoveUnusedModuleElements pass.
    if (!config.passOptions.trapsNeverHappen) {
      Index segmentSize;
      Index parentSize;
      if constexpr (std::is_same_v<decltype(segment), DataSegment*>) {
        segmentSize = segment->data.size();
        auto* memory = primary.getMemory(segment->memory);
        parentSize = memory->initial << memory->pageSizeLog2;
      } else {
        segmentSize = segment->data.size();
        auto* table = primary.getTable(segment->table);
        parentSize = table->initial;
      }

      // Check if this might trap. If it is obviously in bounds then it cannot.
      auto* c = segment->offset->template dynCast<Const>();
      // Check for overflow in the largest possible space of addresses.
      uint64_t maxWritten;
      // If there is no integer, or if there is and the addition overflows, or
      // if the addition leads to a too-large value, then we may trap.
      if (!c ||
          std::ckd_add(&maxWritten,
                       (uint64_t)segmentSize,
                       (uint64_t)c->value.getInteger()) ||
          maxWritten > parentSize) {
        return true;
      }
    }
    return false;
  };

  // Iterate on active data and element segments. If its table or memory is
  // used by a single secondary module, mark it "used" there. Only scan its
  // 'offset' or 'data'(in case of ElementSegment) and add it to that module's
  // used only when it is a sole secondary owner. If not assign it to the
  // primary module and scan it there.
  ModuleUtils::iterActiveDataSegments(primary, [&](DataSegment* segment) {
    UsedNames* owner = tracker.getOwner(segment->memory, tracker.memories);
    // Trapping segments should be kept in the primary module because they are
    // evaluated at the instantiation time.
    if (mayTrap(segment)) {
      owner = &primaryUsed;
    }
    if (!owner) {
      return;
    }
    tracker.insert<DataSegment>(segment->name, owner);
    tracker.insert<Memory>(segment->memory, owner);
    if (segment->offset) {
      NameCollector(*owner, tracker).walk(segment->offset);
    }
  });

  ModuleUtils::iterActiveElementSegments(primary, [&](ElementSegment* segment) {
    UsedNames* owner = tracker.getOwner(segment->table, tracker.tables);

    // If placeholders are NOT used, and if all functions in an element segment
    // belong to a single secondary module, we can move the segment to that
    // secondary module, because those functions aren't available until the
    // secondary module is loaded anyway.
    if (!config.usePlaceholders && segment->type.isFunction() &&
        owner == &primaryUsed) {
      bool foundSecondary = false;
      Index secondaryIndex = 0;
      bool keepInPrimary = false;
      for (auto* item : segment->data) {
        if (item->is<GlobalGet>()) {
          // If we encounter a global.get, we just keep the segment in the
          // primary module. TODO Consider moving this if this global is in a
          // secondary module.
          keepInPrimary = true;
          break;
        } else if (auto* ref = item->dynCast<RefFunc>()) {
          // If this ref.func is in the primary module, keep the segment there.
          auto it = funcToSecondaryIndex.find(ref->func);
          if (it == funcToSecondaryIndex.end()) {
            keepInPrimary = true;
            break;
          }
          // If the segment contains ref.funcs from more than one secondary
          // modules, keep it in the primary.
          if (foundSecondary && secondaryIndex != it->second) {
            keepInPrimary = true;
            break;
          }
          // This segment contains ref.funcs from a single secondary module.
          foundSecondary = true;
          secondaryIndex = it->second;
        }
      }
      if (!keepInPrimary && foundSecondary) {
        owner = &secondaryUsed[secondaryIndex];
      }
    }

    if (mayTrap(segment)) {
      owner = &primaryUsed;
    }
    if (!owner) {
      return;
    }
    tracker.insert<ElementSegment>(segment->name, owner);
    tracker.insert<Table>(segment->table, owner);
    if (segment->offset) {
      NameCollector(*owner, tracker).walk(segment->offset);
    }
    for (auto* item : segment->data) {
      NameCollector(*owner, tracker).walk(item);
    }
  });

  // Passive element segments contain expressions (e.g. global.get) in their
  // data arrays that must be scanned. Since currently functions referring to
  // segments are forced into the primary module, passive segments always belong
  // to the primary module.
  for (auto& segment : primary.elementSegments) {
    if (segment->isPassive() &&
        primaryUsed.elementSegments.contains(segment->name)) {
      for (auto* item : segment->data) {
        NameCollector(primaryUsed, tracker).walk(item);
      }
    }
  }

  // Compute the transitive closure of globals referenced in other globals'
  // initializers. WebAssembly validation requires that global initializers only
  // refer to previously defined globals. Therefore, `primary.globals` is
  // guaranteed to be topologically sorted with respect to its internal
  // dependencies. By iterating in reverse, we are guaranteed to process
  // dependent globals before the globals they depend on, allowing us to
  // propagate ownership in a single pass.
  for (auto it = primary.globals.rbegin(); it != primary.globals.rend(); ++it) {
    auto& global = *it;
    if (!global->init) {
      continue;
    }
    if (UsedNames* owner = tracker.getOwner(global->name, tracker.globals)) {
      for (auto* get : FindAll<GlobalGet>(global->init).list) {
        tracker.insert<Global>(get->name, owner);
      }
    }
  }
}

void ModuleSplitter::shareImportableItems() {
  computeUsedNames();

  // Share module items with secondary modules.
  // 1. Only share an item with the modules that use it
  // 2. If an item is used by only a single secondary module, move the item to
  //    that secondary module. If an item is used by multiple modules (including
  //    the primary and secondary modules), export the item from the primary and
  //    import it from the using secondary modules.

  auto shareElements = [&](auto& elements,
                           auto& trackerElements,
                           auto copyElement,
                           auto removeElements,
                           const char* exportName = nullptr,
                           ExternalKind kind = ExternalKind::Invalid) {
    std::unordered_set<Name> elementsToRemove;
    for (auto& element : elements) {
      if (tracker.isUnused(element->name, trackerElements)) {
        elementsToRemove.insert(element->name);
      } else if (tracker.usedBySingleSecondary(element->name,
                                               trackerElements)) {
        auto* secondary =
          tracker.getUsingSecondaries(element->name, trackerElements)[0];
        copyElement(element.get(), *secondary);
        elementsToRemove.insert(element->name);
      } else {
        // We only import and export Importables, i.e., we don't do this for
        // segments.
        using T = std::remove_pointer_t<decltype(element.get())>;
        if constexpr (std::is_base_of_v<Importable, T>) {
          for (auto* secondary :
               tracker.getUsingSecondaries(element->name, trackerElements)) {
            auto* secondaryElement = copyElement(element.get(), *secondary);
            makeImportExport(*element, *secondaryElement, exportName, kind);
          }
        }
      }
    }
    (primary.*removeElements)(
      [&](auto* element) { return elementsToRemove.count(element->name); });
  };

  shareElements(primary.memories,
                tracker.memories,
                ModuleUtils::copyMemory,
                &Module::removeMemories,
                "memory",
                ExternalKind::Memory);

  shareElements(primary.tables,
                tracker.tables,
                ModuleUtils::copyTable,
                &Module::removeTables,
                "table",
                ExternalKind::Table);

  for (auto& global : primary.globals) {
    if (global->mutable_) {
      assert(primary.features.hasMutableGlobals() &&
             "TODO: add wrapper functions for disallowed mutable globals");
    }
  }
  shareElements(primary.globals,
                tracker.globals,
                ModuleUtils::copyGlobal,
                &Module::removeGlobals,
                "global",
                ExternalKind::Global);

  shareElements(primary.tags,
                tracker.tags,
                ModuleUtils::copyTag,
                &Module::removeTags,
                "tag",
                ExternalKind::Tag);

  // Move segments that are exclusively used in a secondary module. If not, do
  // nothing. (Segments cannot be imported / exported. They will be handled in
  // indirectReferencesToSecondaryFunctions.)

  shareElements(primary.dataSegments,
                tracker.dataSegments,
                ModuleUtils::copyDataSegment,
                &Module::removeDataSegments);

  shareElements(primary.elementSegments,
                tracker.elementSegments,
                ModuleUtils::copyElementSegment,
                &Module::removeElementSegments);
}

void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
  // Turn references to secondary functions into references to thunks that
  // perform a direct call to the original referent. The direct calls in the
  // thunks will be handled like all other cross-module calls later, in
  // |indirectCallsToSecondaryFunctions|.
  struct Gatherer : public PostWalker<Gatherer> {
    ModuleSplitter& parent;

    Gatherer(ModuleSplitter& parent) : parent(parent) {}

    // Collect RefFuncs in a map from the function name to all RefFuncs that
    // refer to it. We only collect this for secondary funcs.
    InsertOrderedMap<Name, std::vector<RefFunc*>> map;

    void visitRefFunc(RefFunc* curr) {
      Module* currModule = getModule();
      // Add ref.func to the map when
      // 1. ref.func's target func is in one of the secondary modules and
      // 2. the current module is a different module (either the primary module
      //    or a different secondary module)
      if (parent.allSecondaryFuncs.contains(curr->func) &&
          (currModule == &parent.primary ||
           parent.secondaries.at(parent.funcToSecondaryIndex.at(curr->func))
               .get() != currModule)) {
        map[curr->func].push_back(curr);
      }
    }
  } gatherer(*this);
  gatherer.walkModule(&primary);
  for (auto& secondaryPtr : secondaries) {
    gatherer.walkModule(secondaryPtr.get());
  }

  // Ignore references to secondary functions that occur in the dispatch
  // segments that will contain the imported placeholders. Indirect calls to
  // table slots initialized by those segments will already go to the right
  // place once the secondary module has been loaded and the table has been
  // patched.
  std::unordered_set<RefFunc*> ignore;
  for (auto* segment : tableManager.dispatchTableSegments) {
    for (auto* expr : segment->data) {
      if (auto* ref = expr->dynCast<RefFunc>()) {
        ignore.insert(ref);
      }
    }
  }

  // Fix up what we found: Generate trampolines as described earlier, and apply
  // them.
  Builder builder(primary);
  // Generate the new trampoline function and add it to the module.
  for (auto& [name, refFuncs] : gatherer.map) {
    // Find the relevant (non-ignored) RefFuncs. If there are none, we can skip
    // creating a thunk entirely.
    std::vector<RefFunc*> relevantRefFuncs;
    for (auto* refFunc : refFuncs) {
      assert(refFunc->func == name);
      if (!ignore.contains(refFunc)) {
        relevantRefFuncs.push_back(refFunc);
      }
    }
    if (relevantRefFuncs.empty()) {
      continue;
    }

    Name trampoline = getTrampoline(name);
    // Update RefFuncs to refer to it.
    for (auto* refFunc : relevantRefFuncs) {
      refFunc->func = trampoline;
    }
  }
}

void ModuleSplitter::indirectCallsToSecondaryFunctions() {
  // Update direct calls of secondary functions to be indirect calls of their
  // corresponding table indices instead.
  struct CallIndirector : public PostWalker<CallIndirector> {
    ModuleSplitter& parent;
    std::unordered_set<Module*> dispatchTableUsingSecondaries;
    CallIndirector(ModuleSplitter& parent) : parent(parent) {}
    void visitCall(Call* curr) {
      // Return if the call's target is not in one of the secondary module.
      if (!parent.allSecondaryFuncs.contains(curr->target)) {
        return;
      }
      // Return if the current module is the same module as the call's target,
      // because we don't need a call_indirect within the same module.
      Module* currModule = getModule();
      Module* calleeModule =
        parent.secondaries.at(parent.funcToSecondaryIndex.at(curr->target))
          .get();
      if (currModule == calleeModule) {
        return;
      }

      Builder builder(*currModule);
      auto* func = calleeModule->getFunction(curr->target);
      auto tableSlot =
        parent.tableManager.getSlot(curr->target, func->type.getHeapType());
      replaceCurrent(
        builder.makeCallIndirect(tableSlot.tableName,
                                 tableSlot.makeExpr(parent.primary),
                                 curr->operands,
                                 func->type.getHeapType(),
                                 curr->isReturn));

      // Share the dispatch table with the current module (caller). We share the
      // dispatch table with with calleeModule later in setupTablePathing.
      if (currModule != &parent.primary) {
        dispatchTableUsingSecondaries.insert(currModule);
      }
    }
  };
  CallIndirector callIndirector(*this);
  callIndirector.walkModule(&primary);
  for (auto& secondaryPtr : secondaries) {
    callIndirector.walkModule(secondaryPtr.get());
  }

  for (auto* secondary : callIndirector.dispatchTableUsingSecondaries) {
    shareDispatchTable(secondary);
  }
}

void ModuleSplitter::exportImportCalledPrimaryFunctions() {
  // Find primary functions called/referred to from the secondary modules.
  using CalledPrimaryToModules = std::map<Name, std::set<Module*>>;
  struct CallCollector : PostWalker<CallCollector> {
    const std::unordered_set<Name>& primaryFuncs;
    CalledPrimaryToModules& calledPrimaryToModules;
    CallCollector(const std::unordered_set<Name>& primaryFuncs,
                  CalledPrimaryToModules& calledPrimaryToModules)
      : primaryFuncs(primaryFuncs),
        calledPrimaryToModules(calledPrimaryToModules) {}
    void visitCall(Call* curr) {
      if (primaryFuncs.contains(curr->target)) {
        calledPrimaryToModules[curr->target].insert(getModule());
      }
    }
    void visitRefFunc(RefFunc* curr) {
      if (primaryFuncs.contains(curr->func)) {
        calledPrimaryToModules[curr->func].insert(getModule());
      }
    }
  };

  for (auto& secondaryPtr : secondaries) {
    Module* secondary = secondaryPtr.get();
    ModuleUtils::ParallelFunctionAnalysis<CalledPrimaryToModules> callCollector(
      *secondary,
      [&](Function* func, CalledPrimaryToModules& calledPrimaryToModules) {
        CallCollector(primaryFuncs, calledPrimaryToModules)
          .walkFunctionInModule(func, secondary);
      });

    CalledPrimaryToModules calledPrimaryToModules;
    for (auto& [_, map] : callCollector.map) {
      calledPrimaryToModules.merge(map);
    }

    CallCollector collector(primaryFuncs, calledPrimaryToModules);
    collector.walkModuleCode(secondary);

    // Ensure each called primary function is exported and imported
    for (auto& [func, modules] : calledPrimaryToModules) {
      exportImportFunction(func, modules);
    }
  }
}

void ModuleSplitter::setupTablePatching() {
  if (!tableManager.dispatchTable) {
    return;
  }

  std::map<Module*, std::map<Index, Function*>> moduleToReplacedElems;
  Name fillerName;
  Type fillerType = Type(Signature(Type::none, Type::none), NonNullable, Exact);
  // Replace table references to secondary functions with an imported
  // placeholder that encodes the table index in its name:
  // `importNamespace`.`index`.
  forEachElement(
    primary, [&](Name table, Name, Index index, Expression*& elem) {
      auto* ref = elem->dynCast<RefFunc>();
      if (!ref) {
        return;
      }
      if (!allSecondaryFuncs.contains(ref->func)) {
        return;
      }
      assert(table == tableManager.dispatchTable->name);

      placeholderMap[table][index] = ref->func;
      Index secondaryIndex = funcToSecondaryIndex.at(ref->func);
      Module& secondary = *secondaries.at(secondaryIndex);
      Name secondaryName = config.secondaryNames.at(secondaryIndex);
      auto* secondaryFunc = secondary.getFunction(ref->func);
      moduleToReplacedElems[&secondary][index] = secondaryFunc;

      if (config.usePlaceholders) {
        auto placeholder = std::make_unique<Function>();
        placeholder->module = config.placeholderNamespacePrefix.toString() +
                              "." + secondaryName.toString();
        placeholder->base = std::to_string(index);
        placeholder->name = Names::getValidFunctionName(
          primary, std::string("placeholder_") + placeholder->base.toString());
        placeholder->hasExplicitName = true;
        placeholder->type = secondaryFunc->type.with(Inexact);
        elem =
          Builder(primary).makeRefFunc(placeholder->name, placeholder->type);
        primary.addFunction(std::move(placeholder));

      } else { // !config.usePlaceholders
        if (primary.features.hasReferenceTypes()) {
          // TODO: This can create dispatch element segments with lots of nulls.
          // We should optimize them like we do data segments with zeros.
          elem = Builder(primary).makeRefNull(HeapType::nofunc);
          return;
        }
        // When reference-types is not enabled, we can't use a ref.null. Put a
        // filler function that contains an unreachable.
        if (!fillerName) {
          fillerName = Names::getValidFunctionName(primary, "filler");
          auto filler = Builder::makeFunction(
            fillerName, fillerType, {}, Builder(primary).makeUnreachable());
          filler->hasExplicitName = true;
          primary.addFunction(std::move(filler));
        }
        elem = Builder(primary).makeRefFunc(fillerName, fillerType);
      }
    });

  if (moduleToReplacedElems.size() == 0) {
    // No placeholders to patch out of the table
    return;
  }

  for (auto& [secondaryPtr, replacedElems] : moduleToReplacedElems) {
    Module& secondary = *secondaryPtr;
    shareDispatchTable(&secondary);
    auto* secondaryTable = secondary.getTable(tableManager.dispatchTable->name);

    if (tableManager.dispatchBase.global) {
      assert(tableManager.dispatchTableSegments.size() == 1 &&
             "Unexpected number of segments with non-const base");
      assert(secondary.tables.size() == 1 && secondary.elementSegments.empty());
      // Since addition is not currently allowed in initializer expressions, we
      // need to start the new secondary segment where the primary segment
      // starts. The secondary segment will contain the same primary functions
      // as the primary module except in positions where it needs to overwrite a
      // placeholder function. All primary functions in the table therefore need
      // to be imported into the second module. TODO: use better strategies
      // here, such as using ref.func in the start function or standardizing
      // addition in initializer expressions.
      ElementSegment* primarySeg = tableManager.dispatchTableSegments.front();
      std::vector<Expression*> secondaryElems;
      secondaryElems.reserve(primarySeg->data.size());

      // Copy functions from the primary segment to the secondary segment,
      // replacing placeholders and creating new exports and imports as
      // necessary.
      auto replacement = replacedElems.begin();
      for (Index i = 0;
           i < primarySeg->data.size() && replacement != replacedElems.end();
           ++i) {
        if (replacement->first == i) {
          // primarySeg->data[i] is a placeholder, so use the secondary
          // function.
          auto* func = replacement->second;
          auto* ref = Builder(secondary).makeRefFunc(func->name, func->type);
          secondaryElems.push_back(ref);
          ++replacement;
        } else if (auto* get = primarySeg->data[i]->dynCast<RefFunc>()) {
          exportImportFunction(get->func, {&secondary});
          auto* copied =
            ExpressionManipulator::copy(primarySeg->data[i], secondary);
          secondaryElems.push_back(copied);
        }
      }

      auto offset = ExpressionManipulator::copy(primarySeg->offset, secondary);
      auto secondarySeg = std::make_unique<ElementSegment>(
        secondaryTable->name, offset, secondaryTable->type, secondaryElems);
      secondarySeg->setName(primarySeg->name, primarySeg->hasExplicitName);
      secondary.addElementSegment(std::move(secondarySeg));
      return;
    }

    // Create dispatch table segments in the secondary module to patch in the
    // original functions when it is instantiated.
    Index currBase = replacedElems.begin()->first;
    std::vector<Expression*> currData;
    auto finishSegment = [&]() {
      auto* offset = Builder(secondary).makeConst(
        Literal::makeFromInt32(currBase, secondaryTable->addressType));
      auto secondarySeg = std::make_unique<ElementSegment>(
        secondaryTable->name, offset, secondaryTable->type, currData);
      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) {
      if (curr->first != currBase + currData.size()) {
        finishSegment();
        currBase = curr->first;
        currData.clear();
      }
      auto* func = curr->second;
      currData.push_back(
        Builder(secondary).makeRefFunc(func->name, func->type));
    }
    if (currData.size()) {
      finishSegment();
    }
  }
}

} // anonymous namespace

Results splitFunctions(Module& primary, const Config& config) {
  ModuleSplitter split(primary, config);
  return {std::move(split.secondaries), std::move(split.placeholderMap)};
}

} // namespace wasm::ModuleSplitting