1 files changed, 91 insertions, 9 deletions

diff --git a/llvm/unittests/XRay/FDRProducerConsumerTest.cpp b/llvm/unittests/XRay/FDRProducerConsumerTest.cpp
index 09ec44db26e..03fabdf2372 100644
--- a/llvm/unittests/XRay/FDRProducerConsumerTest.cpp
+++ b/llvm/unittests/XRay/FDRProducerConsumerTest.cpp
@@ -30,13 +30,10 @@ namespace {
 using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::Not;
+using ::testing::SizeIs;
 
 template <class RecordType> std::unique_ptr<Record> MakeRecord();
 
-template <> std::unique_ptr<Record> MakeRecord<BufferExtents>() {
-  return make_unique<BufferExtents>(1);
-}
-
 template <> std::unique_ptr<Record> MakeRecord<NewBufferRecord>() {
   return make_unique<NewBufferRecord>(1);
 }
@@ -69,10 +66,18 @@ template <> std::unique_ptr<Record> MakeRecord<FunctionRecord>() {
   return make_unique<FunctionRecord>(RecordTypes::ENTER, 1, 2);
 }
 
+template <> std::unique_ptr<Record> MakeRecord<CustomEventRecordV5>() {
+  return make_unique<CustomEventRecordV5>(4, 1, "data");
+}
+
+template <> std::unique_ptr<Record> MakeRecord<TypedEventRecord>() {
+  return make_unique<TypedEventRecord>(4, 1, 2, "data");
+}
+
 template <class T> class RoundTripTest : public ::testing::Test {
 public:
   RoundTripTest() : Data(), OS(Data) {
-    H.Version = 3;
+    H.Version = 4;
     H.Type = 1;
     H.ConstantTSC = true;
     H.NonstopTSC = true;
@@ -92,9 +97,36 @@ protected:
 
 TYPED_TEST_CASE_P(RoundTripTest);
 
+template <class T> class RoundTripTestV5 : public ::testing::Test {
+public:
+  RoundTripTestV5() : Data(), OS(Data) {
+    H.Version = 5;
+    H.Type = 1;
+    H.ConstantTSC = true;
+    H.NonstopTSC = true;
+    H.CycleFrequency = 3e9;
+
+    Writer = make_unique<FDRTraceWriter>(OS, H);
+    Rec = MakeRecord<T>();
+  }
+
+protected:
+  std::string Data;
+  raw_string_ostream OS;
+  XRayFileHeader H;
+  std::unique_ptr<FDRTraceWriter> Writer;
+  std::unique_ptr<Record> Rec;
+};
+
+TYPED_TEST_CASE_P(RoundTripTestV5);
+
 // This test ensures that the writing and reading implementations are in sync --
 // that given write(read(write(R))) == R.
 TYPED_TEST_P(RoundTripTest, RoundTripsSingleValue) {
+  // Always write a buffer extents record which will cover the correct size of
+  // the record, for version 3 and up.
+  BufferExtents BE(200);
+  ASSERT_FALSE(errorToBool(BE.apply(*this->Writer)));
   auto &R = this->Rec;
   ASSERT_FALSE(errorToBool(R->apply(*this->Writer)));
   this->OS.flush();
@@ -125,17 +157,67 @@ TYPED_TEST_P(RoundTripTest, RoundTripsSingleValue) {
 
   EXPECT_EQ(Data2.substr(sizeof(XRayFileHeader)),
             this->Data.substr(sizeof(XRayFileHeader)));
-  EXPECT_THAT(Records[0]->type(), Eq(R->type()));
+  ASSERT_THAT(Records, SizeIs(2));
+  EXPECT_THAT(Records[1]->getRecordType(), Eq(R->getRecordType()));
 }
 
 REGISTER_TYPED_TEST_CASE_P(RoundTripTest, RoundTripsSingleValue);
 
+// We duplicate the above case for the V5 version using different types and
+// encodings.
+TYPED_TEST_P(RoundTripTestV5, RoundTripsSingleValue) {
+  BufferExtents BE(200);
+  ASSERT_FALSE(errorToBool(BE.apply(*this->Writer)));
+  auto &R = this->Rec;
+  ASSERT_FALSE(errorToBool(R->apply(*this->Writer)));
+  this->OS.flush();
+
+  DataExtractor DE(this->Data, sys::IsLittleEndianHost, 8);
+  uint32_t OffsetPtr = 0;
+  auto HeaderOrErr = readBinaryFormatHeader(DE, OffsetPtr);
+  if (!HeaderOrErr)
+    FAIL() << HeaderOrErr.takeError();
+
+  FileBasedRecordProducer P(HeaderOrErr.get(), DE, OffsetPtr);
+  std::vector<std::unique_ptr<Record>> Records;
+  LogBuilderConsumer C(Records);
+  while (DE.isValidOffsetForDataOfSize(OffsetPtr, 1)) {
+    auto R = P.produce();
+    if (!R)
+      FAIL() << R.takeError();
+    if (auto E = C.consume(std::move(R.get())))
+      FAIL() << E;
+  }
+  ASSERT_THAT(Records, Not(IsEmpty()));
+  std::string Data2;
+  raw_string_ostream OS2(Data2);
+  FDRTraceWriter Writer2(OS2, this->H);
+  for (auto &P : Records)
+    ASSERT_FALSE(errorToBool(P->apply(Writer2)));
+  OS2.flush();
+
+  EXPECT_EQ(Data2.substr(sizeof(XRayFileHeader)),
+            this->Data.substr(sizeof(XRayFileHeader)));
+  ASSERT_THAT(Records, SizeIs(2));
+  EXPECT_THAT(Records[1]->getRecordType(), Eq(R->getRecordType()));
+}
+
+REGISTER_TYPED_TEST_CASE_P(RoundTripTestV5, RoundTripsSingleValue);
+
+// These are the record types we support for v4 and below.
 using RecordTypes =
-    ::testing::Types<BufferExtents, NewBufferRecord, NewCPUIDRecord,
-                     TSCWrapRecord, WallclockRecord, CustomEventRecord,
-                     CallArgRecord, BufferExtents, PIDRecord, FunctionRecord>;
+    ::testing::Types<NewBufferRecord, NewCPUIDRecord, TSCWrapRecord,
+                     WallclockRecord, CustomEventRecord, CallArgRecord,
+                     PIDRecord, FunctionRecord>;
 INSTANTIATE_TYPED_TEST_CASE_P(Records, RoundTripTest, RecordTypes);
 
+// For V5, we have two new types we're supporting.
+using RecordTypesV5 =
+    ::testing::Types<NewBufferRecord, NewCPUIDRecord, TSCWrapRecord,
+                     WallclockRecord, CustomEventRecordV5, TypedEventRecord,
+                     CallArgRecord, PIDRecord, FunctionRecord>;
+INSTANTIATE_TYPED_TEST_CASE_P(Records, RoundTripTestV5, RecordTypesV5);
+
 } // namespace
 } // namespace xray
 } // namespace llvm