2 files changed, 30 insertions, 18 deletions

diff --git a/llvm/include/llvm/CodeGen/StackMaps.h b/llvm/include/llvm/CodeGen/StackMaps.h
index 2e50689bb57..7b55b796863 100644
--- a/llvm/include/llvm/CodeGen/StackMaps.h
+++ b/llvm/include/llvm/CodeGen/StackMaps.h
@@ -219,7 +219,7 @@ public:
   void reset() {
     CSInfos.clear();
     ConstPool.clear();
-    FnStackSize.clear();
+    FnInfos.clear();
   }
 
   /// \brief Generate a stackmap record for a stackmap instruction.
@@ -243,7 +243,13 @@ private:
   typedef SmallVector<Location, 8> LocationVec;
   typedef SmallVector<LiveOutReg, 8> LiveOutVec;
   typedef MapVector<uint64_t, uint64_t> ConstantPool;
-  typedef MapVector<const MCSymbol *, uint64_t> FnStackSizeMap;
+
+  struct FunctionInfo {
+    uint64_t StackSize;
+    uint64_t RecordCount;
+    FunctionInfo() : StackSize(0), RecordCount(1) {}
+    explicit FunctionInfo(uint64_t StackSize) : StackSize(StackSize), RecordCount(1) {}
+  };
 
   struct CallsiteInfo {
     const MCExpr *CSOffsetExpr;
@@ -257,12 +263,13 @@ private:
           LiveOuts(std::move(LiveOuts)) {}
   };
 
+  typedef MapVector<const MCSymbol *, FunctionInfo> FnInfoMap;
   typedef std::vector<CallsiteInfo> CallsiteInfoList;
 
   AsmPrinter &AP;
   CallsiteInfoList CSInfos;
   ConstantPool ConstPool;
-  FnStackSizeMap FnStackSize;
+  FnInfoMap FnInfos;
 
   MachineInstr::const_mop_iterator
   parseOperand(MachineInstr::const_mop_iterator MOI,
diff --git a/llvm/include/llvm/Object/StackMapParser.h b/llvm/include/llvm/Object/StackMapParser.h
index e58162de150..efea62bb3cb 100644
--- a/llvm/include/llvm/Object/StackMapParser.h
+++ b/llvm/include/llvm/Object/StackMapParser.h
@@ -17,7 +17,7 @@
 namespace llvm {
 
 template <support::endianness Endianness>
-class StackMapV1Parser {
+class StackMapV2Parser {
 public:
 
   template <typename AccessorT>
@@ -47,7 +47,7 @@ public:
 
   /// Accessor for function records.
   class FunctionAccessor {
-    friend class StackMapV1Parser;
+    friend class StackMapV2Parser;
   public:
 
     /// Get the function address.
@@ -56,14 +56,19 @@ public:
     }
 
     /// Get the function's stack size.
-    uint32_t getStackSize() const {
+    uint64_t getStackSize() const {
       return read<uint64_t>(P + sizeof(uint64_t));
     }
+    
+    /// Get the number of callsite records.
+    uint64_t getRecordCount() const {
+      return read<uint64_t>(P + (2 * sizeof(uint64_t)));
+    }
 
   private:
     FunctionAccessor(const uint8_t *P) : P(P) {}
 
-    const static int FunctionAccessorSize = 2 * sizeof(uint64_t);
+    const static int FunctionAccessorSize = 3 * sizeof(uint64_t);
 
     FunctionAccessor next() const {
       return FunctionAccessor(P + FunctionAccessorSize);
@@ -74,7 +79,7 @@ public:
 
   /// Accessor for constants.
   class ConstantAccessor {
-    friend class StackMapV1Parser;
+    friend class StackMapV2Parser;
   public:
 
     /// Return the value of this constant.
@@ -103,7 +108,7 @@ public:
 
   /// Accessor for location records.
   class LocationAccessor {
-    friend class StackMapV1Parser;
+    friend class StackMapV2Parser;
     friend class RecordAccessor;
   public:
 
@@ -156,7 +161,7 @@ public:
 
   /// Accessor for stackmap live-out fields.
   class LiveOutAccessor {
-    friend class StackMapV1Parser;
+    friend class StackMapV2Parser;
     friend class RecordAccessor;
   public:
 
@@ -188,7 +193,7 @@ public:
 
   /// Accessor for stackmap records.
   class RecordAccessor {
-    friend class StackMapV1Parser;
+    friend class StackMapV2Parser;
   public:
 
     typedef AccessorIterator<LocationAccessor> location_iterator;
@@ -292,14 +297,14 @@ public:
     const uint8_t *P;
   };
 
-  /// Construct a parser for a version-1 stackmap. StackMap data will be read
+  /// Construct a parser for a version-2 stackmap. StackMap data will be read
   /// from the given array.
-  StackMapV1Parser(ArrayRef<uint8_t> StackMapSection)
+  StackMapV2Parser(ArrayRef<uint8_t> StackMapSection)
       : StackMapSection(StackMapSection) {
     ConstantsListOffset = FunctionListOffset + getNumFunctions() * FunctionSize;
 
-    assert(StackMapSection[0] == 1 &&
-           "StackMapV1Parser can only parse version 1 stackmaps");
+    assert(StackMapSection[0] == 2 &&
+           "StackMapV2Parser can only parse version 2 stackmaps");
 
     unsigned CurrentRecordOffset =
       ConstantsListOffset + getNumConstants() * ConstantSize;
@@ -315,8 +320,8 @@ public:
   typedef AccessorIterator<ConstantAccessor> constant_iterator;
   typedef AccessorIterator<RecordAccessor> record_iterator;
 
-  /// Get the version number of this stackmap. (Always returns 1).
-  unsigned getVersion() const { return 1; }
+  /// Get the version number of this stackmap. (Always returns 2).
+  unsigned getVersion() const { return 2; }
 
   /// Get the number of functions in the stack map.
   uint32_t getNumFunctions() const {
@@ -420,7 +425,7 @@ private:
   static const unsigned NumRecordsOffset = NumConstantsOffset + sizeof(uint32_t);
   static const unsigned FunctionListOffset = NumRecordsOffset + sizeof(uint32_t);
 
-  static const unsigned FunctionSize = 2 * sizeof(uint64_t);
+  static const unsigned FunctionSize = 3 * sizeof(uint64_t);
   static const unsigned ConstantSize = sizeof(uint64_t);
 
   std::size_t getFunctionOffset(unsigned FunctionIndex) const {