[WebAssembly] Refactor WasmSignature and use it for MCSymbolWasm

MCContext does not destroy MCSymbols on shutdown. So, rather than putting
SmallVectors (which may heap-allocate) inside MCSymbolWasm, use unowned pointer
to a WasmSignature instead. The signatures are now owned by the AsmPrinter.
Also uses WasmSignature instead of param and result vectors in TargetStreamer,
and leaves some TODOs for further simplification.

 Differential Revision: https://reviews.llvm.org/D52580

19 files changed, 198 insertions, 175 deletions

diff --git a/llvm/include/llvm/BinaryFormat/Wasm.h b/llvm/include/llvm/BinaryFormat/Wasm.h
index aca56af4514..44dd92ea901 100644
--- a/llvm/include/llvm/BinaryFormat/Wasm.h
+++ b/llvm/include/llvm/BinaryFormat/Wasm.h
@@ -16,6 +16,7 @@
 #define LLVM_BINARYFORMAT_WASM_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
 
 namespace llvm {
 namespace wasm {
@@ -34,11 +35,6 @@ struct WasmObjectHeader {
   uint32_t Version;
 };
 
-struct WasmSignature {
-  std::vector<uint8_t> ParamTypes;
-  uint8_t ReturnType;
-};
-
 struct WasmExport {
   StringRef Name;
   uint8_t Kind;
@@ -220,16 +216,6 @@ enum : unsigned {
   WASM_LIMITS_FLAG_HAS_MAX = 0x1,
 };
 
-// Subset of types that a value can have
-enum class ValType {
-  I32 = WASM_TYPE_I32,
-  I64 = WASM_TYPE_I64,
-  F32 = WASM_TYPE_F32,
-  F64 = WASM_TYPE_F64,
-  V128 = WASM_TYPE_V128,
-  EXCEPT_REF = WASM_TYPE_EXCEPT_REF,
-};
-
 // Kind codes used in the custom "name" section
 enum : unsigned {
   WASM_NAMES_FUNCTION = 0x1,
@@ -276,9 +262,32 @@ enum : unsigned {
 
 #undef WASM_RELOC
 
+// Subset of types that a value can have
+enum class ValType {
+  I32 = WASM_TYPE_I32,
+  I64 = WASM_TYPE_I64,
+  F32 = WASM_TYPE_F32,
+  F64 = WASM_TYPE_F64,
+  V128 = WASM_TYPE_V128,
+  EXCEPT_REF = WASM_TYPE_EXCEPT_REF,
+};
+
+struct WasmSignature {
+  SmallVector<wasm::ValType, 1> Returns;
+  SmallVector<wasm::ValType, 4> Params;
+  // Support empty and tombstone instances, needed by DenseMap.
+  enum { Plain, Empty, Tombstone } State = Plain;
+
+  WasmSignature(SmallVector<wasm::ValType, 1> &&InReturns,
+                SmallVector<wasm::ValType, 4> &&InParams)
+      : Returns(InReturns), Params(InParams) {}
+  WasmSignature() = default;
+};
+
 // Useful comparison operators
 inline bool operator==(const WasmSignature &LHS, const WasmSignature &RHS) {
-  return LHS.ReturnType == RHS.ReturnType && LHS.ParamTypes == RHS.ParamTypes;
+  return LHS.State == RHS.State && LHS.Returns == RHS.Returns &&
+         LHS.Params == RHS.Params;
 }
 
 inline bool operator!=(const WasmSignature &LHS, const WasmSignature &RHS) {
diff --git a/llvm/include/llvm/MC/MCSymbolWasm.h b/llvm/include/llvm/MC/MCSymbolWasm.h
index df311e03137..281178aea61 100644
--- a/llvm/include/llvm/MC/MCSymbolWasm.h
+++ b/llvm/include/llvm/MC/MCSymbolWasm.h
@@ -20,11 +20,8 @@ class MCSymbolWasm : public MCSymbol {
   bool IsHidden = false;
   bool IsComdat = false;
   std::string ModuleName;
-  SmallVector<wasm::ValType, 1> Returns;
-  SmallVector<wasm::ValType, 4> Params;
+  wasm::WasmSignature *Signature = nullptr;
   wasm::WasmGlobalType GlobalType;
-  bool ParamsSet = false;
-  bool ReturnsSet = false;
   bool GlobalTypeSet = false;
 
   /// An expression describing how to calculate the size of a symbol. If a
@@ -60,25 +57,8 @@ public:
   const StringRef getModuleName() const { return ModuleName; }
   void setModuleName(StringRef Name) { ModuleName = Name; }
 
-  const SmallVector<wasm::ValType, 1> &getReturns() const {
-    assert(ReturnsSet);
-    return Returns;
-  }
-
-  void setReturns(SmallVectorImpl<wasm::ValType> &&Rets) {
-    ReturnsSet = true;
-    Returns = std::move(Rets);
-  }
-
-  const SmallVector<wasm::ValType, 4> &getParams() const {
-    assert(ParamsSet);
-    return Params;
-  }
-
-  void setParams(SmallVectorImpl<wasm::ValType> &&Pars) {
-    ParamsSet = true;
-    Params = std::move(Pars);
-  }
+  const wasm::WasmSignature *getSignature() const { return Signature; }
+  void setSignature(wasm::WasmSignature *Sig) { Signature = Sig; }
 
   const wasm::WasmGlobalType &getGlobalType() const {
     assert(GlobalTypeSet);
diff --git a/llvm/include/llvm/Object/Wasm.h b/llvm/include/llvm/Object/Wasm.h
index 85cd40aee5b..25903fc9945 100644
--- a/llvm/include/llvm/Object/Wasm.h
+++ b/llvm/include/llvm/Object/Wasm.h
@@ -22,6 +22,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/Wasm.h"
 #include "llvm/Config/llvm-config.h"
+#include "llvm/MC/MCSymbolWasm.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Error.h"
diff --git a/llvm/include/llvm/Object/WasmTraits.h b/llvm/include/llvm/Object/WasmTraits.h
index ebcd00b1522..049d72f79e4 100644
--- a/llvm/include/llvm/Object/WasmTraits.h
+++ b/llvm/include/llvm/Object/WasmTraits.h
@@ -24,14 +24,20 @@ template <typename T> struct DenseMapInfo;
 // Traits for using WasmSignature in a DenseMap.
 template <> struct DenseMapInfo<wasm::WasmSignature> {
   static wasm::WasmSignature getEmptyKey() {
-    return wasm::WasmSignature{{}, 1};
+    wasm::WasmSignature Sig;
+    Sig.State = wasm::WasmSignature::Empty;
+    return Sig;
   }
   static wasm::WasmSignature getTombstoneKey() {
-    return wasm::WasmSignature{{}, 2};
+    wasm::WasmSignature Sig;
+    Sig.State = wasm::WasmSignature::Tombstone;
+    return Sig;
   }
   static unsigned getHashValue(const wasm::WasmSignature &Sig) {
-    unsigned H = hash_value(Sig.ReturnType);
-    for (int32_t Param : Sig.ParamTypes)
+    uintptr_t H = hash_value(Sig.State);
+    for (auto Ret : Sig.Returns)
+      H = hash_combine(H, Ret);
+    for (auto Param : Sig.Params)
       H = hash_combine(H, Param);
     return H;
   }
diff --git a/llvm/lib/MC/WasmObjectWriter.cpp b/llvm/lib/MC/WasmObjectWriter.cpp
index 2c28f01958f..cbbe161ae82 100644
--- a/llvm/lib/MC/WasmObjectWriter.cpp
+++ b/llvm/lib/MC/WasmObjectWriter.cpp
@@ -58,6 +58,7 @@ struct SectionBookkeeping {
 
 // The signature of a wasm function, in a struct capable of being used as a
 // DenseMap key.
+// TODO: Consider using WasmSignature directly instead.
 struct WasmFunctionType {
   // Support empty and tombstone instances, needed by DenseMap.
   enum { Plain, Empty, Tombstone } State;
@@ -1049,8 +1050,10 @@ uint32_t WasmObjectWriter::registerFunctionType(const MCSymbolWasm &Symbol) {
 
   WasmFunctionType F;
   const MCSymbolWasm *ResolvedSym = ResolveSymbol(Symbol);
-  F.Returns = ResolvedSym->getReturns();
-  F.Params = ResolvedSym->getParams();
+  if (auto *Sig = ResolvedSym->getSignature()) {
+    F.Returns = Sig->Returns;
+    F.Params = Sig->Params;
+  }
 
   auto Pair =
       FunctionTypeIndices.insert(std::make_pair(F, FunctionTypes.size()));
diff --git a/llvm/lib/Object/WasmObjectFile.cpp b/llvm/lib/Object/WasmObjectFile.cpp
index 711b415552a..75925a5ea10 100644
--- a/llvm/lib/Object/WasmObjectFile.cpp
+++ b/llvm/lib/Object/WasmObjectFile.cpp
@@ -701,17 +701,16 @@ Error WasmObjectFile::parseTypeSection(ReadContext &Ctx) {
   Signatures.reserve(Count);
   while (Count--) {
     wasm::WasmSignature Sig;
-    Sig.ReturnType = wasm::WASM_TYPE_NORESULT;
     uint8_t Form = readUint8(Ctx);
     if (Form != wasm::WASM_TYPE_FUNC) {
       return make_error<GenericBinaryError>("Invalid signature type",
                                             object_error::parse_failed);
     }
     uint32_t ParamCount = readVaruint32(Ctx);
-    Sig.ParamTypes.reserve(ParamCount);
+    Sig.Params.reserve(ParamCount);
     while (ParamCount--) {
       uint32_t ParamType = readUint8(Ctx);
-      Sig.ParamTypes.push_back(ParamType);
+      Sig.Params.push_back(wasm::ValType(ParamType));
     }
     uint32_t ReturnCount = readVaruint32(Ctx);
     if (ReturnCount) {
@@ -719,9 +718,9 @@ Error WasmObjectFile::parseTypeSection(ReadContext &Ctx) {
         return make_error<GenericBinaryError>(
             "Multiple return types not supported", object_error::parse_failed);
       }
-      Sig.ReturnType = readUint8(Ctx);
+      Sig.Returns.push_back(wasm::ValType(readUint8(Ctx)));
     }
-    Signatures.push_back(Sig);
+    Signatures.push_back(std::move(Sig));
   }
   if (Ctx.Ptr != Ctx.End)
     return make_error<GenericBinaryError>("Type section ended prematurely",
diff --git a/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.cpp b/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.cpp
index 045d4fb4fd6..b7b8c4fe039 100644
--- a/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.cpp
+++ b/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.cpp
@@ -235,26 +235,19 @@ void WebAssemblyInstPrinter::printWebAssemblySignatureOperand(const MCInst *MI,
   }
 }
 
-const char *llvm::WebAssembly::TypeToString(MVT Ty) {
-  switch (Ty.SimpleTy) {
-  case MVT::i32:
+const char *llvm::WebAssembly::TypeToString(wasm::ValType Ty) {
+  switch (Ty) {
+  case wasm::ValType::I32:
     return "i32";
-  case MVT::i64:
+  case wasm::ValType::I64:
     return "i64";
-  case MVT::f32:
+  case wasm::ValType::F32:
     return "f32";
-  case MVT::f64:
+  case wasm::ValType::F64:
     return "f64";
-  case MVT::v16i8:
-  case MVT::v8i16:
-  case MVT::v4i32:
-  case MVT::v2i64:
-  case MVT::v4f32:
-  case MVT::v2f64:
+  case wasm::ValType::V128:
     return "v128";
-  case MVT::ExceptRef:
+  case wasm::ValType::EXCEPT_REF:
     return "except_ref";
-  default:
-    llvm_unreachable("unsupported type");
   }
 }
diff --git a/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.h b/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.h
index 2bb6dcea9e0..18023328b38 100644
--- a/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.h
+++ b/llvm/lib/Target/WebAssembly/InstPrinter/WebAssemblyInstPrinter.h
@@ -50,7 +50,7 @@ public:
 
 namespace WebAssembly {
 
-const char *TypeToString(MVT Ty);
+const char *TypeToString(wasm::ValType Ty);
 
 } // end namespace WebAssembly
 
diff --git a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp
index 961ef8146a9..2158ee3be04 100644
--- a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp
+++ b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp
@@ -46,7 +46,7 @@ static void PrintTypes(formatted_raw_ostream &OS, ArrayRef<MVT> Types) {
       First = false;
     else
       OS << ", ";
-    OS << WebAssembly::TypeToString(Type);
+    OS << WebAssembly::TypeToString(WebAssembly::toValType(Type));
   }
   OS << '\n';
 }
@@ -85,16 +85,16 @@ void WebAssemblyTargetAsmStreamer::emitLocal(ArrayRef<MVT> Types) {
 void WebAssemblyTargetAsmStreamer::emitEndFunc() { OS << "\t.endfunc\n"; }
 
 void WebAssemblyTargetAsmStreamer::emitIndirectFunctionType(
-    MCSymbol *Symbol, SmallVectorImpl<MVT> &Params,
-    SmallVectorImpl<MVT> &Results) {
+    MCSymbolWasm *Symbol) {
   OS << "\t.functype\t" << Symbol->getName();
-  if (Results.empty())
+  if (Symbol->getSignature()->Returns.empty())
     OS << ", void";
   else {
-    assert(Results.size() == 1);
-    OS << ", " << WebAssembly::TypeToString(Results.front());
+    assert(Symbol->getSignature()->Returns.size() == 1);
+    OS << ", "
+       << WebAssembly::TypeToString(Symbol->getSignature()->Returns.front());
   }
-  for (auto Ty : Params)
+  for (auto Ty : Symbol->getSignature()->Params)
     OS << ", " << WebAssembly::TypeToString(Ty);
   OS << '\n';
 }
@@ -114,20 +114,12 @@ void WebAssemblyTargetAsmStreamer::emitIndIdx(const MCExpr *Value) {
 
 void WebAssemblyTargetWasmStreamer::emitParam(MCSymbol *Symbol,
                                               ArrayRef<MVT> Types) {
-  SmallVector<wasm::ValType, 4> Params;
-  for (MVT Ty : Types)
-    Params.push_back(WebAssembly::toValType(Ty));
-
-  cast<MCSymbolWasm>(Symbol)->setParams(std::move(Params));
+  // The Symbol already has its signature
 }
 
 void WebAssemblyTargetWasmStreamer::emitResult(MCSymbol *Symbol,
                                                ArrayRef<MVT> Types) {
-  SmallVector<wasm::ValType, 4> Returns;
-  for (MVT Ty : Types)
-    Returns.push_back(WebAssembly::toValType(Ty));
-
-  cast<MCSymbolWasm>(Symbol)->setReturns(std::move(Returns));
+  // The Symbol already has its signature
 }
 
 void WebAssemblyTargetWasmStreamer::emitLocal(ArrayRef<MVT> Types) {
@@ -155,25 +147,9 @@ void WebAssemblyTargetWasmStreamer::emitIndIdx(const MCExpr *Value) {
 }
 
 void WebAssemblyTargetWasmStreamer::emitIndirectFunctionType(
-    MCSymbol *Symbol, SmallVectorImpl<MVT> &Params,
-    SmallVectorImpl<MVT> &Results) {
-  MCSymbolWasm *WasmSym = cast<MCSymbolWasm>(Symbol);
-  if (WasmSym->isFunction()) {
-    // Symbol already has its arguments and result set.
-    return;
-  }
-
-  SmallVector<wasm::ValType, 4> ValParams;
-  for (MVT Ty : Params)
-    ValParams.push_back(WebAssembly::toValType(Ty));
-
-  SmallVector<wasm::ValType, 1> ValResults;
-  for (MVT Ty : Results)
-    ValResults.push_back(WebAssembly::toValType(Ty));
-
-  WasmSym->setParams(std::move(ValParams));
-  WasmSym->setReturns(std::move(ValResults));
-  WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
+    MCSymbolWasm *Symbol) {
+  // Symbol already has its arguments and result set.
+  Symbol->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
 }
 
 void WebAssemblyTargetWasmStreamer::emitGlobalImport(StringRef name) {
diff --git a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h
index 124161c5579..43c422d593a 100644
--- a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h
+++ b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h
@@ -40,9 +40,7 @@ public:
   /// .endfunc
   virtual void emitEndFunc() = 0;
   /// .functype
-  virtual void emitIndirectFunctionType(MCSymbol *Symbol,
-                                        SmallVectorImpl<MVT> &Params,
-                                        SmallVectorImpl<MVT> &Results) = 0;
+  virtual void emitIndirectFunctionType(MCSymbolWasm *Symbol) = 0;
   /// .indidx
   virtual void emitIndIdx(const MCExpr *Value) = 0;
   /// .import_global
@@ -65,8 +63,7 @@ public:
   void emitResult(MCSymbol *Symbol, ArrayRef<MVT> Types) override;
   void emitLocal(ArrayRef<MVT> Types) override;
   void emitEndFunc() override;
-  void emitIndirectFunctionType(MCSymbol *Symbol, SmallVectorImpl<MVT> &Params,
-                                SmallVectorImpl<MVT> &Results) override;
+  void emitIndirectFunctionType(MCSymbolWasm *Symbol) override;
   void emitIndIdx(const MCExpr *Value) override;
   void emitGlobalImport(StringRef name) override;
   void emitImportModule(MCSymbolWasm *Sym, StringRef ModuleName) override;
@@ -81,8 +78,7 @@ public:
   void emitResult(MCSymbol *Symbol, ArrayRef<MVT> Types) override;
   void emitLocal(ArrayRef<MVT> Types) override;
   void emitEndFunc() override;
-  void emitIndirectFunctionType(MCSymbol *Symbol, SmallVectorImpl<MVT> &Params,
-                                SmallVectorImpl<MVT> &Results) override;
+  void emitIndirectFunctionType(MCSymbolWasm *Symbol) override;
   void emitIndIdx(const MCExpr *Value) override;
   void emitGlobalImport(StringRef name) override;
   void emitImportModule(MCSymbolWasm *Sym, StringRef ModuleName) override;
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp
index 20b41e1ffab..2ea3760b923 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp
@@ -83,16 +83,25 @@ void WebAssemblyAsmPrinter::EmitEndOfAsmFile(Module &M) {
     if (F.isDeclarationForLinker() && !F.isIntrinsic()) {
       SmallVector<MVT, 4> Results;
       SmallVector<MVT, 4> Params;
-      ComputeSignatureVTs(F, TM, Params, Results);
-      MCSymbol *Sym = getSymbol(&F);
-      getTargetStreamer()->emitIndirectFunctionType(Sym, Params, Results);
+      ComputeSignatureVTs(F.getFunctionType(), F, TM, Params, Results);
+      auto *Sym = cast<MCSymbolWasm>(getSymbol(&F));
+      if (!Sym->getSignature()) {
+        auto Signature = SignatureFromMVTs(Results, Params);
+        Sym->setSignature(Signature.get());
+        addSignature(std::move(Signature));
+      }
+      // FIXME: this was originally intended for post-linking and was only used
+      // for imports that were only called indirectly (i.e. s2wasm could not
+      // infer the type from a call). With object files it applies to all
+      // imports. so fix the names and the tests, or rethink how import
+      // delcarations work in asm files.
+      getTargetStreamer()->emitIndirectFunctionType(Sym);
 
       if (TM.getTargetTriple().isOSBinFormatWasm() &&
           F.hasFnAttribute("wasm-import-module")) {
-        MCSymbolWasm *WasmSym = cast<MCSymbolWasm>(Sym);
         StringRef Name =
             F.getFnAttribute("wasm-import-module").getValueAsString();
-        getTargetStreamer()->emitImportModule(WasmSym, Name);
+        getTargetStreamer()->emitImportModule(Sym, Name);
       }
     }
   }
@@ -137,10 +146,17 @@ void WebAssemblyAsmPrinter::EmitJumpTableInfo() {
 }
 
 void WebAssemblyAsmPrinter::EmitFunctionBodyStart() {
-  getTargetStreamer()->emitParam(CurrentFnSym, MFI->getParams());
-
-  SmallVector<MVT, 4> ResultVTs;
   const Function &F = MF->getFunction();
+  SmallVector<MVT, 1> ResultVTs;
+  SmallVector<MVT, 4> ParamVTs;
+  ComputeSignatureVTs(F.getFunctionType(), F, TM, ParamVTs, ResultVTs);
+  auto Signature = SignatureFromMVTs(ResultVTs, ParamVTs);
+  auto *WasmSym = cast<MCSymbolWasm>(CurrentFnSym);
+  WasmSym->setSignature(Signature.get());
+  addSignature(std::move(Signature));
+
+  // FIXME: clean up how params and results are emitted (use signatures)
+  getTargetStreamer()->emitParam(CurrentFnSym, ParamVTs);
 
   // Emit the function index.
   if (MDNode *Idx = F.getMetadata("wasm.index")) {
@@ -150,15 +166,7 @@ void WebAssemblyAsmPrinter::EmitFunctionBodyStart() {
         cast<ConstantAsMetadata>(Idx->getOperand(0))->getValue()));
   }
 
-  ComputeLegalValueVTs(F, TM, F.getReturnType(), ResultVTs);
-
-  // If the return type needs to be legalized it will get converted into
-  // passing a pointer.
-  if (ResultVTs.size() == 1)
-    getTargetStreamer()->emitResult(CurrentFnSym, ResultVTs);
-  else
-    getTargetStreamer()->emitResult(CurrentFnSym, ArrayRef<MVT>());
-
+  getTargetStreamer()->emitResult(CurrentFnSym, ResultVTs);
   getTargetStreamer()->emitLocal(MFI->getLocals());
 
   AsmPrinter::EmitFunctionBodyStart();
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.h b/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.h
index 97b1ae578b5..f6cb5610bad 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.h
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyAsmPrinter.h
@@ -25,6 +25,8 @@ class LLVM_LIBRARY_VISIBILITY WebAssemblyAsmPrinter final : public AsmPrinter {
   const WebAssemblySubtarget *Subtarget;
   const MachineRegisterInfo *MRI;
   WebAssemblyFunctionInfo *MFI;
+  // TODO: Do the uniquing of Signatures here instead of ObjectFileWriter?
+  std::vector<std::unique_ptr<wasm::WasmSignature>> Signatures;
 
 public:
   explicit WebAssemblyAsmPrinter(TargetMachine &TM,
@@ -37,6 +39,9 @@ public:
   }
 
   const WebAssemblySubtarget &getSubtarget() const { return *Subtarget; }
+  void addSignature(std::unique_ptr<wasm::WasmSignature> &&Sig) {
+    Signatures.push_back(std::move(Sig));
+  }
 
   //===------------------------------------------------------------------===//
   // MachineFunctionPass Implementation.
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
index 55c963b7b40..a8992d89ca3 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
@@ -769,12 +769,18 @@ SDValue WebAssemblyTargetLowering::LowerFormalArguments(
     MFI->addParam(PtrVT);
   }
 
-  // Record the number and types of results.
+  // Record the number and types of arguments and results.
   SmallVector<MVT, 4> Params;
   SmallVector<MVT, 4> Results;
-  ComputeSignatureVTs(MF.getFunction(), DAG.getTarget(), Params, Results);
+  ComputeSignatureVTs(MF.getFunction().getFunctionType(), MF.getFunction(),
+                      DAG.getTarget(), Params, Results);
   for (MVT VT : Results)
     MFI->addResult(VT);
+  // TODO: Use signatures in WebAssemblyMachineFunctionInfo too and unify
+  // the param logic here with ComputeSignatureVTs
+  assert(MFI->getParams().size() == Params.size() &&
+         std::equal(MFI->getParams().begin(), MFI->getParams().end(),
+                    Params.begin()));
 
   return Chain;
 }
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp
index a52bedd9f29..e9a0cf51905 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp
@@ -51,34 +51,13 @@ WebAssemblyMCInstLower::GetGlobalAddressSymbol(const MachineOperand &MO) const {
     const TargetMachine &TM = MF.getTarget();
     const Function &CurrentFunc = MF.getFunction();
 
-    SmallVector<wasm::ValType, 4> Returns;
-    SmallVector<wasm::ValType, 4> Params;
+    SmallVector<MVT, 1> ResultMVTs;
+    SmallVector<MVT, 4> ParamMVTs;
+    ComputeSignatureVTs(FuncTy, CurrentFunc, TM, ParamMVTs, ResultMVTs);
 
-    wasm::ValType iPTR = MF.getSubtarget<WebAssemblySubtarget>().hasAddr64()
-                             ? wasm::ValType::I64
-                             : wasm::ValType::I32;
-
-    SmallVector<MVT, 4> ResultMVTs;
-    ComputeLegalValueVTs(CurrentFunc, TM, FuncTy->getReturnType(), ResultMVTs);
-    // WebAssembly can't currently handle returning tuples.
-    if (ResultMVTs.size() <= 1)
-      for (MVT ResultMVT : ResultMVTs)
-        Returns.push_back(WebAssembly::toValType(ResultMVT));
-    else
-      Params.push_back(iPTR);
-
-    for (Type *Ty : FuncTy->params()) {
-      SmallVector<MVT, 4> ParamMVTs;
-      ComputeLegalValueVTs(CurrentFunc, TM, Ty, ParamMVTs);
-      for (MVT ParamMVT : ParamMVTs)
-        Params.push_back(WebAssembly::toValType(ParamMVT));
-    }
-
-    if (FuncTy->isVarArg())
-      Params.push_back(iPTR);
-
-    WasmSym->setReturns(std::move(Returns));
-    WasmSym->setParams(std::move(Params));
+    auto Signature = SignatureFromMVTs(ResultMVTs, ParamMVTs);
+    WasmSym->setSignature(Signature.get());
+    Printer.addSignature(std::move(Signature));
     WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
   }
 
@@ -108,9 +87,10 @@ MCSymbol *WebAssemblyMCInstLower::GetExternalSymbolSymbol(
   SmallVector<wasm::ValType, 4> Returns;
   SmallVector<wasm::ValType, 4> Params;
   GetLibcallSignature(Subtarget, Name, Returns, Params);
-
-  WasmSym->setReturns(std::move(Returns));
-  WasmSym->setParams(std::move(Params));
+  auto Signature =
+      make_unique<wasm::WasmSignature>(std::move(Returns), std::move(Params));
+  WasmSym->setSignature(Signature.get());
+  Printer.addSignature(std::move(Signature));
   WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
 
   return WasmSym;
@@ -203,8 +183,10 @@ void WebAssemblyMCInstLower::Lower(const MachineInstr *MI,
             Params.pop_back();
 
           MCSymbolWasm *WasmSym = cast<MCSymbolWasm>(Sym);
-          WasmSym->setReturns(std::move(Returns));
-          WasmSym->setParams(std::move(Params));
+          auto Signature = make_unique<wasm::WasmSignature>(std::move(Returns),
+                                                            std::move(Params));
+          WasmSym->setSignature(Signature.get());
+          Printer.addSignature(std::move(Signature));
           WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
 
           const MCExpr *Expr = MCSymbolRefExpr::create(
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.cpp
index e511e574050..073706e567e 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.cpp
@@ -43,20 +43,34 @@ void llvm::ComputeLegalValueVTs(const Function &F, const TargetMachine &TM,
   }
 }
 
-void llvm::ComputeSignatureVTs(const Function &F, const TargetMachine &TM,
+void llvm::ComputeSignatureVTs(const FunctionType *Ty, const Function &F,
+                               const TargetMachine &TM,
                                SmallVectorImpl<MVT> &Params,
                                SmallVectorImpl<MVT> &Results) {
-  ComputeLegalValueVTs(F, TM, F.getReturnType(), Results);
+  ComputeLegalValueVTs(F, TM, Ty->getReturnType(), Results);
 
+  MVT PtrVT = MVT::getIntegerVT(TM.createDataLayout().getPointerSizeInBits());
   if (Results.size() > 1) {
     // WebAssembly currently can't lower returns of multiple values without
     // demoting to sret (see WebAssemblyTargetLowering::CanLowerReturn). So
     // replace multiple return values with a pointer parameter.
     Results.clear();
-    Params.push_back(
-        MVT::getIntegerVT(TM.createDataLayout().getPointerSizeInBits()));
+    Params.push_back(PtrVT);
   }
 
-  for (auto &Arg : F.args())
-    ComputeLegalValueVTs(F, TM, Arg.getType(), Params);
+  for (auto *Param : Ty->params())
+    ComputeLegalValueVTs(F, TM, Param, Params);
+  if (Ty->isVarArg())
+    Params.push_back(PtrVT);
+}
+
+std::unique_ptr<wasm::WasmSignature>
+llvm::SignatureFromMVTs(const SmallVectorImpl<MVT> &Results,
+                        const SmallVectorImpl<MVT> &Params) {
+  auto Sig = make_unique<wasm::WasmSignature>();
+  for (MVT Ty : Results)
+    Sig->Returns.push_back(WebAssembly::toValType(Ty));
+  for (MVT Ty : Params)
+    Sig->Params.push_back(WebAssembly::toValType(Ty));
+  return Sig;
 }
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h b/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h
index 5a2602f87ee..cde44d24599 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h
@@ -17,7 +17,9 @@
 #define LLVM_LIB_TARGET_WEBASSEMBLY_WEBASSEMBLYMACHINEFUNCTIONINFO_H
 
 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
+#include "llvm/BinaryFormat/Wasm.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/MC/MCSymbolWasm.h"
 
 namespace llvm {
 
@@ -121,10 +123,16 @@ public:
 void ComputeLegalValueVTs(const Function &F, const TargetMachine &TM, Type *Ty,
                           SmallVectorImpl<MVT> &ValueVTs);
 
-void ComputeSignatureVTs(const Function &F, const TargetMachine &TM,
-                         SmallVectorImpl<MVT> &Params,
+// Compute the signature for a given FunctionType (Ty). Note that it's not the
+// signature for F (F is just used to get varous context)
+void ComputeSignatureVTs(const FunctionType *Ty, const Function &F,
+                         const TargetMachine &TM, SmallVectorImpl<MVT> &Params,
                          SmallVectorImpl<MVT> &Results);
 
+std::unique_ptr<wasm::WasmSignature>
+SignatureFromMVTs(const SmallVectorImpl<MVT> &Results,
+                  const SmallVectorImpl<MVT> &Params);
+
 } // end namespace llvm
 
 #endif
diff --git a/llvm/test/MC/WebAssembly/external-func-address.ll b/llvm/test/MC/WebAssembly/external-func-address.ll
index 4e39feb6341..60ec23a1a8e 100644
--- a/llvm/test/MC/WebAssembly/external-func-address.ll
+++ b/llvm/test/MC/WebAssembly/external-func-address.ll
@@ -3,7 +3,7 @@
 target triple = "wasm32-unknown-unknown"
 
 ; Verify that addresses of external functions generate correctly typed
-; imports and relocations or type R_TABLE_INDEX_I32.
+; imports and relocations of type R_TABLE_INDEX_I32.
 
 declare void @f0(i32) #0
 @ptr_to_f0 = hidden global void (i32)* @f0, align 4
@@ -13,6 +13,13 @@ attributes #0 = { "wasm-import-module"="somewhere" }
 declare void @f1(i32) #1
 @ptr_to_f1 = hidden global void (i32)* @f1, align 4
 
+; Check that varargs functions have correctly typed imports
+declare i32 @varargs(i32, i32, ...)
+define void @call(i32) {
+ %a = call i32 (i32, i32, ...) @varargs(i32 0, i32 0)
+ ret void
+}
+
 ; CHECK:      --- !WASM
 ; CHECK-NEXT: FileHeader:      
 ; CHECK-NEXT:   Version:         0x00000001
@@ -21,7 +28,13 @@ declare void @f1(i32) #1
 ; CHECK-NEXT:     Signatures:      
 ; CHECK-NEXT:       - Index:           0
 ; CHECK-NEXT:         ReturnType:      NORESULT
-; CHECK-NEXT:         ParamTypes:      
+; CHECK-NEXT:         ParamTypes:
+; CHECK-NEXT:           - I32
+; CHECK-NEXT:       - Index:           1
+; CHECK-NEXT:         ReturnType:      I32
+; CHECK-NEXT:         ParamTypes:
+; CHECK-NEXT:           - I32
+; CHECK-NEXT:           - I32
 ; CHECK-NEXT:           - I32
 ; CHECK:        - Type:            IMPORT
 ; CHECK-NEXT:     Imports:
@@ -29,6 +42,10 @@ declare void @f1(i32) #1
 ; CHECK-NEXT:         Field:           __linear_memory
 ; CHECK:            - Module:          env
 ; CHECK-NEXT:         Field:           __indirect_function_table
+; CHECK:            - Module:          env
+; CHECK-NEXT:         Field:           varargs
+; CHECK-NEXT:         Kind:            FUNCTION
+; CHECK-NEXT:         SigIndex:        1
 ; CHECK:            - Module:          somewhere
 ; CHECK-NEXT:         Field:           f0
 ; CHECK:            - Module:          env
@@ -40,9 +57,9 @@ declare void @f1(i32) #1
 ; CHECK-NEXT:       - Offset:
 ; CHECK-NEXT:           Opcode:          I32_CONST
 ; CHECK-NEXT:           Value:           1
-; CHECK-NEXT:         Functions:       [ 0, 1 ]
+; CHECK-NEXT:         Functions:       [ 1, 2 ]
 ; CHECK:        - Type:            DATA
 ; CHECK-NEXT:     Relocations:
 ; CHECK-NEXT:       - Type:            R_WEBASSEMBLY_TABLE_INDEX_I32
-; CHECK-NEXT:         Index:           1
+; CHECK-NEXT:         Index:           3
 ; CHECK-NEXT:         Offset:          0x00000006
diff --git a/llvm/test/MC/WebAssembly/types.ll b/llvm/test/MC/WebAssembly/types.ll
index 1f85f85695a..b8012420a8f 100644
--- a/llvm/test/MC/WebAssembly/types.ll
+++ b/llvm/test/MC/WebAssembly/types.ll
@@ -12,6 +12,7 @@ declare <4 x i32> @v4i32()
 declare <2 x i64> @v2i64()
 declare <4 x float> @v4f32()
 declare <2 x double> @v2f64()
+declare { i32, i32, i32 } @structret()
 
 define void @f1() {
 entry:
@@ -25,6 +26,11 @@ entry:
   %tmp8 = call <2 x i64> @v2i64()
   %tmp9 = call <4 x float> @v4f32()
   %tmp10 = call <2 x double> @v2f64()
+  %tmp11 = call {i32, i32, i32} @structret()
+  ret void
+}
+
+define void @vararg(i32, i32, ...) {
   ret void
 }
 
@@ -48,5 +54,15 @@ entry:
 ; CHECK-NEXT:       - Index: 5
 ; CHECK-NEXT:         ReturnType: V128
 ; CHECK-NEXT:         ParamTypes:
+; CHECK-NEXT:       - Index: 6
+; CHECK-NEXT:         ReturnType: NORESULT
+; CHECK-NEXT:         ParamTypes:
+; CHECK-NEXT:           - I32
+; CHECK-NEXT:       - Index: 7
+; CHECK-NEXT:         ReturnType: NORESULT
+; CHECK-NEXT:         ParamTypes:
+; CHECK-NEXT:           - I32
+; CHECK-NEXT:           - I32
+; CHECK-NEXT:           - I32
 ; should be no additional types
 ; CHECK-NOT: ReturnType
diff --git a/llvm/tools/obj2yaml/wasm2yaml.cpp b/llvm/tools/obj2yaml/wasm2yaml.cpp
index ad91acaeb26..a1120a96571 100644
--- a/llvm/tools/obj2yaml/wasm2yaml.cpp
+++ b/llvm/tools/obj2yaml/wasm2yaml.cpp
@@ -155,9 +155,13 @@ ErrorOr<WasmYAML::Object *> WasmDumper::dump() {
       for (const auto &FunctionSig : Obj.types()) {
         WasmYAML::Signature Sig;
         Sig.Index = Index++;
-        Sig.ReturnType = FunctionSig.ReturnType;
-        for (const auto &ParamType : FunctionSig.ParamTypes)
-          Sig.ParamTypes.push_back(ParamType);
+        Sig.ReturnType = wasm::WASM_TYPE_NORESULT;
+        assert(FunctionSig.Returns.size() <= 1 &&
+               "Functions with multiple returns are not supported");
+        if (FunctionSig.Returns.size())
+          Sig.ReturnType = static_cast<uint32_t>(FunctionSig.Returns[0]);
+        for (const auto &ParamType : FunctionSig.Params)
+          Sig.ParamTypes.push_back(static_cast<uint32_t>(ParamType));
         TypeSec->Signatures.push_back(Sig);
       }
       S = std::move(TypeSec);