2 files changed, 53 insertions, 44 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index a87daee9e..666718c92 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -476,65 +476,63 @@ inline void collectHeapTypes(Module& wasm,
     }
   };
 
-  // Collect the type use counts for a single function
-  auto updateCounts = [&](Function* func, Counts& counts) {
-    if (func->imported()) {
-      return;
-    }
-    struct TypeCounter
-      : PostWalker<TypeCounter, UnifiedExpressionVisitor<TypeCounter>> {
-      Counts& counts;
-
-      TypeCounter(Counts& counts) : counts(counts) {}
-
-      void visitExpression(Expression* curr) {
-        if (auto* call = curr->dynCast<CallIndirect>()) {
-          counts.note(call->sig);
-        } else if (curr->is<RefNull>()) {
+  struct CodeScanner
+    : PostWalker<CodeScanner, UnifiedExpressionVisitor<CodeScanner>> {
+    Counts& counts;
+
+    CodeScanner(Counts& counts) : counts(counts) {}
+
+    void visitExpression(Expression* curr) {
+      if (auto* call = curr->dynCast<CallIndirect>()) {
+        counts.note(call->sig);
+      } else if (curr->is<RefNull>()) {
+        counts.maybeNote(curr->type);
+      } else if (curr->is<RttCanon>() || curr->is<RttSub>()) {
+        counts.note(curr->type.getRtt().heapType);
+      } else if (auto* get = curr->dynCast<StructGet>()) {
+        counts.maybeNote(get->ref->type);
+      } else if (auto* set = curr->dynCast<StructSet>()) {
+        counts.maybeNote(set->ref->type);
+      } else if (Properties::isControlFlowStructure(curr)) {
+        if (curr->type.isTuple()) {
+          // TODO: Allow control flow to have input types as well
+          counts.note(Signature(Type::none, curr->type));
+        } else {
           counts.maybeNote(curr->type);
-        } else if (curr->is<RttCanon>() || curr->is<RttSub>()) {
-          counts.note(curr->type.getRtt().heapType);
-        } else if (auto* get = curr->dynCast<StructGet>()) {
-          counts.maybeNote(get->ref->type);
-        } else if (auto* set = curr->dynCast<StructSet>()) {
-          counts.maybeNote(set->ref->type);
-        } else if (Properties::isControlFlowStructure(curr)) {
-          if (curr->type.isTuple()) {
-            // TODO: Allow control flow to have input types as well
-            counts.note(Signature(Type::none, curr->type));
-          } else {
-            counts.maybeNote(curr->type);
-          }
         }
       }
-    };
-    TypeCounter(counts).walk(func->body);
+    }
   };
 
-  ModuleUtils::ParallelFunctionAnalysis<Counts> analysis(wasm, updateCounts);
-
-  // Collect all the counts.
+  // Collect module-level info.
   Counts counts;
-  for (auto& curr : wasm.functions) {
-    counts.note(curr->sig);
-    for (auto type : curr->vars) {
-      for (auto t : type) {
-        counts.maybeNote(t);
-      }
-    }
-  }
+  CodeScanner(counts).walkModuleCode(&wasm);
   for (auto& curr : wasm.events) {
     counts.note(curr->sig);
   }
-  for (auto& curr : wasm.globals) {
-    counts.maybeNote(curr->type);
-  }
+
+  // Collect info from functions in parallel.
+  ModuleUtils::ParallelFunctionAnalysis<Counts> analysis(
+    wasm, [&](Function* func, Counts& counts) {
+      counts.note(func->sig);
+      for (auto type : func->vars) {
+        for (auto t : type) {
+          counts.maybeNote(t);
+        }
+      }
+      if (!func->imported()) {
+        CodeScanner(counts).walk(func->body);
+      }
+    });
+
+  // Combine the function info with the module info.
   for (auto& pair : analysis.map) {
     Counts& functionCounts = pair.second;
     for (auto& innerPair : functionCounts) {
       counts[innerPair.first] += innerPair.second;
     }
   }
+
   // A generic utility to traverse the child types of a type.
   // TODO: work with tlively to refactor this to a shared place
   auto walkRelevantChildren = [&](HeapType type, auto callback) {
diff --git a/test/lit/passes/roundtrip-table.wast b/test/lit/passes/roundtrip-table.wast
new file mode 100644
index 000000000..21eb1197c
--- /dev/null
+++ b/test/lit/passes/roundtrip-table.wast
@@ -0,0 +1,11 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s -all --roundtrip -S -o -
+
+(module
+ (type $functype (func))
+ (table $0 48 funcref)
+ ;; a type that appears in the table and nowhere else. this test checks that
+ ;; we do not crash during the roundtrip on seeing an unexpected type that
+ ;; collectHeapTypes() did not scan.
+ (elem (table $0) (i32.const 0) funcref (ref.null $functype))
+)