1 files changed, 39 insertions, 2 deletions
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index 09ffe8f..38dc863 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -9,6 +9,7 @@
 #include "llvm/ProfileData/RawMemProfReader.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MD5.h"
+#include "llvm/Support/raw_ostream.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
@@ -24,6 +25,9 @@ using ::llvm::DILocal;
 using ::llvm::memprof::CallStackMap;
 using ::llvm::memprof::MemInfoBlock;
 using ::llvm::memprof::MemProfRecord;
+using ::llvm::memprof::MemProfSchema;
+using ::llvm::memprof::Meta;
+using ::llvm::memprof::PortableMemInfoBlock;
 using ::llvm::memprof::RawMemProfReader;
 using ::llvm::memprof::SegmentEntry;
 using ::llvm::object::SectionedAddress;
@@ -99,6 +103,14 @@ MATCHER_P4(FrameContains, Function, LineOffset, Column, Inline, "") {
   return false;
 }
 
+MemProfSchema getFullSchema() {
+  MemProfSchema Schema;
+#define MIBEntryDef(NameTag, Name, Type) Schema.push_back(Meta::Name);
+#include "llvm/ProfileData/MIBEntryDef.inc"
+#undef MIBEntryDef
+  return Schema;
+}
+
 TEST(MemProf, FillsValue) {
   std::unique_ptr<MockSymbolizer> Symbolizer(new MockSymbolizer());
 
@@ -136,8 +148,8 @@ TEST(MemProf, FillsValue) {
   }
   EXPECT_EQ(Records.size(), 2U);
 
-  EXPECT_EQ(Records[0].Info.AllocCount, 1U);
-  EXPECT_EQ(Records[1].Info.AllocCount, 2U);
+  EXPECT_EQ(Records[0].Info.getAllocCount(), 1U);
+  EXPECT_EQ(Records[1].Info.getAllocCount(), 2U);
   EXPECT_THAT(Records[0].CallStack[0], FrameContains("foo", 5U, 30U, false));
   EXPECT_THAT(Records[0].CallStack[1], FrameContains("bar", 51U, 20U, true));
 
@@ -147,4 +159,29 @@ TEST(MemProf, FillsValue) {
   EXPECT_THAT(Records[1].CallStack[3], FrameContains("bar", 51U, 20U, true));
 }
 
+TEST(MemProf, PortableWrapper) {
+  MemInfoBlock Info(/*size=*/16, /*access_count=*/7, /*alloc_timestamp=*/1000,
+                    /*dealloc_timestamp=*/2000, /*alloc_cpu=*/3,
+                    /*dealloc_cpu=*/4);
+
+  const auto Schema = getFullSchema();
+  PortableMemInfoBlock WriteBlock(Info);
+
+  std::string Buffer;
+  llvm::raw_string_ostream OS(Buffer);
+  WriteBlock.serialize(Schema, OS);
+  OS.flush();
+
+  PortableMemInfoBlock ReadBlock(
+      Schema, reinterpret_cast<const unsigned char *>(Buffer.data()));
+
+  EXPECT_EQ(ReadBlock, WriteBlock);
+  // Here we compare directly with the actual counts instead of MemInfoBlock
+  // members. Since the MemInfoBlock struct is packed and the EXPECT_EQ macros
+  // take a reference to the params, this results in unaligned accesses.
+  EXPECT_EQ(1, ReadBlock.getAllocCount());
+  EXPECT_EQ(7, ReadBlock.getTotalAccessCount());
+  EXPECT_EQ(3, ReadBlock.getAllocCpuId());
+}
+
 } // namespace