[NFC] RemoveUnusedModuleElements: Optimize repeated table lookups

src/passes/RemoveUnusedModuleElements.cpp

@@ -268,11 +268,36 @@ struct Analyzer {
   std::unordered_map<StructField, std::vector<Expression*>>
     unreadStructFieldExprMap;
 
+  // Cached segment data. Each time we see a new indirect call, we must scan all
+  // the segments of the table it refers to, find the functions in that segment,
+  // and check their types. If the number of segments is immense, we may end up
+  // doing a massive amount of function lookups (N * M where N = number of
+  // unique indirect call forms and M = size of the table's segments). To avoid
+  // that, precompute the function lookups in advance by "flattening" the data.
+  struct FlatElemInfo {
+    // The name of the element segment.
+    Name name;
+
+    // The data in the element segment.
+    struct Item {
+      // The function the element segment's item refers to.
+      Name func;
+      // The type of function.
+      Type type;
+    };
+    std::vector<Item> data;
+  };
+  // Each table tracks all its elems.
+  using FlatTableInfo = std::vector<FlatElemInfo>;
+  std::unordered_map<Name, FlatTableInfo> flatTableInfoMap;
+
   Analyzer(Module* module,
            const PassOptions& options,
            const std::vector<ModuleElement>& roots)
     : module(module), options(options) {
 
+    prepare();
+
     // All roots are used.
     for (auto& element : roots) {
       use(element);
@@ -283,6 +308,26 @@ struct Analyzer {
     }
   }
 
+  void prepare() {
+    for (auto& elem : module->elementSegments) {
+      if (!elem->table) {
+        continue;
+      }
+      FlatElemInfo elemInfo;
+      elemInfo.name = elem->name;
+      auto& data = elemInfo.data;
+      for (auto* item : elem->data) {
+        if (auto* refFunc = item->dynCast<RefFunc>()) {
+          auto* func = module->getFunction(refFunc->func);
+          data.emplace_back(FlatElemInfo::Item{func->name, func->type});
+        }
+      }
+      if (!elemInfo.data.empty()) {
+        flatTableInfoMap[elem->table].push_back(std::move(elemInfo));
+      }
+    }
+  }
+
   // Process expressions in the expression queue while we have any, visiting
   // them (using their contents) and adding children. Returns whether we did any
   // work.
@@ -367,27 +412,21 @@ struct Analyzer {
       return;
     }
 
-    // TODO: use structured bindings with c++20, needed for the capture below
-    auto table = call.first;
-    auto type = call.second;
+    auto [table, type] = call;
 
     // Any function in the table of that signature may be called.
-    ModuleUtils::iterTableSegments(
-      *module, table, [&](ElementSegment* segment) {
-        auto segmentReferenced = false;
-        for (auto* item : segment->data) {
-          if (auto* refFunc = item->dynCast<RefFunc>()) {
-            auto* func = module->getFunction(refFunc->func);
-            if (HeapType::isSubType(func->type.getHeapType(), type)) {
-              use({ModuleElementKind::Function, refFunc->func});
-              segmentReferenced = true;
-            }
-          }
-        }
-        if (segmentReferenced) {
-          reference({ModuleElementKind::ElementSegment, segment->name});
+    for (auto& elemInfo : flatTableInfoMap[table]) {
+      auto elemReferenced = false;
+      for (auto& [func, funcType] : elemInfo.data) {
+        if (HeapType::isSubType(funcType.getHeapType(), type)) {
+          use({ModuleElementKind::Function, func});
+          elemReferenced = true;
         }
-      });
+      }
+      if (elemReferenced) {
+        reference({ModuleElementKind::ElementSegment, elemInfo.name});
+      }
+    }
   }
 
   void useRefFunc(Name func) {