[WebAssembly] Support CFI for WebAssembly target

Summary: This patch implements CFI for WebAssembly. It modifies the
LowerTypeTest pass to pre-assign table indexes to functions that are
called indirectly, and lowers type checks to test against the
appropriate table indexes. It also modifies the WebAssembly backend to
support a special ".indidx" assembly directive that propagates the table
index assignments out to the linker.

Patch by Dominic Chen

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@277398 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 96 insertions, 28 deletions

diff --git a/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCTargetDesc.h b/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCTargetDesc.h
index 001bd7f1fc4..7a32ddd5793 100644
--- a/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCTargetDesc.h
+++ b/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCTargetDesc.h
@@ -59,6 +59,7 @@ enum Directive {
   DotResult = UINT64_MAX - 1,  ///< .result
   DotLocal = UINT64_MAX - 2,   ///< .local
   DotEndFunc = UINT64_MAX - 3, ///< .endfunc
+  DotIndIdx = UINT64_MAX - 4,  /// < .indidx
 };
 
 } // end namespace WebAssembly
@@ -123,7 +124,8 @@ inline unsigned GetDefaultP2Align(unsigned Opcode) {
   case WebAssembly::STORE_I64:
   case WebAssembly::STORE_F64:
     return 3;
-  default: llvm_unreachable("Only loads and stores have p2align values");
+  default:
+    llvm_unreachable("Only loads and stores have p2align values");
   }
 }
 
diff --git a/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp b/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp
index 3d61c15717b..6ca95aaa03b 100644
--- a/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp
+++ b/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.cpp
@@ -67,13 +67,18 @@ void WebAssemblyTargetAsmStreamer::emitEndFunc() { OS << "\t.endfunc\n"; }
 void WebAssemblyTargetAsmStreamer::emitIndirectFunctionType(
     StringRef name, SmallVectorImpl<MVT> &SignatureVTs, size_t NumResults) {
   OS << "\t.functype\t" << name;
-  if (NumResults == 0) OS << ", void";
+  if (NumResults == 0)
+    OS << ", void";
   for (auto Ty : SignatureVTs) {
     OS << ", " << WebAssembly::TypeToString(Ty);
   }
   OS << "\n";
 }
 
+void WebAssemblyTargetAsmStreamer::emitIndIdx(const MCExpr *Value) {
+  OS << "\t.indidx  \t" << *Value << '\n';
+}
+
 // FIXME: What follows is not the real binary encoding.
 
 static void EncodeTypes(MCStreamer &Streamer, ArrayRef<MVT> Types) {
@@ -100,3 +105,8 @@ void WebAssemblyTargetELFStreamer::emitLocal(ArrayRef<MVT> Types) {
 void WebAssemblyTargetELFStreamer::emitEndFunc() {
   Streamer.EmitIntValue(WebAssembly::DotEndFunc, sizeof(uint64_t));
 }
+
+void WebAssemblyTargetELFStreamer::emitIndIdx(const MCExpr *Value) {
+  Streamer.EmitIntValue(WebAssembly::DotIndIdx, sizeof(uint64_t));
+  Streamer.EmitValue(Value, sizeof(uint64_t));
+}
diff --git a/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h b/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h
index 51354ef22d7..2e3e9bccdb1 100644
--- a/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h
+++ b/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyTargetStreamer.h
@@ -43,6 +43,8 @@ public:
                                         size_t NumResults) {
     llvm_unreachable("emitIndirectFunctionType not implemented");
   }
+  /// .indidx
+  virtual void emitIndIdx(const MCExpr *Value) = 0;
 };
 
 /// This part is for ascii assembly output
@@ -59,6 +61,7 @@ public:
   void emitIndirectFunctionType(StringRef name,
                                 SmallVectorImpl<MVT> &SignatureVTs,
                                 size_t NumResults) override;
+  void emitIndIdx(const MCExpr *Value) override;
 };
 
 /// This part is for ELF object output
@@ -70,6 +73,7 @@ public:
   void emitResult(ArrayRef<MVT> Types) override;
   void emitLocal(ArrayRef<MVT> Types) override;
   void emitEndFunc() override;
+  void emitIndIdx(const MCExpr *Value) override;
 };
 
 } // end namespace llvm
diff --git a/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp b/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp
index 54e9f7f5290..b95d77b770a 100644
--- a/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp
+++ b/lib/Target/WebAssembly/WebAssemblyAsmPrinter.cpp
@@ -14,10 +14,10 @@
 ///
 //===----------------------------------------------------------------------===//
 
-#include "WebAssembly.h"
 #include "InstPrinter/WebAssemblyInstPrinter.h"
 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
 #include "MCTargetDesc/WebAssemblyTargetStreamer.h"
+#include "WebAssembly.h"
 #include "WebAssemblyMCInstLower.h"
 #include "WebAssemblyMachineFunctionInfo.h"
 #include "WebAssemblyRegisterInfo.h"
@@ -183,6 +183,15 @@ void WebAssemblyAsmPrinter::EmitFunctionBodyStart() {
 
   SmallVector<MVT, 4> ResultVTs;
   const Function &F(*MF->getFunction());
+
+  // Emit the function index.
+  if (MDNode *Idx = F.getMetadata("wasm.index")) {
+    assert(Idx->getNumOperands() == 1);
+
+    getTargetStreamer()->emitIndIdx(AsmPrinter::lowerConstant(
+        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
diff --git a/lib/Transforms/IPO/LowerTypeTests.cpp b/lib/Transforms/IPO/LowerTypeTests.cpp
index 36089f0a880..8e343bcbb6e 100644
--- a/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -13,7 +13,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
-#include "llvm/Transforms/IPO.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Triple.h"
@@ -30,6 +29,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace llvm;
@@ -79,8 +79,7 @@ bool BitSetInfo::containsValue(
     if (!Result)
       return false;
     COffset += APOffset.getZExtValue();
-    return containsValue(DL, GlobalLayout, GEP->getPointerOperand(),
-                         COffset);
+    return containsValue(DL, GlobalLayout, GEP->getPointerOperand(), COffset);
   }
 
   if (auto Op = dyn_cast<Operator>(V)) {
@@ -222,6 +221,9 @@ struct LowerTypeTests : public ModulePass {
   IntegerType *Int64Ty;
   IntegerType *IntPtrTy;
 
+  // Indirect function call index assignment counter for WebAssembly
+  uint64_t IndirectIndex;
+
   // Mapping from type identifiers to the call sites that test them.
   DenseMap<Metadata *, std::vector<CallInst *>> TypeTestCallSites;
 
@@ -250,6 +252,10 @@ struct LowerTypeTests : public ModulePass {
   void verifyTypeMDNode(GlobalObject *GO, MDNode *Type);
   void buildBitSetsFromFunctions(ArrayRef<Metadata *> TypeIds,
                                  ArrayRef<Function *> Functions);
+  void buildBitSetsFromFunctionsX86(ArrayRef<Metadata *> TypeIds,
+                                    ArrayRef<Function *> Functions);
+  void buildBitSetsFromFunctionsWASM(ArrayRef<Metadata *> TypeIds,
+                                     ArrayRef<Function *> Functions);
   void buildBitSetsFromDisjointSet(ArrayRef<Metadata *> TypeIds,
                                    ArrayRef<GlobalObject *> Globals);
   bool lower();
@@ -267,8 +273,7 @@ ModulePass *llvm::createLowerTypeTestsPass() { return new LowerTypeTests; }
 /// Build a bit set for TypeId using the object layouts in
 /// GlobalLayout.
 BitSetInfo LowerTypeTests::buildBitSet(
-    Metadata *TypeId,
-    const DenseMap<GlobalObject *, uint64_t> &GlobalLayout) {
+    Metadata *TypeId, const DenseMap<GlobalObject *, uint64_t> &GlobalLayout) {
   BitSetBuilder BSB;
 
   // Compute the byte offset of each address associated with this type
@@ -281,8 +286,9 @@ BitSetInfo LowerTypeTests::buildBitSet(
       if (Type->getOperand(1) != TypeId)
         continue;
       uint64_t Offset =
-          cast<ConstantInt>(cast<ConstantAsMetadata>(Type->getOperand(0))
-                                ->getValue())->getZExtValue();
+          cast<ConstantInt>(
+              cast<ConstantAsMetadata>(Type->getOperand(0))->getValue())
+              ->getZExtValue();
       BSB.addOffset(GlobalAndOffset.second + Offset);
     }
   }
@@ -311,8 +317,8 @@ ByteArrayInfo *LowerTypeTests::createByteArray(BitSetInfo &BSI) {
   // we know the offset and mask to use.
   auto ByteArrayGlobal = new GlobalVariable(
       *M, Int8Ty, /*isConstant=*/true, GlobalValue::PrivateLinkage, nullptr);
-  auto MaskGlobal = new GlobalVariable(
-      *M, Int8Ty, /*isConstant=*/true, GlobalValue::PrivateLinkage, nullptr);
+  auto MaskGlobal = new GlobalVariable(*M, Int8Ty, /*isConstant=*/true,
+                                       GlobalValue::PrivateLinkage, nullptr);
 
   ByteArrayInfos.emplace_back();
   ByteArrayInfo *BAI = &ByteArrayInfos.back();
@@ -588,8 +594,7 @@ void LowerTypeTests::lowerTypeTestCalls(
 
 void LowerTypeTests::verifyTypeMDNode(GlobalObject *GO, MDNode *Type) {
   if (Type->getNumOperands() != 2)
-    report_fatal_error(
-        "All operands of type metadata must have 2 elements");
+    report_fatal_error("All operands of type metadata must have 2 elements");
 
   if (GO->isThreadLocal())
     report_fatal_error("Bit set element may not be thread-local");
@@ -612,9 +617,6 @@ void LowerTypeTests::verifyTypeMDNode(GlobalObject *GO, MDNode *Type) {
 static const unsigned kX86JumpTableEntrySize = 8;
 
 unsigned LowerTypeTests::getJumpTableEntrySize() {
-  if (Arch != Triple::x86 && Arch != Triple::x86_64)
-    report_fatal_error("Unsupported architecture for jump tables");
-
   return kX86JumpTableEntrySize;
 }
 
@@ -625,9 +627,6 @@ unsigned LowerTypeTests::getJumpTableEntrySize() {
 Constant *LowerTypeTests::createJumpTableEntry(GlobalObject *Src,
                                                Function *Dest,
                                                unsigned Distance) {
-  if (Arch != Triple::x86 && Arch != Triple::x86_64)
-    report_fatal_error("Unsupported architecture for jump tables");
-
   const unsigned kJmpPCRel32Code = 0xe9;
   const unsigned kInt3Code = 0xcc;
 
@@ -652,18 +651,27 @@ Constant *LowerTypeTests::createJumpTableEntry(GlobalObject *Src,
 }
 
 Type *LowerTypeTests::getJumpTableEntryType() {
-  if (Arch != Triple::x86 && Arch != Triple::x86_64)
-    report_fatal_error("Unsupported architecture for jump tables");
-
   return StructType::get(M->getContext(),
                          {Int8Ty, Int32Ty, Int8Ty, Int8Ty, Int8Ty},
                          /*Packed=*/true);
 }
 
-/// Given a disjoint set of type identifiers and functions, build a jump table
-/// for the functions, build the bit sets and lower the llvm.type.test calls.
+/// Given a disjoint set of type identifiers and functions, build the bit sets
+/// and lower the llvm.type.test calls, architecture dependently.
 void LowerTypeTests::buildBitSetsFromFunctions(ArrayRef<Metadata *> TypeIds,
                                                ArrayRef<Function *> Functions) {
+  if (Arch == Triple::x86 || Arch == Triple::x86_64)
+    buildBitSetsFromFunctionsX86(TypeIds, Functions);
+  else if (Arch == Triple::wasm32 || Arch == Triple::wasm64)
+    buildBitSetsFromFunctionsWASM(TypeIds, Functions);
+  else
+    report_fatal_error("Unsupported architecture for jump tables");
+}
+
+/// Given a disjoint set of type identifiers and functions, build a jump table
+/// for the functions, build the bit sets and lower the llvm.type.test calls.
+void LowerTypeTests::buildBitSetsFromFunctionsX86(
+    ArrayRef<Metadata *> TypeIds, ArrayRef<Function *> Functions) {
   // Unlike the global bitset builder, the function bitset builder cannot
   // re-arrange functions in a particular order and base its calculations on the
   // layout of the functions' entry points, as we have no idea how large a
@@ -795,6 +803,40 @@ void LowerTypeTests::buildBitSetsFromFunctions(ArrayRef<Metadata *> TypeIds,
       ConstantArray::get(JumpTableType, JumpTableEntries));
 }
 
+/// Assign a dummy layout using an incrementing counter, tag each function
+/// with its index represented as metadata, and lower each type test to an
+/// integer range comparison. During generation of the indirect function call
+/// table in the backend, it will assign the given indexes.
+/// Note: Dynamic linking is not supported, as the WebAssembly ABI has not yet
+/// been finalized.
+void LowerTypeTests::buildBitSetsFromFunctionsWASM(
+    ArrayRef<Metadata *> TypeIds, ArrayRef<Function *> Functions) {
+  assert(!Functions.empty());
+
+  // Build consecutive monotonic integer ranges for each call target set
+  DenseMap<GlobalObject *, uint64_t> GlobalLayout;
+  for (Function *F : Functions) {
+    // Skip functions that are not address taken, to avoid bloating the table
+    if (!F->hasAddressTaken())
+      continue;
+
+    // Store metadata with the index for each function
+    MDNode *MD = MDNode::get(F->getContext(),
+                             ArrayRef<Metadata *>(ConstantAsMetadata::get(
+                                 ConstantInt::get(Int64Ty, IndirectIndex))));
+    F->setMetadata("wasm.index", MD);
+
+    // Assign the counter value
+    GlobalLayout[F] = IndirectIndex++;
+  }
+
+  // The indirect function table index space starts at zero, so pass a NULL
+  // pointer as the subtracted "jump table" offset.
+  lowerTypeTestCalls(TypeIds,
+                     ConstantPointerNull::get(cast<PointerType>(Int32PtrTy)),
+                     GlobalLayout);
+}
+
 void LowerTypeTests::buildBitSetsFromDisjointSet(
     ArrayRef<Metadata *> TypeIds, ArrayRef<GlobalObject *> Globals) {
   llvm::DenseMap<Metadata *, uint64_t> TypeIdIndices;
@@ -900,8 +942,7 @@ bool LowerTypeTests::lower() {
 
     auto BitSetMDVal = dyn_cast<MetadataAsValue>(CI->getArgOperand(1));
     if (!BitSetMDVal)
-      report_fatal_error(
-          "Second argument of llvm.type.test must be metadata");
+      report_fatal_error("Second argument of llvm.type.test must be metadata");
     auto BitSet = BitSetMDVal->getMetadata();
 
     // Add the call site to the list of call sites for this type identifier. We
@@ -939,7 +980,8 @@ bool LowerTypeTests::lower() {
   for (GlobalClassesTy::iterator I = GlobalClasses.begin(),
                                  E = GlobalClasses.end();
        I != E; ++I) {
-    if (!I->isLeader()) continue;
+    if (!I->isLeader())
+      continue;
     ++NumTypeIdDisjointSets;
 
     unsigned MaxIndex = 0;
@@ -1000,6 +1042,7 @@ static void init(LowerTypeTests *LTT, Module &M) {
   LTT->Int64Ty = Type::getInt64Ty(M.getContext());
   LTT->IntPtrTy = DL.getIntPtrType(M.getContext(), 0);
   LTT->TypeTestCallSites.clear();
+  LTT->IndirectIndex = 0;
 }
 
 bool LowerTypeTests::runOnModule(Module &M) {