[NFC] RemoveUnusedModuleElements: Optimize repeated table lookups (#8048)

Each time we see a new call_indirect form (a new combo of table + type that
is called), we scan the table to see which functions might be called. In a large
Dart testcase I am looking at, we have 500K items in a table (!), and the
constant scanning of items in the segment into RefFuncs and then looking up
their functions on the module is by far the slowest part of the pass. Reduce that
overhead by precomputing a "flat", friendly form of segments, which has the
function name and type on each element segment item.

This makes the slowest pass on that testcase over 2x faster, from 16.05 to 6.85
seconds. The pass runs 3 times in -O3 so this lets us save almost 30
seconds from the total time of 184, making us overall 15% faster.

Author: kripken

src/passes/RemoveUnusedModuleElements.cpp

@@ -275,11 +275,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);
@@ -290,6 +315,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.
@@ -374,27 +419,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) {