5 files changed, 47 insertions, 7 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 5d23278fa..13fa76fff 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -488,11 +488,13 @@ inline void collectHeapTypes(Module& wasm,
     : PostWalker<CodeScanner, UnifiedExpressionVisitor<CodeScanner>> {
     Counts& counts;
 
-    CodeScanner(Counts& counts) : counts(counts) {}
+    CodeScanner(Module& wasm, Counts& counts) : counts(counts) {
+      setModule(&wasm);
+    }
 
     void visitExpression(Expression* curr) {
       if (auto* call = curr->dynCast<CallIndirect>()) {
-        counts.note(call->sig);
+        counts.note(call->getHeapType(getModule()));
       } else if (curr->is<RefNull>()) {
         counts.note(curr->type);
       } else if (curr->is<RttCanon>() || curr->is<RttSub>()) {
@@ -542,7 +544,7 @@ inline void collectHeapTypes(Module& wasm,
 
   // Collect module-level info.
   Counts counts;
-  CodeScanner(counts).walkModuleCode(&wasm);
+  CodeScanner(wasm, counts).walkModuleCode(&wasm);
   for (auto& curr : wasm.tags) {
     counts.note(curr->sig);
   }
@@ -561,7 +563,7 @@ inline void collectHeapTypes(Module& wasm,
         counts.note(type);
       }
       if (!func->imported()) {
-        CodeScanner(counts).walk(func->body);
+        CodeScanner(wasm, counts).walk(func->body);
       }
     });
 
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index a5db73132..bea704b1d 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -435,7 +435,9 @@ struct PrintExpressionContents
 
     o << '(';
     printMinor(o, "type ");
-    TypeNamePrinter(o, wasm).print(HeapType(curr->sig));
+
+    TypeNamePrinter(o, wasm).print(curr->getHeapType(wasm));
+
     o << ')';
   }
   void visitLocalGet(LocalGet* curr) {
diff --git a/src/wasm.h b/src/wasm.h
index fbec9247a..c87b8c197 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -556,6 +556,9 @@ enum BrOnOp {
   BrOnNonI31,
 };
 
+// Forward declaration for methods that receive a Module as a parameter.
+class Module;
+
 //
 // Expressions
 //
@@ -828,6 +831,15 @@ public:
   bool isReturn = false;
 
   void finalize();
+
+  // FIXME We should probably store a heap type here, and not a signature, see
+  //       https://github.com/WebAssembly/binaryen/issues/4220
+  //       For now, copy the heap type from the table if it matches - then a
+  //       nominal check will succeed too. If it does not match, then just
+  //       emit something for it like we always used to, using
+  //       HeapType(sig) (also do that if no module is provided).
+  // FIXME When we remove this, also remove the forward decl of Module, above.
+  HeapType getHeapType(Module* module = nullptr);
 };
 
 class LocalGet : public SpecificExpression<Expression::LocalGetId> {
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index b74b73ba1..8f6581956 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -81,10 +81,10 @@ void BinaryInstWriter::visitCall(Call* curr) {
 
 void BinaryInstWriter::visitCallIndirect(CallIndirect* curr) {
   Index tableIdx = parent.getTableIndex(curr->table);
-
   int8_t op =
     curr->isReturn ? BinaryConsts::RetCallIndirect : BinaryConsts::CallIndirect;
-  o << op << U32LEB(parent.getTypeIndex(curr->sig)) << U32LEB(tableIdx);
+  o << op << U32LEB(parent.getTypeIndex(curr->getHeapType(parent.getModule())))
+    << U32LEB(tableIdx);
 }
 
 void BinaryInstWriter::visitLocalGet(LocalGet* curr) {
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 26e6b369b..4975f4597 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -300,6 +300,30 @@ void CallIndirect::finalize() {
   }
 }
 
+HeapType CallIndirect::getHeapType(Module* module) {
+  auto heapType = HeapType(sig);
+  // See comment in wasm.h
+  if (module) {
+    // The table may not yet exist if the wasm module is still being
+    // constructed. This should perhaps be an error, but as this is a hack for
+    // the time being, handle this the same as the case where module is null.
+    // Note: table_ (with underscore) is needed as |table| is a field on |this|.
+    if (auto* table_ = module->getTableOrNull(table)) {
+      // The wasm spec may allow more things eventually, and if so we'd need to
+      // add more checking here.
+      assert(table_->type.isRef());
+      auto tableHeapType = table_->type.getHeapType();
+      if (tableHeapType.isSignature()) {
+        auto tableSig = tableHeapType.getSignature();
+        if (sig == tableSig) {
+          heapType = tableHeapType;
+        }
+      }
+    }
+  }
+  return heapType;
+}
+
 bool LocalSet::isTee() const { return type != Type::none; }
 
 // Changes to local.tee. The type of the local should be given.