[XRay] Add a BlockIndexer visitor for FDR Records.

[dberris]

Summary

This change adds a BlockIndexer type which maintains pointers to
records that belong to the same process+thread pairs. The indexing
happens with order of appearance of records as they are visited.

This version of the indexer currently only supports FDR version 3 logs,
which contain BufferExtent records. We will add support for v2 and v1
logs in follow-up patches.

This is another part of D50441.

[mboerger]

LGTM

llvm/lib/XRay/BlockIndexer.cpp

74

Decide for an order and be consistent.

llvm/unittests/XRay/FDRBlockIndexerTest.cpp

65

consider SizeIs() matcher.

72

consider using just ElementsAre() matcher without the assertion first.

[dberris]

Address comments by @mboerger.

llvm/include/llvm/XRay/BlockIndexer.h

+//===- BlockIndexer.h - FDR Block Indexing Visitor ------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// An implementation of the RecordVisitor which generates a mapping between a
+// thread and a range of records representing a block.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_LIB_XRAY_BLOCKINDEXER_H_
+#define LLVM_LIB_XRAY_BLOCKINDEXER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/XRay/FDRRecords.h"
+#include <cstdint>
+#include <vector>
+
+namespace llvm {
+namespace xray {
+
+// The BlockIndexer will gather all related records associated with a
+// process+thread and group them by 'Block'.
+class BlockIndexer : public RecordVisitor {
+public:
+  struct Block {
+    uint64_t ProcessID;
+    int32_t ThreadID;
+    std::vector<Record *> Records;
+  };
+
+  // This maps the process + thread combination to a sequence of blocks.
+  using Index = DenseMap<std::pair<uint64_t, int32_t>, std::vector<Block>>;
+
+private:
+  Index &Indices;
+
+  enum class State : unsigned { SeekExtents, ExtentsFound, ThreadIDFound };
+
+  State CurrentState = State::SeekExtents;
+  Block CurrentBlock{0, 0, {}};
+
+public:
+  explicit BlockIndexer(Index &I) : RecordVisitor(), Indices(I) {}
+
+  Error visit(BufferExtents &) override;
+  Error visit(WallclockRecord &) override;
+  Error visit(NewCPUIDRecord &) override;
+  Error visit(TSCWrapRecord &) override;
+  Error visit(CustomEventRecord &) override;
+  Error visit(CallArgRecord &) override;
+  Error visit(PIDRecord &) override;
+  Error visit(NewBufferRecord &) override;
+  Error visit(EndBufferRecord &) override;
+  Error visit(FunctionRecord &) override;
+
+  /// The flush() function will clear out the current state of the visitor, to
+  /// allow for explicitly flushing a block's records to the currently
+  /// recognized thread and process combination.
+  Error flush();
+};
+
+} // namespace xray
+} // namespace llvm
+
+#endif // LLVM_LIB_XRAY_BLOCKINDEXER_H_

llvm/lib/XRay/BlockIndexer.cpp

+//===- BlockIndexer.cpp - FDR Block Indexing VIsitor ----------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// An implementation of the RecordVisitor which generates a mapping between a
+// thread and a range of records representing a block.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/XRay/BlockIndexer.h"
+
+namespace llvm {
+namespace xray {
+
+Error BlockIndexer::visit(BufferExtents &) {
+  if (CurrentState == State::ThreadIDFound) {
+    Index::iterator It;
+    std::tie(It, std::ignore) =
+        Indices.insert({{CurrentBlock.ProcessID, CurrentBlock.ThreadID}, {}});
+    It->second.push_back({CurrentBlock.ProcessID, CurrentBlock.ThreadID,
+                          std::move(CurrentBlock.Records)});
+    CurrentBlock.ProcessID = 0;
+    CurrentBlock.ThreadID = 0;
+    CurrentBlock.Records = {};
+  }
+  CurrentState = State::ExtentsFound;
+  return Error::success();
+}
+
+Error BlockIndexer::visit(WallclockRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::visit(NewCPUIDRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::visit(TSCWrapRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::visit(CustomEventRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::visit(CallArgRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+};
+
+Error BlockIndexer::visit(PIDRecord &R) {
+  CurrentBlock.ProcessID = R.pid();
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::visit(NewBufferRecord &R) {
+  CurrentState = State::ThreadIDFound;
+  CurrentBlock.ThreadID = R.tid();
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::visit(EndBufferRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  CurrentState = State::SeekExtents;
+  return Error::success();
+}
+
+Error BlockIndexer::visit(FunctionRecord &R) {
+  CurrentBlock.Records.push_back(&R);
+  return Error::success();
+}
+
+Error BlockIndexer::flush() {
+  Index::iterator It;
+  std::tie(It, std::ignore) =
+      Indices.insert({{CurrentBlock.ProcessID, CurrentBlock.ThreadID}, {}});
+  It->second.push_back({CurrentBlock.ProcessID, CurrentBlock.ThreadID,
+                        std::move(CurrentBlock.Records)});
+  CurrentBlock.ProcessID = 0;
+  CurrentBlock.ThreadID = 0;
+  CurrentBlock.Records = {};
+  CurrentState = State::SeekExtents;
+  return Error::success();
+}
+
+} // namespace xray
+} // namespace llvm

llvm/lib/XRay/CMakeLists.txt

 add_llvm_library(LLVMXRay
+  BlockIndexer.cpp
   FDRRecordProducer.cpp
   FDRRecords.cpp
   FDRTraceWriter.cpp
   FileHeaderReader.cpp
   InstrumentationMap.cpp
   LogBuilderConsumer.cpp
   Profile.cpp
   RecordInitializer.cpp

llvm/unittests/XRay/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Support
   XRay
   )
 
 add_llvm_unittest(XRayTests
   ProfileTest.cpp
+  FDRBlockIndexerTest.cpp
   FDRProducerConsumerTest.cpp
   FDRRecordPrinterTest.cpp
   FDRTraceWriterTest.cpp
   GraphTest.cpp
   )
 
 add_dependencies(XRayTests intrinsics_gen)

llvm/unittests/XRay/FDRBlockIndexerTest.cpp

+//===- llvm/unittest/XRay/FDRTraceWriterTest.cpp ----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/XRay/BlockIndexer.h"
+#include "llvm/XRay/FDRLogBuilder.h"
+#include "llvm/XRay/FDRRecords.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace llvm {
+namespace xray {
+namespace {
+
+using ::testing::Eq;
+using ::testing::Not;
+
+// This test ensures that we can index blocks that follow version 3 of the log
+// format.
+TEST(FDRBlockIndexerTest, IndexBlocksV3) {
+  auto Block0 = LogBuilder()
+                    .add<BufferExtents>(80)
+                    .add<NewBufferRecord>(1)
+                    .add<WallclockRecord>(1, 2)
+                    .add<PIDRecord>(1)
+                    .add<NewCPUIDRecord>(1)
+                    .add<FunctionRecord>(RecordTypes::ENTER, 1, 1)
+                    .add<FunctionRecord>(RecordTypes::EXIT, 1, 100)
+                    .consume();
+  auto Block1 = LogBuilder()
+                    .add<BufferExtents>(80)
+                    .add<NewBufferRecord>(1)
+                    .add<WallclockRecord>(1, 2)
+                    .add<PIDRecord>(1)
+                    .add<NewCPUIDRecord>(1)
+                    .add<FunctionRecord>(RecordTypes::ENTER, 1, 1)
+                    .add<FunctionRecord>(RecordTypes::EXIT, 1, 100)
+                    .consume();
+  auto Block2 = LogBuilder()
+                    .add<BufferExtents>(80)
+                    .add<NewBufferRecord>(2)
+                    .add<WallclockRecord>(1, 2)
+                    .add<PIDRecord>(1)
+                    .add<NewCPUIDRecord>(2)
+                    .add<FunctionRecord>(RecordTypes::ENTER, 1, 1)
+                    .add<FunctionRecord>(RecordTypes::EXIT, 1, 100)
+                    .consume();
+  BlockIndexer::Index Index;
+  BlockIndexer Indexer(Index);
+  // Iterate through the contrived blocks we have created above.
+  for (auto B : {std::ref(Block0), std::ref(Block1), std::ref(Block2)}) {
+    // For each record in the block, we apply the indexer.
+    for (auto &R : B.get())
+      ASSERT_FALSE(errorToBool(R->apply(Indexer)));
+    ASSERT_FALSE(errorToBool(Indexer.flush()));
+  }
+
+  ASSERT_THAT(Index.size(), Eq(2u));
+  auto T1Blocks = Index.find({1, 1});
+  ASSERT_THAT(T1Blocks, Not(Eq(Index.end())));
+  ASSERT_THAT(T1Blocks->second.size(), Eq(2u));
+
+  // Expect only six records, because we're ignoring the BufferExtents record.
+  EXPECT_THAT(T1Blocks->second[0].Records.size(), Eq(6u));
+  EXPECT_THAT(T1Blocks->second[1].Records.size(), Eq(6u));
+  auto T2Blocks = Index.find({1, 2});
+  ASSERT_THAT(T2Blocks, Not(Eq(Index.end())));
+  ASSERT_THAT(T2Blocks->second.size(), Eq(1u));
+  EXPECT_THAT(T2Blocks->second[0].Records.size(), Eq(6u));
+}
+
+// FIXME: Support indexing V2 and V1 blocks.
+
+} // namespace
+} // namespace xray
+} // namespace llvm