Preserve function heap types during text parsing (#3935)

Previously, ref.func instructions would be assigned the canonical (i.e. first
parsed) heap type for the referenced function signature rather than the HeapType
actually specified in the type definition. In nominal mode, this could cause
validation failures because the types assigned to ref.func instructions would
not be correct.

Fix the problem by tracking function HeapTypes rather than function Signatures
when parsing the text format.

There can still be validation failures when round-tripping modules because
function HeapTypes are not properly preserved after parsing, but that will be
addressed in a follow-up PR.

2 files changed, 42 insertions, 34 deletions

diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 4145d781a..e5ba2996e 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -135,7 +135,7 @@ class SExpressionWasmBuilder {
   int elemCounter = 0;
   int memoryCounter = 0;
   // we need to know function return types before we parse their contents
-  std::map<Name, Signature> functionSignatures;
+  std::map<Name, HeapType> functionTypes;
   std::unordered_map<cashew::IString, Index> debugInfoFileIndices;
 
   // Maps type indexes to a mapping of field index => name. This is not the same
@@ -297,13 +297,12 @@ private:
   std::vector<Type> parseParamOrLocal(Element& s);
   std::vector<NameType> parseParamOrLocal(Element& s, size_t& localIndex);
   std::vector<Type> parseResults(Element& s);
-  Signature parseTypeRef(Element& s);
+  HeapType parseTypeRef(Element& s);
   size_t parseTypeUse(Element& s,
                       size_t startPos,
-                      Signature& functionSignature,
+                      HeapType& functionType,
                       std::vector<NameType>& namedParams);
-  size_t
-  parseTypeUse(Element& s, size_t startPos, Signature& functionSignature);
+  size_t parseTypeUse(Element& s, size_t startPos, HeapType& functionType);
 
   void stringToBinary(const char* input, size_t size, std::vector<char>& data);
   void parseMemory(Element& s, bool preParseImport = false);
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 8bedeab9f..bda94dec5 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -580,7 +580,7 @@ std::vector<Type> SExpressionWasmBuilder::parseResults(Element& s) {
 // Parses an element that references an entry in the type section. The element
 // should be in the form of (type name) or (type index).
 // (e.g. (type $a), (type 0))
-Signature SExpressionWasmBuilder::parseTypeRef(Element& s) {
+HeapType SExpressionWasmBuilder::parseTypeRef(Element& s) {
   assert(elementStartsWith(s, TYPE));
   if (s.size() != 2) {
     throw ParseException("invalid type reference", s.line, s.col);
@@ -589,7 +589,7 @@ Signature SExpressionWasmBuilder::parseTypeRef(Element& s) {
   if (!heapType.isSignature()) {
     throw ParseException("expected signature type", s.line, s.col);
   }
-  return heapType.getSignature();
+  return heapType;
 }
 
 // Prases typeuse, a reference to a type definition. It is in the form of either
@@ -602,7 +602,7 @@ Signature SExpressionWasmBuilder::parseTypeRef(Element& s) {
 size_t
 SExpressionWasmBuilder::parseTypeUse(Element& s,
                                      size_t startPos,
-                                     Signature& functionSignature,
+                                     HeapType& functionType,
                                      std::vector<NameType>& namedParams) {
   std::vector<Type> params, results;
   size_t i = startPos;
@@ -610,7 +610,7 @@ SExpressionWasmBuilder::parseTypeUse(Element& s,
   bool typeExists = false, paramsOrResultsExist = false;
   if (i < s.size() && elementStartsWith(*s[i], TYPE)) {
     typeExists = true;
-    functionSignature = parseTypeRef(*s[i++]);
+    functionType = parseTypeRef(*s[i++]);
   }
 
   size_t paramPos = i;
@@ -640,10 +640,10 @@ SExpressionWasmBuilder::parseTypeUse(Element& s,
   }
 
   if (!typeExists) {
-    functionSignature = inlineSig;
+    functionType = inlineSig;
   } else if (paramsOrResultsExist) {
     // verify that (type) and (params)/(result) match
-    if (inlineSig != functionSignature) {
+    if (inlineSig != functionType.getSignature()) {
       throw ParseException("type and param/result don't match",
                            s[paramPos]->line,
                            s[paramPos]->col);
@@ -651,15 +651,16 @@ SExpressionWasmBuilder::parseTypeUse(Element& s,
   }
 
   // Add implicitly defined type to global list so it has an index
-  auto heapType = HeapType(functionSignature);
-  if (std::find(types.begin(), types.end(), heapType) == types.end()) {
-    types.push_back(heapType);
+  if (std::find(types.begin(), types.end(), functionType) == types.end()) {
+    types.push_back(functionType);
   }
 
   // If only (type) is specified, populate `namedParams`
   if (!paramsOrResultsExist) {
     size_t index = 0;
-    for (const auto& param : functionSignature.params) {
+    assert(functionType.isSignature());
+    Signature sig = functionType.getSignature();
+    for (const auto& param : sig.params) {
       namedParams.emplace_back(Name::fromInt(index++), param);
     }
   }
@@ -670,9 +671,9 @@ SExpressionWasmBuilder::parseTypeUse(Element& s,
 // Parses a typeuse. Use this when only FunctionType* is needed.
 size_t SExpressionWasmBuilder::parseTypeUse(Element& s,
                                             size_t startPos,
-                                            Signature& functionSignature) {
+                                            HeapType& functionType) {
   std::vector<NameType> params;
-  return parseTypeUse(s, startPos, functionSignature, params);
+  return parseTypeUse(s, startPos, functionType, params);
 }
 
 void SExpressionWasmBuilder::preParseHeapTypes(Element& module) {
@@ -913,9 +914,9 @@ void SExpressionWasmBuilder::preParseFunctionType(Element& s) {
   }
   functionNames.push_back(name);
   functionCounter++;
-  Signature sig;
-  parseTypeUse(s, i, sig);
-  functionSignatures[name] = sig;
+  HeapType type;
+  parseTypeUse(s, i, type);
+  functionTypes[name] = type;
 }
 
 size_t SExpressionWasmBuilder::parseFunctionNames(Element& s,
@@ -989,9 +990,9 @@ void SExpressionWasmBuilder::parseFunction(Element& s, bool preParseImport) {
   }
 
   // parse typeuse: type/param/result
-  Signature sig;
+  HeapType type;
   std::vector<NameType> params;
-  i = parseTypeUse(s, i, sig, params);
+  i = parseTypeUse(s, i, type, params);
 
   // when (import) is inside a (func) element, this is not a function definition
   // but an import.
@@ -1006,8 +1007,8 @@ void SExpressionWasmBuilder::parseFunction(Element& s, bool preParseImport) {
     im->setName(name, hasExplicitName);
     im->module = importModule;
     im->base = importBase;
-    im->sig = sig;
-    functionSignatures[name] = sig;
+    im->sig = type.getSignature();
+    functionTypes[name] = type;
     if (wasm.getFunctionOrNull(im->name)) {
       throw ParseException("duplicate import", s.line, s.col);
     }
@@ -1034,7 +1035,7 @@ void SExpressionWasmBuilder::parseFunction(Element& s, bool preParseImport) {
 
   // make a new function
   currFunction = std::unique_ptr<Function>(Builder(wasm).makeFunction(
-    name, std::move(params), sig.results, std::move(vars)));
+    name, std::move(params), type.getSignature().results, std::move(vars)));
   currFunction->profile = profile;
 
   // parse body
@@ -1061,7 +1062,7 @@ void SExpressionWasmBuilder::parseFunction(Element& s, bool preParseImport) {
     autoBlock->name = FAKE_RETURN;
   }
   if (autoBlock) {
-    autoBlock->finalize(sig.results);
+    autoBlock->finalize(type.getSignature().results);
   }
   if (!currFunction->body) {
     currFunction->body = allocator.alloc<Nop>();
@@ -2258,7 +2259,7 @@ Expression* SExpressionWasmBuilder::makeCall(Element& s, bool isReturn) {
   auto target = getFunctionName(*s[1]);
   auto ret = allocator.alloc<Call>();
   ret->target = target;
-  ret->type = functionSignatures[ret->target].results;
+  ret->type = functionTypes[ret->target].getSignature().results;
   parseCallOperands(s, 2, s.size(), ret);
   ret->isReturn = isReturn;
   ret->finalize();
@@ -2277,7 +2278,9 @@ Expression* SExpressionWasmBuilder::makeCallIndirect(Element& s,
   } else {
     ret->table = wasm.tables.front()->name;
   }
-  i = parseTypeUse(s, i, ret->sig);
+  HeapType callType;
+  i = parseTypeUse(s, i, callType);
+  ret->sig = callType.getSignature();
   parseCallOperands(s, i, s.size() - 1, ret);
   ret->target = parseExpression(s[s.size() - 1]);
   ret->isReturn = isReturn;
@@ -2393,7 +2396,7 @@ Expression* SExpressionWasmBuilder::makeRefFunc(Element& s) {
   ret->func = func;
   // To support typed function refs, we give the reference not just a general
   // funcref, but a specific subtype with the actual signature.
-  ret->finalize(Type(HeapType(functionSignatures[func]), NonNullable));
+  ret->finalize(Type(functionTypes[func], NonNullable));
   return ret;
 }
 
@@ -3042,11 +3045,13 @@ void SExpressionWasmBuilder::parseImport(Element& s) {
   if (kind == ExternalKind::Function) {
     auto func = make_unique<Function>();
 
-    j = parseTypeUse(inner, j, func->sig);
+    HeapType funcType;
+    j = parseTypeUse(inner, j, funcType);
+    func->sig = funcType.getSignature();
     func->setName(name, hasExplicitName);
     func->module = module;
     func->base = base;
-    functionSignatures[name] = func->sig;
+    functionTypes[name] = funcType;
     wasm.addFunction(func.release());
   } else if (kind == ExternalKind::Global) {
     Type type;
@@ -3109,7 +3114,9 @@ void SExpressionWasmBuilder::parseImport(Element& s) {
     }
   } else if (kind == ExternalKind::Tag) {
     auto tag = make_unique<Tag>();
-    j = parseTypeUse(inner, j, tag->sig);
+    HeapType tagType;
+    j = parseTypeUse(inner, j, tagType);
+    tag->sig = tagType.getSignature();
     tag->setName(name, hasExplicitName);
     tag->module = module;
     tag->base = base;
@@ -3396,7 +3403,7 @@ ElementSegment* SExpressionWasmBuilder::parseElemFinish(
     for (; i < s.size(); i++) {
       auto func = getFunctionName(*s[i]);
       segment->data.push_back(
-        Builder(wasm).makeRefFunc(func, functionSignatures[func]));
+        Builder(wasm).makeRefFunc(func, functionTypes[func].getSignature()));
     }
   }
   return wasm.addElementSegment(std::move(segment));
@@ -3491,7 +3498,9 @@ void SExpressionWasmBuilder::parseTag(Element& s, bool preParseImport) {
   }
 
   // Parse typeuse
-  i = parseTypeUse(s, i, tag->sig);
+  HeapType tagType;
+  i = parseTypeUse(s, i, tagType);
+  tag->sig = tagType.getSignature();
 
   // If there are more elements, they are invalid
   if (i < s.size()) {