From e22e0e95953064b419eba8a3bdad49bce1d9224d Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 13 Nov 2025 14:22:39 -0800
Subject: [PATCH 1/3] fast

---
 src/passes/RemoveUnusedModuleElements.cpp | 74 +++++++++++++++++------
 1 file changed, 57 insertions(+), 17 deletions(-)

diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
index cf933c639f4..0f4897ae5bf 100644
--- a/src/passes/RemoveUnusedModuleElements.cpp
+++ b/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.
@@ -368,26 +413,21 @@ struct Analyzer {
     }
 
     // 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) {

From 66defc6364bf324de1a7a68a16235698f5dcf26e Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 13 Nov 2025 14:47:19 -0800
Subject: [PATCH 2/3] todone

---
 src/passes/RemoveUnusedModuleElements.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
index 0f4897ae5bf..2af1a5bc951 100644
--- a/src/passes/RemoveUnusedModuleElements.cpp
+++ b/src/passes/RemoveUnusedModuleElements.cpp
@@ -412,7 +412,6 @@ struct Analyzer {
       return;
     }
 
-    // TODO: use structured bindings with c++20, needed for the capture below
     auto [table, type] = call;
 
     // Any function in the table of that signature may be called.

From 88c81206cf21629ff6e1e9d94b2f9d7334334adf Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 13 Nov 2025 14:47:50 -0800
Subject: [PATCH 3/3] fix

---
 src/passes/RemoveUnusedModuleElements.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
index 2af1a5bc951..7abe83d7da3 100644
--- a/src/passes/RemoveUnusedModuleElements.cpp
+++ b/src/passes/RemoveUnusedModuleElements.cpp
@@ -417,7 +417,7 @@ struct Analyzer {
     // Any function in the table of that signature may be called.
     for (auto& elemInfo : flatTableInfoMap[table]) {
       auto elemReferenced = false;
-      for (auto [func, funcType] : elemInfo.data) {
+      for (auto& [func, funcType] : elemInfo.data) {
         if (HeapType::isSubType(funcType.getHeapType(), type)) {
           use({ModuleElementKind::Function, func});
           elemReferenced = true;