3 files changed, 172 insertions, 77 deletions

diff --git a/include/llvm/CodeGen/StackMaps.h b/include/llvm/CodeGen/StackMaps.h
index f832132..86d55cf 100644
--- a/include/llvm/CodeGen/StackMaps.h
+++ b/include/llvm/CodeGen/StackMaps.h
@@ -20,6 +20,60 @@ namespace llvm {
 class AsmPrinter;
 class MCExpr;
 
+/// \brief MI-level patchpoint operands.
+///
+/// MI patchpoint operations take the form:
+/// [<def>], <id>, <numBytes>, <target>, <numArgs>, <cc>, ...
+///
+/// Note that IR/SD patchpoints do not have the <def> or <cc> operands.
+///
+/// Patchpoints following the anyregcc convention are handled specially. For
+/// these, the stack map also records the location of the return value and
+/// arguments.
+class PatchPointOpers {
+public:
+  /// Enumerate the meta operands.
+  enum { IDPos, NBytesPos, TargetPos, NArgPos, CCPos, MetaEnd };
+private:
+  const MachineInstr *MI;
+  bool HasDef;
+  bool IsAnyReg;
+public:
+  explicit PatchPointOpers(const MachineInstr *MI);
+
+  bool isAnyReg() const { return IsAnyReg; }
+  bool hasDef() const { return HasDef; }
+
+  unsigned getMetaIdx(unsigned Pos = 0) const {
+    assert(Pos < MetaEnd && "Meta operand index out of range.");
+    return (HasDef ? 1 : 0) + Pos;
+  }
+
+  const MachineOperand &getMetaOper(unsigned Pos) {
+    return MI->getOperand(getMetaIdx(Pos));
+  }
+
+  unsigned getArgIdx() const { return getMetaIdx() + MetaEnd; }
+
+  /// Get the operand index of the variable list of non-argument operands.
+  /// These hold the "live state".
+  unsigned getVarIdx() const {
+    return getMetaIdx() + MetaEnd
+      + MI->getOperand(getMetaIdx(NArgPos)).getImm();
+  }
+
+  /// Get the index at which stack map locations will be recorded.
+  /// Arguments are not recorded unless the anyregcc convention is used.
+  unsigned getStackMapStartIdx() const {
+    if (IsAnyReg)
+      return getArgIdx();
+    return getVarIdx();
+  }
+
+  /// \brief Get the next scratch register operand index.
+  unsigned getNextScratchIdx(unsigned StartIdx = 0) const;
+};
+
 class StackMaps {
 public:
   struct Location {
@@ -48,15 +102,13 @@ public:
   StackMaps(AsmPrinter &AP, OperandParser OpParser)
     : AP(AP), OpParser(OpParser) {}
 
-  /// This should be called by the MC lowering code _immediately_ before
-  /// lowering the MI to an MCInst. It records where the operands for the
-  /// instruction are stored, and outputs a label to record the offset of
-  /// the call from the start of the text section. In special cases (e.g. AnyReg
-  /// calling convention) the return register is also recorded if requested.
-  void recordStackMap(const MachineInstr &MI, uint32_t ID,
-                      MachineInstr::const_mop_iterator MOI,
-                      MachineInstr::const_mop_iterator MOE,
-                      bool recordResult = false);
+  /// \brief Generate a stackmap record for a stackmap instruction.
+  ///
+  /// MI must be a raw STACKMAP, not a PATCHPOINT.
+  void recordStackMap(const MachineInstr &MI);
+
+  /// \brief Generate a stackmap record for a patchpoint instruction.
+  void recordPatchPoint(const MachineInstr &MI);
 
   /// If there is any stack map data, create a stack map section and serialize
   /// the map info into it. This clears the stack map data structures
@@ -64,7 +116,6 @@ public:
   void serializeToStackMapSection();
 
 private:
-
   typedef SmallVector<Location, 8> LocationVec;
 
   struct CallsiteInfo {
@@ -103,6 +154,16 @@ private:
   OperandParser OpParser;
   CallsiteInfoList CSInfos;
   ConstantPool ConstPool;
+
+  /// This should be called by the MC lowering code _immediately_ before
+  /// lowering the MI to an MCInst. It records where the operands for the
+  /// instruction are stored, and outputs a label to record the offset of
+  /// the call from the start of the text section. In special cases (e.g. AnyReg
+  /// calling convention) the return register is also recorded if requested.
+  void recordStackMapOpers(const MachineInstr &MI, uint32_t ID,
+                           MachineInstr::const_mop_iterator MOI,
+                           MachineInstr::const_mop_iterator MOE,
+                           bool recordResult = false);
 };
 
 }
diff --git a/lib/CodeGen/StackMaps.cpp b/lib/CodeGen/StackMaps.cpp
index 0eeec83..86e51a1 100644
--- a/lib/CodeGen/StackMaps.cpp
+++ b/lib/CodeGen/StackMaps.cpp
@@ -28,10 +28,47 @@
 
 using namespace llvm;
 
-void StackMaps::recordStackMap(const MachineInstr &MI, uint32_t ID,
-                               MachineInstr::const_mop_iterator MOI,
-                               MachineInstr::const_mop_iterator MOE,
-                               bool recordResult) {
+PatchPointOpers::PatchPointOpers(const MachineInstr *MI):
+  MI(MI),
+  HasDef(MI->getOperand(0).isReg() && MI->getOperand(0).isDef() &&
+         !MI->getOperand(0).isImplicit()),
+  IsAnyReg(MI->getOperand(getMetaIdx(CCPos)).getImm() == CallingConv::AnyReg) {
+
+#ifndef NDEBUG
+  {
+  unsigned CheckStartIdx = 0, e = MI->getNumOperands();
+  while (CheckStartIdx < e && MI->getOperand(CheckStartIdx).isReg() &&
+         MI->getOperand(CheckStartIdx).isDef() &&
+         !MI->getOperand(CheckStartIdx).isImplicit())
+    ++CheckStartIdx;
+
+  assert(getMetaIdx() == CheckStartIdx &&
+         "Unexpected additonal definition in Patchpoint intrinsic.");
+  }
+#endif
+}
+
+unsigned PatchPointOpers::getNextScratchIdx(unsigned StartIdx) const {
+  if (!StartIdx)
+    StartIdx = getVarIdx();
+
+  // Find the next scratch register (implicit def and early clobber)
+  unsigned ScratchIdx = StartIdx, e = MI->getNumOperands();
+  while (ScratchIdx < e &&
+         !(MI->getOperand(ScratchIdx).isReg() &&
+           MI->getOperand(ScratchIdx).isDef() &&
+           MI->getOperand(ScratchIdx).isImplicit() &&
+           MI->getOperand(ScratchIdx).isEarlyClobber()))
+    ++ScratchIdx;
+
+  assert(ScratchIdx != e && "No scratch register available");
+  return ScratchIdx;
+}
+
+void StackMaps::recordStackMapOpers(const MachineInstr &MI, uint32_t ID,
+                                    MachineInstr::const_mop_iterator MOI,
+                                    MachineInstr::const_mop_iterator MOE,
+                                    bool recordResult) {
 
   MCContext &OutContext = AP.OutStreamer.getContext();
   MCSymbol *MILabel = OutContext.CreateTempSymbol();
@@ -73,6 +110,49 @@ void StackMaps::recordStackMap(const MachineInstr &MI, uint32_t ID,
   CSInfos.push_back(CallsiteInfo(CSOffsetExpr, ID, CallsiteLocs));
 }
 
+static MachineInstr::const_mop_iterator
+getStackMapEndMOP(MachineInstr::const_mop_iterator MOI,
+                  MachineInstr::const_mop_iterator MOE) {
+  for (; MOI != MOE; ++MOI)
+    if (MOI->isRegMask() || (MOI->isReg() && MOI->isImplicit()))
+      break;
+
+  return MOI;
+}
+
+void StackMaps::recordStackMap(const MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::STACKMAP && "exected stackmap");
+
+  int64_t ID = MI.getOperand(0).getImm();
+  assert((int32_t)ID == ID && "Stack maps hold 32-bit IDs");
+  recordStackMapOpers(MI, ID, llvm::next(MI.operands_begin(), 2),
+                      getStackMapEndMOP(MI.operands_begin(),
+                                        MI.operands_end()));
+}
+
+void StackMaps::recordPatchPoint(const MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::PATCHPOINT && "exected stackmap");
+
+  PatchPointOpers opers(&MI);
+  int64_t ID = opers.getMetaOper(PatchPointOpers::IDPos).getImm();
+  assert((int32_t)ID == ID && "Stack maps hold 32-bit IDs");
+  MachineInstr::const_mop_iterator MOI =
+    llvm::next(MI.operands_begin(), opers.getStackMapStartIdx());
+  recordStackMapOpers(MI, ID, MOI, getStackMapEndMOP(MOI, MI.operands_end()),
+                      opers.isAnyReg() && opers.hasDef());
+
+#ifndef NDEBUG
+  // verify anyregcc
+  LocationVec &Locations = CSInfos.back().Locations;
+  if (opers.isAnyReg()) {
+    unsigned NArgs = opers.getMetaOper(PatchPointOpers::NArgPos).getImm();
+    for (unsigned i = 0, e = (opers.hasDef() ? NArgs+1 : NArgs); i != e; ++i)
+      assert(Locations[i].LocType == Location::Register &&
+             "anyreg arg must be in reg.");
+  }
+#endif
+}
+
 /// serializeToStackMapSection conceptually populates the following fields:
 ///
 /// uint32 : Reserved (header)
diff --git a/lib/Target/X86/X86MCInstLower.cpp b/lib/Target/X86/X86MCInstLower.cpp
index 4e27ef0..6649c82 100644
--- a/lib/Target/X86/X86MCInstLower.cpp
+++ b/lib/Target/X86/X86MCInstLower.cpp
@@ -716,6 +716,9 @@ X86AsmPrinter::stackmapOperandParser(MachineInstr::const_mop_iterator MOI,
          "Register mask and implicit operands should not be processed.");
 
   if (MOP.isImm()) {
+    // Verify anyregcc
+    // [<def>], <id>, <numBytes>, <target>, <numArgs>, <cc>, ...
+
     switch (MOP.getImm()) {
     default: llvm_unreachable("Unrecognized operand type.");
     case StackMaps::DirectMemRefOp: {
@@ -756,82 +759,33 @@ X86AsmPrinter::stackmapOperandParser(MachineInstr::const_mop_iterator MOI,
     Location(Location::Register, RC->getSize(), MOP.getReg(), 0), ++MOI);
 }
 
-static MachineInstr::const_mop_iterator
-getStackMapEndMOP(MachineInstr::const_mop_iterator MOI,
-                  MachineInstr::const_mop_iterator MOE) {
-  for (; MOI != MOE; ++MOI)
-    if (MOI->isRegMask() || (MOI->isReg() && MOI->isImplicit()))
-      break;
-
-  return MOI;
-}
-
+// Lower a stackmap of the form:
+// <id>, <shadowBytes>, ...
 static void LowerSTACKMAP(MCStreamer &OutStreamer,
                           StackMaps &SM,
                           const MachineInstr &MI)
 {
-  int64_t ID = MI.getOperand(0).getImm();
   unsigned NumNOPBytes = MI.getOperand(1).getImm();
-
-  assert((int32_t)ID == ID && "Stack maps hold 32-bit IDs");
-  SM.recordStackMap(MI, ID, llvm::next(MI.operands_begin(), 2),
-                    getStackMapEndMOP(MI.operands_begin(), MI.operands_end()));
+  SM.recordStackMap(MI);
   // Emit padding.
+  // FIXME: These nops ensure that the stackmap's shadow is covered by
+  // instructions from the same basic block, but the nops should not be
+  // necessary if instructions from the same block follow the stackmap.
   for (unsigned i = 0; i < NumNOPBytes; ++i)
     OutStreamer.EmitInstruction(MCInstBuilder(X86::NOOP));
 }
 
 // Lower a patchpoint of the form:
-// [<def>], <id>, <numBytes>, <target>, <numArgs>
+// [<def>], <id>, <numBytes>, <target>, <numArgs>, <cc>, ...
 static void LowerPATCHPOINT(MCStreamer &OutStreamer,
                             StackMaps &SM,
                             const MachineInstr &MI) {
-  bool hasDef = MI.getOperand(0).isReg() && MI.getOperand(0).isDef() &&
-                !MI.getOperand(0).isImplicit();
-  unsigned StartIdx = hasDef ? 1 : 0;
-#ifndef NDEBUG
-  {
-  unsigned StartIdx2 = 0, e = MI.getNumOperands();
-  while (StartIdx2 < e && MI.getOperand(StartIdx2).isReg() &&
-         MI.getOperand(StartIdx2).isDef() &&
-         !MI.getOperand(StartIdx2).isImplicit())
-    ++StartIdx2;
-
-  assert(StartIdx == StartIdx2 &&
-         "Unexpected additonal definition in Patchpoint intrinsic.");
-  }
-#endif
-
-  // Find the first scratch register (implicit def and early clobber)
-  unsigned ScratchIdx = StartIdx, e = MI.getNumOperands();
-  while (ScratchIdx < e &&
-         !(MI.getOperand(ScratchIdx).isReg() &&
-           MI.getOperand(ScratchIdx).isDef() &&
-           MI.getOperand(ScratchIdx).isImplicit() &&
-           MI.getOperand(ScratchIdx).isEarlyClobber()))
-    ++ScratchIdx;
-
-  assert(ScratchIdx != e && "No scratch register available");
-
-  int64_t ID = MI.getOperand(StartIdx).getImm();
-  assert((int32_t)ID == ID && "Stack maps hold 32-bit IDs");
-
-  // Get the number of arguments participating in the call. This number was
-  // adjusted during call lowering by subtracting stack args.
-  bool isAnyRegCC = MI.getOperand(StartIdx + 4).getImm() == CallingConv::AnyReg;
-  assert(((hasDef && isAnyRegCC) || !hasDef) &&
-         "Only Patchpoints with AnyReg calling convention may have a result");
-  int64_t StackMapIdx = isAnyRegCC ? StartIdx + 5 :
-    StartIdx + 5 + MI.getOperand(StartIdx + 3).getImm();
-  assert(StackMapIdx <= MI.getNumOperands() &&
-         "Patchpoint intrinsic dropped arguments.");
-
-  SM.recordStackMap(MI, ID, llvm::next(MI.operands_begin(), StackMapIdx),
-                    getStackMapEndMOP(MI.operands_begin(), MI.operands_end()),
-                    isAnyRegCC && hasDef);
+  SM.recordPatchPoint(MI);
 
+  PatchPointOpers opers(&MI);
+  unsigned ScratchIdx = opers.getNextScratchIdx();
   unsigned EncodedBytes = 0;
-  int64_t CallTarget = MI.getOperand(StartIdx + 2).getImm();
+  int64_t CallTarget = opers.getMetaOper(PatchPointOpers::TargetPos).getImm();
   if (CallTarget) {
     // Emit MOV to materialize the target address and the CALL to target.
     // This is encoded with 12-13 bytes, depending on which register is used.
@@ -845,11 +799,11 @@ static void LowerPATCHPOINT(MCStreamer &OutStreamer,
                                 .addReg(MI.getOperand(ScratchIdx).getReg()));
   }
   // Emit padding.
-  unsigned NumNOPBytes = MI.getOperand(StartIdx + 1).getImm();
-  assert(NumNOPBytes >= EncodedBytes &&
+  unsigned NumBytes = opers.getMetaOper(PatchPointOpers::NBytesPos).getImm();
+  assert(NumBytes >= EncodedBytes &&
          "Patchpoint can't request size less than the length of a call.");
 
-  for (unsigned i = EncodedBytes; i < NumNOPBytes; ++i)
+  for (unsigned i = EncodedBytes; i < NumBytes; ++i)
     OutStreamer.EmitInstruction(MCInstBuilder(X86::NOOP));
 }