[GC] Fix parsing/printing of ref types using i31 (#3469)

This lets us parse (ref null i31) and (ref i31) and not just i31ref.

It also fixes the parsing of i31ref, making it nullable for now, which we
need to do until we support non-nullability.

Fix some internal handling of i31 where we had just i31ref (which meant we
just handled the non-nullable type).

After fixing a bug in printing (where we didn't print out (ref null i31)
properly), I found some a simplification, to remove TypeName.

7 files changed, 74 insertions, 63 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 0095fb407..f33d41266 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -405,7 +405,10 @@ inline void collectHeapTypes(Module& wasm,
                              std::unordered_map<HeapType, Index>& typeIndices) {
   struct Counts : public std::unordered_map<HeapType, size_t> {
     bool isRelevant(Type type) {
-      return !type.isBasic() && (type.isRef() || type.isRtt());
+      if (type.isRef()) {
+        return !type.getHeapType().isBasic();
+      }
+      return type.isRtt();
     }
     void note(HeapType type) { (*this)[type]++; }
     void maybeNote(Type type) {
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index 8b028c25f..e1d80a11a 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -67,17 +67,10 @@ static std::ostream& printLocal(Index index, Function* func, std::ostream& o) {
   return printName(name, o);
 }
 
-// Wrapper for printing a type when we try to print the type name as much as
-// possible. For example, for a signature we will print the signature's name,
-// not its contents.
-struct TypeName {
-  Type type;
-  TypeName(Type type) : type(type) {}
-};
-
 static void
 printHeapTypeName(std::ostream& os, HeapType type, bool first = true);
 
+// Prints the name of a type. This output is guaranteed to not contain spaces.
 static void printTypeName(std::ostream& os, Type type) {
   if (type.isBasic()) {
     os << type;
@@ -131,6 +124,8 @@ static void printFieldName(std::ostream& os, const Field& field) {
   }
 }
 
+// Prints the name of a heap type. As with printTypeName, this output is
+// guaranteed to not contain spaces.
 static void printHeapTypeName(std::ostream& os, HeapType type, bool first) {
   if (type.isBasic()) {
     os << type;
@@ -174,8 +169,8 @@ struct SExprType {
   SExprType(Type type) : type(type){};
 };
 
-static std::ostream& operator<<(std::ostream& o, const SExprType& localType) {
-  Type type = localType.type;
+static std::ostream& operator<<(std::ostream& o, const SExprType& sType) {
+  Type type = sType.type;
   if (type.isTuple()) {
     o << '(';
     auto sep = "";
@@ -192,24 +187,17 @@ static std::ostream& operator<<(std::ostream& o, const SExprType& localType) {
     }
     printHeapTypeName(o, rtt.heapType);
     o << ')';
-  } else {
-    printTypeName(o, localType.type);
-  }
-  return o;
-}
-
-std::ostream& operator<<(std::ostream& os, TypeName typeName) {
-  auto type = typeName.type;
-  if (type.isRef() && !type.isBasic()) {
-    os << "(ref ";
+  } else if (type.isRef() && !type.isBasic()) {
+    o << "(ref ";
     if (type.isNullable()) {
-      os << "null ";
+      o << "null ";
     }
-    printHeapTypeName(os, type.getHeapType());
-    os << ')';
-    return os;
+    printHeapTypeName(o, type.getHeapType());
+    o << ')';
+  } else {
+    printTypeName(o, sType.type);
   }
-  return os << SExprType(typeName.type);
+  return o;
 }
 
 // TODO: try to simplify or even remove this, as we may be able to do the same
@@ -229,10 +217,10 @@ std::ostream& operator<<(std::ostream& os, ResultTypeName typeName) {
     for (auto t : type) {
       os << sep;
       sep = " ";
-      os << TypeName(t);
+      os << SExprType(t);
     }
   } else {
-    os << TypeName(type);
+    os << SExprType(type);
   }
   os << ')';
   return os;
@@ -2571,7 +2559,7 @@ struct PrintSExpression : public OverriddenVisitor<PrintSExpression> {
       o << "(param ";
       auto sep = "";
       for (auto type : curr.params) {
-        o << sep << TypeName(type);
+        o << sep << SExprType(type);
         sep = " ";
       }
       o << ')';
@@ -2581,7 +2569,7 @@ struct PrintSExpression : public OverriddenVisitor<PrintSExpression> {
       o << "(result ";
       auto sep = "";
       for (auto type : curr.results) {
-        o << sep << TypeName(type);
+        o << sep << SExprType(type);
         sep = " ";
       }
       o << ')';
@@ -2601,7 +2589,7 @@ struct PrintSExpression : public OverriddenVisitor<PrintSExpression> {
         WASM_UNREACHABLE("invalid packed type");
       }
     } else {
-      o << TypeName(field.type);
+      o << SExprType(field.type);
     }
     if (field.mutable_) {
       o << ')';
@@ -2736,7 +2724,7 @@ struct PrintSExpression : public OverriddenVisitor<PrintSExpression> {
         o << '(';
         printMinor(o, "param ");
         printLocal(i, currFunction, o);
-        o << ' ' << TypeName(param) << ')';
+        o << ' ' << SExprType(param) << ')';
         ++i;
       }
     }
@@ -2750,7 +2738,7 @@ struct PrintSExpression : public OverriddenVisitor<PrintSExpression> {
       o << '(';
       printMinor(o, "local ");
       printLocal(i, currFunction, o)
-        << ' ' << TypeName(curr->getLocalType(i)) << ')';
+        << ' ' << SExprType(curr->getLocalType(i)) << ')';
       o << maybeNewLine;
     }
     // Print the body.
diff --git a/src/support/string.h b/src/support/string.h
index b93d3b363..1708f4d94 100644
--- a/src/support/string.h
+++ b/src/support/string.h
@@ -115,6 +115,10 @@ inline std::string trim(const std::string& input) {
   return input.substr(0, size);
 }
 
+inline bool isNumber(const std::string& str) {
+  return !str.empty() && std::all_of(str.begin(), str.end(), ::isdigit);
+}
+
 } // namespace String
 
 } // namespace wasm
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index f3160c461..bcf7c15fe 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -1416,7 +1416,8 @@ Type WasmBinaryBuilder::getType(int initial) {
       // FIXME: for now, force all inputs to be nullable
       return Type(getHeapType(), Nullable);
     case BinaryConsts::EncodedType::i31ref:
-      return Type::i31ref;
+      // FIXME: for now, force all inputs to be nullable
+      return Type(HeapType::BasicHeapType::i31, Nullable);
     case BinaryConsts::EncodedType::rtt_n: {
       auto depth = getU32LEB();
       auto heapType = getHeapType();
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 20ab4f3d2..d76a3c49a 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -22,6 +22,7 @@
 
 #include "ir/branch-utils.h"
 #include "shared-constants.h"
+#include "support/string.h"
 #include "wasm-binary.h"
 #include "wasm-builder.h"
 
@@ -867,7 +868,8 @@ Type SExpressionWasmBuilder::stringToType(const char* str,
     return Type::eqref;
   }
   if (strncmp(str, "i31ref", 6) == 0 && (prefix || str[6] == 0)) {
-    return Type::i31ref;
+    // FIXME: for now, force all inputs to be nullable
+    return Type(HeapType::BasicHeapType::i31, Nullable);
   }
   if (allowError) {
     return Type::none;
@@ -2802,12 +2804,17 @@ HeapType SExpressionWasmBuilder::parseHeapType(Element& s) {
       }
       return types[it->second];
     } else {
-      // index
-      size_t offset = atoi(s.str().c_str());
-      if (offset >= types.size()) {
-        throw ParseException("unknown indexed function type", s.line, s.col);
+      // It may be a numerical index, or it may be a built-in type name like
+      // "i31".
+      auto* str = s.str().c_str();
+      if (String::isNumber(str)) {
+        size_t offset = atoi(str);
+        if (offset >= types.size()) {
+          throw ParseException("unknown indexed function type", s.line, s.col);
+        }
+        return types[offset];
       }
-      return types[offset];
+      return stringToHeapType(str, /* prefix = */ false);
     }
   }
   // It's a list.
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 2eef7eb3d..31290c745 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -481,18 +481,29 @@ Type Type::reinterpret() const {
 FeatureSet Type::getFeatures() const {
   auto getSingleFeatures = [](Type t) -> FeatureSet {
     if (t.isRef()) {
-      if (t != Type::funcref && t.isFunction()) {
-        // Strictly speaking, typed function references require the typed
-        // function references feature, however, we use these types internally
-        // regardless of the presence of features (in particular, since during
-        // load of the wasm we don't know the features yet, so we apply the more
-        // refined types).
-        return FeatureSet::ReferenceTypes;
-      }
+      // A reference type implies we need that feature. Some also require more,
+      // such as GC or exceptions.
       auto heapType = t.getHeapType();
       if (heapType.isStruct() || heapType.isArray()) {
         return FeatureSet::ReferenceTypes | FeatureSet::GC;
       }
+      if (heapType.isBasic()) {
+        switch (heapType.getBasic()) {
+          case HeapType::BasicHeapType::exn:
+            return FeatureSet::ReferenceTypes | FeatureSet::ExceptionHandling;
+          case HeapType::BasicHeapType::any:
+          case HeapType::BasicHeapType::eq:
+          case HeapType::BasicHeapType::i31:
+            return FeatureSet::ReferenceTypes | FeatureSet::GC;
+          default: {}
+        }
+      }
+      // Note: Technically typed function references also require the typed
+      // function references feature, however, we use these types internally
+      // regardless of the presence of features (in particular, since during
+      // load of the wasm we don't know the features yet, so we apply the more
+      // refined types), so we don't add that in any case here.
+      return FeatureSet::ReferenceTypes;
     } else if (t.isRtt()) {
       return FeatureSet::ReferenceTypes | FeatureSet::GC;
     }
@@ -500,15 +511,6 @@ FeatureSet Type::getFeatures() const {
     switch (t.getBasic()) {
       case Type::v128:
         return FeatureSet::SIMD;
-      case Type::funcref:
-      case Type::externref:
-        return FeatureSet::ReferenceTypes;
-      case Type::exnref:
-        return FeatureSet::ReferenceTypes | FeatureSet::ExceptionHandling;
-      case Type::anyref:
-      case Type::eqref:
-      case Type::i31ref:
-        return FeatureSet::ReferenceTypes | FeatureSet::GC;
       default:
         return FeatureSet::MVP;
     }
@@ -594,17 +596,21 @@ bool Type::isSubType(Type left, Type right) {
       return true;
     }
     // Various things are subtypes of eqref.
-    if ((left == Type::i31ref || left.getHeapType().isArray() ||
-         left.getHeapType().isStruct()) &&
-        right == Type::eqref) {
+    auto leftHeap = left.getHeapType();
+    auto rightHeap = right.getHeapType();
+    if ((leftHeap == HeapType::i31 || leftHeap.isArray() ||
+         leftHeap.isStruct()) &&
+        rightHeap == HeapType::eq &&
+        (!left.isNullable() || right.isNullable())) {
       return true;
     }
     // All typed function signatures are subtypes of funcref.
-    if (left.getHeapType().isSignature() && right == Type::funcref) {
+    if (leftHeap.isSignature() && rightHeap == HeapType::func &&
+        (!left.isNullable() || right.isNullable())) {
       return true;
     }
     // A non-nullable type is a supertype of a nullable one
-    if (left.getHeapType() == right.getHeapType() && !left.isNullable()) {
+    if (leftHeap == rightHeap && !left.isNullable()) {
       // The only difference is the nullability.
       assert(right.isNullable());
       return true;
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index a45be1fd3..06f77b093 100644
--- a/src/wasm/wasm-validator.cpp
+++ b/src/wasm/wasm-validator.cpp
@@ -2191,9 +2191,11 @@ void FunctionValidator::visitI31Get(I31Get* curr) {
   shouldBeTrue(getModule()->features.hasGC(),
                curr,
                "i31.get_s/u requires gc to be enabled");
+  // FIXME: use i31ref here, which is non-nullable, when we support non-
+  // nullability.
   shouldBeSubTypeOrFirstIsUnreachable(
     curr->i31->type,
-    Type::i31ref,
+    Type(HeapType::i31, Nullable),
     curr->i31,
     "i31.get_s/u's argument should be i31ref");
 }