[GC types] Refactoring to allow future heap type parsing. NFC (#3409)

Defined types in wasm are really one of the "heap types": a signature type, or
(with GC) a struct or an array type. This refactors the binary and text parsers
to load the defined types into an array of heap types, so that we can start to
parse GC types. This replaces the existing array of signature types (which
could not support a struct or an array).

Locally this PR can parse and print as text simple GC types. For that it was
necessary to also fix Type::getFeatures for GC.

5 files changed, 105 insertions, 110 deletions

diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index b0f41e69c..532cfb7ec 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1220,8 +1220,8 @@ class WasmBinaryBuilder {
 
   std::set<BinaryConsts::Section> seenSections;
 
-  // All signatures present in the type section
-  std::vector<Signature> signatures;
+  // All types defined in the type section
+  std::vector<HeapType> types;
 
 public:
   WasmBinaryBuilder(Module& wasm, const std::vector<char>& input)
@@ -1261,7 +1261,7 @@ public:
   void readHeader();
   void readStart();
   void readMemory();
-  void readSignatures();
+  void readTypes();
 
   // gets a name in the combined import+defined space
   Name getFunctionName(Index index);
@@ -1279,7 +1279,8 @@ public:
   std::vector<Signature> functionSignatures;
 
   void readFunctionSignatures();
-  Signature getFunctionSignatureByIndex(Index index);
+  Signature getSignatureByFunctionIndex(Index index);
+  Signature getSignatureByTypeIndex(Index index);
 
   size_t nextLabel;
 
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 9a501171d..88237249a 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -117,8 +117,11 @@ class SExpressionWasmBuilder {
   Module& wasm;
   MixedArena& allocator;
   IRProfile profile;
-  std::vector<Signature> signatures;
-  std::unordered_map<std::string, size_t> signatureIndices;
+
+  // The main list of types declared in the module
+  std::vector<HeapType> types;
+  std::unordered_map<std::string, size_t> typeIndices;
+
   std::vector<Name> functionNames;
   std::vector<Name> globalNames;
   std::vector<Name> eventNames;
@@ -149,8 +152,6 @@ private:
 
   UniqueNameMapper nameMapper;
 
-  // Given a function signature type's name, return the signature
-  Signature getFunctionSignature(Element& s);
   Name getFunctionName(Element& s);
   Name getGlobalName(Element& s);
   Name getEventName(Element& s);
@@ -295,7 +296,10 @@ private:
   void parseTable(Element& s, bool preParseImport = false);
   void parseElem(Element& s);
   void parseInnerElem(Element& s, Index i = 1, Expression* offset = nullptr);
-  Signature parseInlineFunctionSignature(Element& s);
+
+  // Parses something like (func ..), (array ..), (struct)
+  HeapType parseHeapType(Element& s);
+
   void parseType(Element& s);
   void parseEvent(Element& s, bool preParseImport = false);
 
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 170aa17bb..7a6d4d35d 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -1121,7 +1121,7 @@ void WasmBinaryBuilder::read() {
         readMemory();
         break;
       case BinaryConsts::Section::Type:
-        readSignatures();
+        readTypes();
         break;
       case BinaryConsts::Section::Import:
         readImports();
@@ -1324,11 +1324,8 @@ Type WasmBinaryBuilder::getType() {
   int type = getS32LEB();
   // Single value types are negative; signature indices are non-negative
   if (type >= 0) {
-    // TODO: Handle block input types properly
-    if (size_t(type) >= signatures.size()) {
-      throwError("invalid signature index: " + std::to_string(type));
-    }
-    return signatures[type].results;
+    // TODO: Handle block input types properly.
+    return getSignatureByTypeIndex(type).results;
   }
   switch (type) {
     // None only used for block signatures. TODO: Separate out?
@@ -1371,10 +1368,10 @@ HeapType WasmBinaryBuilder::getHeapType() {
   int type = getS32LEB(); // TODO: Actually encoded as s33
   // Single heap types are negative; heap type indices are non-negative
   if (type >= 0) {
-    if (size_t(type) >= signatures.size()) {
+    if (size_t(type) >= types.size()) {
       throwError("invalid signature index: " + std::to_string(type));
     }
-    return HeapType(signatures[type]);
+    return types[type];
   }
   switch (type) {
     case BinaryConsts::EncodedHeapType::func:
@@ -1485,8 +1482,8 @@ void WasmBinaryBuilder::readMemory() {
                      Memory::kUnlimitedSize);
 }
 
-void WasmBinaryBuilder::readSignatures() {
-  BYN_TRACE("== readSignatures\n");
+void WasmBinaryBuilder::readTypes() {
+  BYN_TRACE("== readTypes\n");
   size_t numTypes = getU32LEB();
   BYN_TRACE("num: " << numTypes << std::endl);
   for (size_t i = 0; i < numTypes; i++) {
@@ -1507,7 +1504,7 @@ void WasmBinaryBuilder::readSignatures() {
     for (size_t j = 0; j < numResults; j++) {
       results.push_back(getConcreteType());
     }
-    signatures.emplace_back(Type(params), Type(results));
+    types.emplace_back(Signature(Type(params), Type(results)));
   }
 }
 
@@ -1576,11 +1573,8 @@ void WasmBinaryBuilder::readImports() {
       case ExternalKind::Function: {
         Name name(std::string("fimport$") + std::to_string(functionCounter++));
         auto index = getU32LEB();
-        if (index >= signatures.size()) {
-          throwError("invalid function index " + std::to_string(index) + " / " +
-                     std::to_string(signatures.size()));
-        }
-        auto curr = builder.makeFunction(name, signatures[index], {});
+        auto curr =
+          builder.makeFunction(name, getSignatureByTypeIndex(index), {});
         curr->module = module;
         curr->base = base;
         functionImports.push_back(curr.get());
@@ -1645,11 +1639,8 @@ void WasmBinaryBuilder::readImports() {
         Name name(std::string("eimport$") + std::to_string(eventCounter++));
         auto attribute = getU32LEB();
         auto index = getU32LEB();
-        if (index >= signatures.size()) {
-          throwError("invalid event index " + std::to_string(index) + " / " +
-                     std::to_string(signatures.size()));
-        }
-        auto curr = builder.makeEvent(name, attribute, signatures[index]);
+        auto curr =
+          builder.makeEvent(name, attribute, getSignatureByTypeIndex(index));
         curr->module = module;
         curr->base = base;
         wasm.addEvent(std::move(curr));
@@ -1680,14 +1671,11 @@ void WasmBinaryBuilder::readFunctionSignatures() {
   for (size_t i = 0; i < num; i++) {
     BYN_TRACE("read one\n");
     auto index = getU32LEB();
-    if (index >= signatures.size()) {
-      throwError("invalid function type index for function");
-    }
-    functionSignatures.push_back(signatures[index]);
+    functionSignatures.push_back(getSignatureByTypeIndex(index));
   }
 }
 
-Signature WasmBinaryBuilder::getFunctionSignatureByIndex(Index index) {
+Signature WasmBinaryBuilder::getSignatureByFunctionIndex(Index index) {
   Signature sig;
   if (index < functionImports.size()) {
     return functionImports[index]->sig;
@@ -1699,6 +1687,18 @@ Signature WasmBinaryBuilder::getFunctionSignatureByIndex(Index index) {
   return functionSignatures[adjustedIndex];
 }
 
+Signature WasmBinaryBuilder::getSignatureByTypeIndex(Index index) {
+  if (index >= types.size()) {
+    throwError("invalid type index " + std::to_string(index) + " / " +
+               std::to_string(types.size()));
+  }
+  auto heapType = types[index];
+  if (!heapType.isSignature()) {
+    throwError("invalid signature type " + heapType.toString());
+  }
+  return heapType.getSignature();
+}
+
 void WasmBinaryBuilder::readFunctions() {
   BYN_TRACE("== readFunctions\n");
   size_t total = getU32LEB();
@@ -2352,12 +2352,9 @@ void WasmBinaryBuilder::readEvents() {
     BYN_TRACE("read one\n");
     auto attribute = getU32LEB();
     auto typeIndex = getU32LEB();
-    if (typeIndex >= signatures.size()) {
-      throwError("invalid event index " + std::to_string(typeIndex) + " / " +
-                 std::to_string(signatures.size()));
-    }
-    wasm.addEvent(Builder::makeEvent(
-      "event$" + std::to_string(i), attribute, signatures[typeIndex]));
+    wasm.addEvent(Builder::makeEvent("event$" + std::to_string(i),
+                                     attribute,
+                                     getSignatureByTypeIndex(typeIndex)));
   }
 }
 
@@ -3173,7 +3170,7 @@ void WasmBinaryBuilder::visitSwitch(Switch* curr) {
 void WasmBinaryBuilder::visitCall(Call* curr) {
   BYN_TRACE("zz node: Call\n");
   auto index = getU32LEB();
-  auto sig = getFunctionSignatureByIndex(index);
+  auto sig = getSignatureByFunctionIndex(index);
   auto num = sig.params.size();
   curr->operands.resize(num);
   for (size_t i = 0; i < num; i++) {
@@ -3187,10 +3184,7 @@ void WasmBinaryBuilder::visitCall(Call* curr) {
 void WasmBinaryBuilder::visitCallIndirect(CallIndirect* curr) {
   BYN_TRACE("zz node: CallIndirect\n");
   auto index = getU32LEB();
-  if (index >= signatures.size()) {
-    throwError("bad call_indirect function index");
-  }
-  curr->sig = signatures[index];
+  curr->sig = getSignatureByTypeIndex(index);
   auto reserved = getU32LEB();
   if (reserved != 0) {
     throwError("Invalid flags field in call_indirect");
@@ -5293,8 +5287,8 @@ void WasmBinaryBuilder::visitRefFunc(RefFunc* curr) {
   // To support typed function refs, we give the reference not just a general
   // funcref, but a specific subtype with the actual signature.
   // FIXME: for now, emit a nullable type here
-  curr->finalize(
-    Type(HeapType(getFunctionSignatureByIndex(index)), /* nullable = */ true));
+  curr->finalize(Type(HeapType(getSignatureByFunctionIndex(index)),
+                      /* nullable = */ true));
 }
 
 void WasmBinaryBuilder::visitRefEq(RefEq* curr) {
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 434dec9d7..47ef33371 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -449,25 +449,6 @@ Name SExpressionWasmBuilder::getFunctionName(Element& s) {
   }
 }
 
-Signature SExpressionWasmBuilder::getFunctionSignature(Element& s) {
-  if (s.dollared()) {
-    auto it = signatureIndices.find(s.str().str);
-    if (it == signatureIndices.end()) {
-      throw ParseException(
-        "unknown function type in getFunctionSignature", s.line, s.col);
-    }
-    return signatures[it->second];
-  } else {
-    // index
-    size_t offset = atoi(s.str().c_str());
-    if (offset >= signatures.size()) {
-      throw ParseException(
-        "unknown function type in getFunctionSignature", s.line, s.col);
-    }
-    return signatures[offset];
-  }
-}
-
 Name SExpressionWasmBuilder::getGlobalName(Element& s) {
   if (s.dollared()) {
     return s.str();
@@ -564,7 +545,11 @@ Signature SExpressionWasmBuilder::parseTypeRef(Element& s) {
   if (s.size() != 2) {
     throw ParseException("invalid type reference", s.line, s.col);
   }
-  return getFunctionSignature(*s[1]);
+  auto heapType = parseHeapType(*s[1]);
+  if (!heapType.isSignature()) {
+    throw ParseException("expected signature type", s.line, s.col);
+  }
+  return heapType.getSignature();
 }
 
 // Prases typeuse, a reference to a type definition. It is in the form of either
@@ -626,9 +611,9 @@ SExpressionWasmBuilder::parseTypeUse(Element& s,
   }
 
   // Add implicitly defined type to global list so it has an index
-  if (std::find(signatures.begin(), signatures.end(), functionSignature) ==
-      signatures.end()) {
-    signatures.push_back(functionSignature);
+  auto heapType = HeapType(functionSignature);
+  if (std::find(types.begin(), types.end(), heapType) == types.end()) {
+    types.push_back(heapType);
   }
 
   // If only (type) is specified, populate `namedParams`
@@ -945,20 +930,7 @@ Type SExpressionWasmBuilder::elementToType(Element& s) {
       nullable = true;
       i++;
     }
-    Signature sig;
-    auto& last = *s[i];
-    if (last.isStr()) {
-      // A string name of a signature.
-      sig = getFunctionSignature(last);
-    } else {
-      // A signature written out in full in-line.
-      if (*last[0] != FUNC) {
-        throw ParseException(
-          std::string("invalid reference type type"), s.line, s.col);
-      }
-      sig = parseInlineFunctionSignature(last);
-    }
-    return Type(HeapType(sig), nullable);
+    return Type(parseHeapType(*s[i]), nullable);
   }
   // It's a tuple.
   std::vector<Type> types;
@@ -2784,36 +2756,54 @@ void SExpressionWasmBuilder::parseInnerElem(Element& s,
   wasm.table.segments.push_back(segment);
 }
 
-Signature SExpressionWasmBuilder::parseInlineFunctionSignature(Element& s) {
-  if (*s[0] != FUNC) {
-    throw ParseException("invalid inline function signature", s.line, s.col);
-  }
-  std::vector<Type> params;
-  std::vector<Type> results;
-  for (size_t k = 1; k < s.size(); k++) {
-    Element& curr = *s[k];
-    if (elementStartsWith(curr, PARAM)) {
-      auto newParams = parseParamOrLocal(curr);
-      params.insert(params.end(), newParams.begin(), newParams.end());
-    } else if (elementStartsWith(curr, RESULT)) {
-      auto newResults = parseResults(curr);
-      results.insert(results.end(), newResults.begin(), newResults.end());
+HeapType SExpressionWasmBuilder::parseHeapType(Element& s) {
+  if (s.isStr()) {
+    // It's a string.
+    if (s.dollared()) {
+      auto it = typeIndices.find(s.str().str);
+      if (it == typeIndices.end()) {
+        throw ParseException("unknown dollared function type", s.line, s.col);
+      }
+      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);
+      }
+      return types[offset];
+    }
+  }
+  // It's a list.
+  if (*s[0] == FUNC) {
+    std::vector<Type> params;
+    std::vector<Type> results;
+    for (size_t k = 1; k < s.size(); k++) {
+      Element& curr = *s[k];
+      if (elementStartsWith(curr, PARAM)) {
+        auto newParams = parseParamOrLocal(curr);
+        params.insert(params.end(), newParams.begin(), newParams.end());
+      } else if (elementStartsWith(curr, RESULT)) {
+        auto newResults = parseResults(curr);
+        results.insert(results.end(), newResults.begin(), newResults.end());
+      }
     }
+    return Signature(Type(params), Type(results));
   }
-  return Signature(Type(params), Type(results));
+  throw ParseException("invalid heap type", s.line, s.col);
 }
 
 void SExpressionWasmBuilder::parseType(Element& s) {
   size_t i = 1;
   if (s[i]->isStr()) {
     std::string name = s[i]->str().str;
-    if (signatureIndices.find(name) != signatureIndices.end()) {
+    if (typeIndices.find(name) != typeIndices.end()) {
       throw ParseException("duplicate function type", s.line, s.col);
     }
-    signatureIndices[name] = signatures.size();
+    typeIndices[name] = types.size();
     i++;
   }
-  signatures.emplace_back(parseInlineFunctionSignature(*s[i]));
+  types.emplace_back(parseHeapType(*s[i]));
 }
 
 void SExpressionWasmBuilder::parseEvent(Element& s, bool preParseImport) {
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 4c506beda..b323dcc06 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -471,13 +471,19 @@ Type Type::reinterpret() const {
 
 FeatureSet Type::getFeatures() const {
   auto getSingleFeatures = [](Type t) -> FeatureSet {
-    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;
+    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;
+      }
+      auto heapType = t.getHeapType();
+      if (heapType.isStruct() || heapType.isArray()) {
+        return FeatureSet::ReferenceTypes | FeatureSet::GC;
+      }
     }
     TODO_SINGLE_COMPOUND(t);
     switch (t.getBasic()) {