Emit heap types for call_indirect that match the table (#4221)

See #4220 - this lets us handle the common case for now of simply having
an identical heap type to the table when the signature is identical.

With this PR, #4207's optimization of call_ref + table.get into
call_indirect now leads to a binary that works in V8 in nominal mode.

7 files changed, 137 insertions, 9 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.
diff --git a/test/lit/passes/directize_all-features.wast b/test/lit/passes/directize_all-features.wast
index c792797af..9220b24e2 100644
--- a/test/lit/passes/directize_all-features.wast
+++ b/test/lit/passes/directize_all-features.wast
@@ -821,10 +821,10 @@
  (table $no-set 5 5 funcref)
 
  ;; CHECK:      (elem $0 (table $has-set) (i32.const 1) func $foo)
- (elem (table $has-set) (i32.const 1) $foo)
+ (elem $0 (table $has-set) (i32.const 1) $foo)
 
  ;; CHECK:      (elem $1 (table $no-set) (i32.const 1) func $foo)
- (elem (table $no-set) (i32.const 1) $foo)
+ (elem $1 (table $no-set) (i32.const 1) $foo)
 
  ;; CHECK:      (func $foo
  ;; CHECK-NEXT:  (table.set $has-set
diff --git a/test/lit/passes/optimize-instructions-call_ref-roundtrip.wast b/test/lit/passes/optimize-instructions-call_ref-roundtrip.wast
new file mode 100644
index 000000000..5bbec4b64
--- /dev/null
+++ b/test/lit/passes/optimize-instructions-call_ref-roundtrip.wast
@@ -0,0 +1,88 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_test.py and could be cleaned up.
+
+;; RUN: wasm-opt %s --optimize-instructions --nominal --roundtrip --all-features -S -o - | filecheck %s
+;; roundtrip to see the effects on heap types in the binary format, specifically
+;; regarding nominal heap types
+
+(module
+ ;; Create multiple different signature types, all identical structurally but
+ ;; distinct nominally. The three tables will use different ones, and the
+ ;; emitted call_indirects should use the corresponding ones.
+
+ ;; CHECK:      (type $v1 (func))
+ (type $v1 (func))
+
+ ;; CHECK:      (type $v2 (func))
+ (type $v2 (func))
+
+ ;; CHECK:      (type $v3 (func))
+ (type $v3 (func))
+
+ ;; CHECK:      (type $i32_=>_none (func (param i32)))
+
+ ;; CHECK:      (table $table-1 10 (ref null $v1))
+ (table $table-1 10 (ref null $v1))
+
+ ;; CHECK:      (table $table-2 10 (ref null $v2))
+ (table $table-2 10 (ref null $v2))
+
+ ;; CHECK:      (table $table-3 10 (ref null $v3))
+ (table $table-3 10 (ref null $v3))
+
+ ;; CHECK:      (elem $elem-1 (table $table-1) (i32.const 0) (ref null $v1) (ref.func $helper-1))
+ (elem $elem-1 (table $table-1) (i32.const 0) (ref null $v1)
+  (ref.func $helper-1))
+
+ ;; CHECK:      (elem $elem-2 (table $table-2) (i32.const 0) (ref null $v2) (ref.func $helper-2))
+ (elem $elem-2 (table $table-2) (i32.const 0) (ref null $v2)
+  (ref.func $helper-2))
+
+ ;; CHECK:      (elem $elem-3 (table $table-3) (i32.const 0) (ref null $v3) (ref.func $helper-3))
+ (elem $elem-3 (table $table-3) (i32.const 0) (ref null $v3)
+  (ref.func $helper-3))
+
+ ;; CHECK:      (func $helper-1
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $helper-1 (type $v1))
+ ;; CHECK:      (func $helper-2
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $helper-2 (type $v2))
+ ;; CHECK:      (func $helper-3
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $helper-3 (type $v3))
+
+ ;; CHECK:      (func $call-table-get (param $x i32)
+ ;; CHECK-NEXT:  (call_indirect $table-1 (type $v1)
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call_indirect $table-2 (type $v2)
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call_indirect $table-3 (type $v3)
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-table-get (param $x i32)
+  ;; The heap type of the call_indirects that we emit here should be the
+  ;; identical one as on the table that they correspond to.
+  (call_ref
+   (table.get $table-1
+    (local.get $x)
+   )
+  )
+  (call_ref
+   (table.get $table-2
+    (local.get $x)
+   )
+  )
+  (call_ref
+   (table.get $table-3
+    (local.get $x)
+   )
+  )
+ )
+)