Preserve Function HeapTypes (#3952)

When using nominal types, func.ref of two functions with identical signatures
but different HeapTypes will yield different types. To preserve these semantics,
Functions need to track their HeapTypes, not just their Signatures.

This PR replaces the Signature field in Function with a HeapType field and adds
new utility methods to make it almost as simple to update and query the function
HeapType as it was to update and query the Function Signature.

6 files changed, 17 insertions, 14 deletions

diff --git a/test/example/local-graph.cpp b/test/example/local-graph.cpp
index 0ac435883..5f82982e3 100644
--- a/test/example/local-graph.cpp
+++ b/test/example/local-graph.cpp
@@ -12,6 +12,7 @@ int main() {
 
   {
     Function foo;
+    foo.type = Signature(Type::none, Type::none);
     foo.vars = {Type::i32};
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(0, Type::i32);
@@ -27,6 +28,7 @@ int main() {
 
   {
     Function foo;
+    foo.type = Signature(Type::none, Type::none);
     foo.vars = {Type::i32};
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(0, Type::i32);
@@ -42,7 +44,7 @@ int main() {
 
   {
     Function foo;
-    foo.sig = Signature({Type::i32}, Type::none);
+    foo.type = Signature({Type::i32}, Type::none);
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(0, Type::i32);
     foo.body = builder.makeBlock({
@@ -56,7 +58,7 @@ int main() {
 
   {
     Function foo;
-    foo.sig = Signature({Type::i32}, Type::none);
+    foo.type = Signature({Type::i32}, Type::none);
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(0, Type::i32);
     foo.body = builder.makeBlock({
@@ -71,7 +73,7 @@ int main() {
 
   {
     Function foo;
-    foo.sig = Signature({Type::i32, Type::i32}, Type::none);
+    foo.type = Signature({Type::i32, Type::i32}, Type::none);
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(1, Type::i32);
     foo.body = builder.makeBlock({
@@ -85,6 +87,7 @@ int main() {
 
   {
     Function foo;
+    foo.type = Signature(Type::none, Type::none);
     foo.vars = {Type::i32, Type::i32};
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(1, Type::i32);
@@ -100,6 +103,7 @@ int main() {
 
   {
     Function foo;
+    foo.type = Signature(Type::none, Type::none);
     foo.vars = {Type::i32, Type::f64};
     auto* get1 = builder.makeLocalGet(0, Type::i32);
     auto* get2 = builder.makeLocalGet(1, Type::f64);
diff --git a/test/example/module-splitting.cpp b/test/example/module-splitting.cpp
index bbe57a657..95d6de44c 100644
--- a/test/example/module-splitting.cpp
+++ b/test/example/module-splitting.cpp
@@ -451,10 +451,12 @@ void test_minimized_exports() {
   Expression* callBody = nullptr;
 
   Builder builder(primary);
+  auto funcType = Signature(Type::none, Type::none);
 
   for (size_t i = 0; i < 10; ++i) {
     Name name = std::to_string(i);
-    primary.addFunction(Builder::makeFunction(name, {}, {}, builder.makeNop()));
+    primary.addFunction(
+      Builder::makeFunction(name, funcType, {}, builder.makeNop()));
     keep.insert(name);
     callBody =
       builder.blockify(callBody, builder.makeCall(name, {}, Type::none));
@@ -469,7 +471,7 @@ void test_minimized_exports() {
     }
   }
 
-  primary.addFunction(Builder::makeFunction("call", {}, {}, callBody));
+  primary.addFunction(Builder::makeFunction("call", funcType, {}, callBody));
 
   ModuleSplitting::Config config;
   config.primaryFuncs = std::move(keep);
diff --git a/test/example/type-builder-nominal.cpp b/test/example/type-builder-nominal.cpp
index bd0791ae9..f3eed5810 100644
--- a/test/example/type-builder-nominal.cpp
+++ b/test/example/type-builder-nominal.cpp
@@ -110,10 +110,10 @@ void test_signatures(bool warm) {
   builder[1] = Signature(tempRef, tempRef);
   std::vector<HeapType> built = builder.build();
 
-  HeapType small = HeapType(Signature(Type::anyref, Type::i31ref));
+  HeapType small = Signature(Type::anyref, Type::i31ref);
   HeapType big =
-    HeapType(Signature(Type(Signature(Type::anyref, Type::i31ref), Nullable),
-                       Type(Signature(Type::anyref, Type::i31ref), Nullable)));
+    Signature(Type(Signature(Type::anyref, Type::i31ref), Nullable),
+              Type(Signature(Type::anyref, Type::i31ref), Nullable));
   if (warm) {
     assert(built[0] != small);
     assert(built[1] != big);
diff --git a/test/example/type-builder.cpp b/test/example/type-builder.cpp
index 8d87c5054..d86643cf0 100644
--- a/test/example/type-builder.cpp
+++ b/test/example/type-builder.cpp
@@ -121,7 +121,7 @@ void test_basic() {
 
   std::vector<HeapType> built = builder.build();
 
-  assert(built[0] == HeapType(Signature(Type::anyref, Type::i31ref)));
+  assert(built[0] == Signature(Type::anyref, Type::i31ref));
   assert(built[1] == built[0]);
   assert(built[2] == built[1]);
   assert(built[3] == built[2]);
diff --git a/test/example/typeinfo.cpp b/test/example/typeinfo.cpp
index f4306b391..ae40ad188 100644
--- a/test/example/typeinfo.cpp
+++ b/test/example/typeinfo.cpp
@@ -152,7 +152,7 @@ void test_printing() {
     std::cout << HeapType(HeapType::i31) << "\n";
     std::cout << Type(HeapType::i31, Nullable) << "\n";
     std::cout << Type(HeapType::i31, NonNullable) << "\n";
-    std::cout << HeapType(Signature(Type::none, Type::none)) << "\n";
+    std::cout << Signature(Type::none, Type::none) << "\n";
     std::cout << HeapType(Struct{}) << "\n";
     std::cout << HeapType(Array({Type::i32, Immutable})) << "\n";
   }
diff --git a/test/lit/nominal-func.wast b/test/lit/nominal-func.wast
index 4aeef2eb7..b5f266db3 100644
--- a/test/lit/nominal-func.wast
+++ b/test/lit/nominal-func.wast
@@ -1,12 +1,9 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
 ;; RUN: wasm-opt %s -all --nominal -S -o - | filecheck %s
-
-;; TODO: Preserve function heap types through round tripping as well.
-;; RUNX: wasm-opt %s -all --nominal --roundtrip -S -o - | filecheck %s
+;; RUN: wasm-opt %s -all --nominal --roundtrip -S -o - | filecheck %s
 
 (module
   ;; This will be the "canonical" function type rather than $foo_t
-  ;; CHECK:      (type $bad_t (func))
   (type $bad_t (func))
 
   ;; CHECK:      (type $foo_t (func))