llvm/lib/MC/MCPseudoProbe.cpp

   1 //===- lib/MC/MCPseudoProbe.cpp - Pseudo probe encoding support ----------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "llvm/MC/MCPseudoProbe.h"
  10 #include "llvm/ADT/STLExtras.h"
  11 #include "llvm/IR/PseudoProbe.h"
  12 #include "llvm/MC/MCAsmInfo.h"
  13 #include "llvm/MC/MCAssembler.h"
  14 #include "llvm/MC/MCContext.h"
  15 #include "llvm/MC/MCExpr.h"
  16 #include "llvm/MC/MCFragment.h"
  17 #include "llvm/MC/MCObjectFileInfo.h"
  18 #include "llvm/MC/MCObjectStreamer.h"
  19 #include "llvm/MC/MCSymbol.h"
  20 #include "llvm/Support/Endian.h"
  21 #include "llvm/Support/Error.h"
  22 #include "llvm/Support/LEB128.h"
  23 #include "llvm/Support/MD5.h"
  24 #include "llvm/Support/raw_ostream.h"
  25 #include <algorithm>
  26 #include <cassert>
  27 #include <limits>
  28 #include <memory>
  29 #include <sstream>
  30 #include <vector>
  31
  32 #define DEBUG_TYPE "mcpseudoprobe"
  33
  34 using namespace llvm;
  35 using namespace support;
  36
  37 #ifndef NDEBUG
  38 int MCPseudoProbeTable::DdgPrintIndent = 0;
  39 #endif
  40
  41 static const MCExpr *buildSymbolDiff(MCObjectStreamer *MCOS, const MCSymbol *A,
  42                                      const MCSymbol *B) {
  43   MCContext &Context = MCOS->getContext();
  44   MCSymbolRefExpr::VariantKind Variant = MCSymbolRefExpr::VK_None;
  45   const MCExpr *ARef = MCSymbolRefExpr::create(A, Variant, Context);
  46   const MCExpr *BRef = MCSymbolRefExpr::create(B, Variant, Context);
  47   const MCExpr *AddrDelta =
  48       MCBinaryExpr::create(MCBinaryExpr::Sub, ARef, BRef, Context);
  49   return AddrDelta;
  50 }
  51
  52 uint64_t MCDecodedPseudoProbe::getGuid() const { return InlineTree->Guid; }
  53
  54 void MCPseudoProbe::emit(MCObjectStreamer *MCOS,
  55                          const MCPseudoProbe *LastProbe) const {
  56   bool IsSentinel = isSentinelProbe(getAttributes());
  57   assert((LastProbe || IsSentinel) &&
  58          "Last probe should not be null for non-sentinel probes");
  59
  60   // Emit Index
  61   MCOS->emitULEB128IntValue(Index);
  62   // Emit Type and the flag:
  63   // Type (bit 0 to 3), with bit 4 to 6 for attributes.
  64   // Flag (bit 7, 0 - code address, 1 - address delta). This indicates whether
  65   // the following field is a symbolic code address or an address delta.
  66   // Emit FS discriminator
  67   assert(Type <= 0xF && "Probe type too big to encode, exceeding 15");
  68   auto NewAttributes = Attributes;
  69   if (Discriminator)
  70     NewAttributes |= (uint32_t)PseudoProbeAttributes::HasDiscriminator;
  71   assert(NewAttributes <= 0x7 &&
  72          "Probe attributes too big to encode, exceeding 7");
  73   uint8_t PackedType = Type | (NewAttributes << 4);
  74   uint8_t Flag =
  75       !IsSentinel ? ((int8_t)MCPseudoProbeFlag::AddressDelta << 7) : 0;
  76   MCOS->emitInt8(Flag | PackedType);
  77
  78   if (!IsSentinel) {
  79     // Emit the delta between the address label and LastProbe.
  80     const MCExpr *AddrDelta =
  81         buildSymbolDiff(MCOS, Label, LastProbe->getLabel());
  82     int64_t Delta;
  83     if (AddrDelta->evaluateAsAbsolute(Delta, MCOS->getAssemblerPtr())) {
  84       MCOS->emitSLEB128IntValue(Delta);
  85     } else {
  86       MCOS->insert(MCOS->getContext().allocFragment<MCPseudoProbeAddrFragment>(
  87           AddrDelta));
  88     }
  89   } else {
  90     // Emit the GUID of the split function that the sentinel probe represents.
  91     MCOS->emitInt64(Guid);
  92   }
  93
  94   if (Discriminator)
  95     MCOS->emitULEB128IntValue(Discriminator);
  96
  97   LLVM_DEBUG({
  98     dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
  99     dbgs() << "Probe: " << Index << "\n";
 100   });
 101 }
 102
 103 void MCPseudoProbeInlineTree::addPseudoProbe(
 104     const MCPseudoProbe &Probe, const MCPseudoProbeInlineStack &InlineStack) {
 105   // The function should not be called on the root.
 106   assert(isRoot() && "Should only be called on root");
 107
 108   // When it comes here, the input look like:
 109   //    Probe: GUID of C, ...
 110   //    InlineStack: [88, A], [66, B]
 111   // which means, Function A inlines function B at call site with a probe id of
 112   // 88, and B inlines C at probe 66. The tri-tree expects a tree path like {[0,
 113   // A], [88, B], [66, C]} to locate the tree node where the probe should be
 114   // added. Note that the edge [0, A] means A is the top-level function we are
 115   // emitting probes for.
 116
 117   // Make a [0, A] edge.
 118   // An empty inline stack means the function that the probe originates from
 119   // is a top-level function.
 120   InlineSite Top;
 121   if (InlineStack.empty()) {
 122     Top = InlineSite(Probe.getGuid(), 0);
 123   } else {
 124     Top = InlineSite(std::get<0>(InlineStack.front()), 0);
 125   }
 126
 127   auto *Cur = getOrAddNode(Top);
 128
 129   // Make interior edges by walking the inline stack. Once it's done, Cur should
 130   // point to the node that the probe originates from.
 131   if (!InlineStack.empty()) {
 132     auto Iter = InlineStack.begin();
 133     auto Index = std::get<1>(*Iter);
 134     Iter++;
 135     for (; Iter != InlineStack.end(); Iter++) {
 136       // Make an edge by using the previous probe id and current GUID.
 137       Cur = Cur->getOrAddNode(InlineSite(std::get<0>(*Iter), Index));
 138       Index = std::get<1>(*Iter);
 139     }
 140     Cur = Cur->getOrAddNode(InlineSite(Probe.getGuid(), Index));
 141   }
 142
 143   Cur->Probes.push_back(Probe);
 144 }
 145
 146 void MCPseudoProbeInlineTree::emit(MCObjectStreamer *MCOS,
 147                                    const MCPseudoProbe *&LastProbe) {
 148   LLVM_DEBUG({
 149     dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
 150     dbgs() << "Group [\n";
 151     MCPseudoProbeTable::DdgPrintIndent += 2;
 152   });
 153   assert(!isRoot() && "Root should be handled separately");
 154
 155   // Emit probes grouped by GUID.
 156   LLVM_DEBUG({
 157     dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
 158     dbgs() << "GUID: " << Guid << "\n";
 159   });
 160   // Emit Guid
 161   MCOS->emitInt64(Guid);
 162   // Emit number of probes in this node, including a sentinel probe for
 163   // top-level functions if needed.
 164   bool NeedSentinel = false;
 165   if (Parent->isRoot()) {
 166     assert(isSentinelProbe(LastProbe->getAttributes()) &&
 167            "Starting probe of a top-level function should be a sentinel probe");
 168     // The main body of a split function doesn't need a sentinel probe.
 169     if (LastProbe->getGuid() != Guid)
 170       NeedSentinel = true;
 171   }
 172
 173   MCOS->emitULEB128IntValue(Probes.size() + NeedSentinel);
 174   // Emit number of direct inlinees
 175   MCOS->emitULEB128IntValue(Children.size());
 176   // Emit sentinel probe for top-level functions
 177   if (NeedSentinel)
 178     LastProbe->emit(MCOS, nullptr);
 179
 180   // Emit probes in this group
 181   for (const auto &Probe : Probes) {
 182     Probe.emit(MCOS, LastProbe);
 183     LastProbe = &Probe;
 184   }
 185
 186   // Emit sorted descendant. InlineSite is unique for each pair, so there will
 187   // be no ordering of Inlinee based on MCPseudoProbeInlineTree*
 188   using InlineeType = std::pair<InlineSite, MCPseudoProbeInlineTree *>;
 189   std::vector<InlineeType> Inlinees;
 190   for (const auto &Child : Children)
 191     Inlinees.emplace_back(Child.first, Child.second.get());
 192   llvm::sort(Inlinees, llvm::less_first());
 193
 194   for (const auto &Inlinee : Inlinees) {
 195     // Emit probe index
 196     MCOS->emitULEB128IntValue(std::get<1>(Inlinee.first));
 197     LLVM_DEBUG({
 198       dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
 199       dbgs() << "InlineSite: " << std::get<1>(Inlinee.first) << "\n";
 200     });
 201     // Emit the group
 202     Inlinee.second->emit(MCOS, LastProbe);
 203   }
 204
 205   LLVM_DEBUG({
 206     MCPseudoProbeTable::DdgPrintIndent -= 2;
 207     dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
 208     dbgs() << "]\n";
 209   });
 210 }
 211
 212 void MCPseudoProbeSections::emit(MCObjectStreamer *MCOS) {
 213   MCContext &Ctx = MCOS->getContext();
 214   SmallVector<std::pair<MCSymbol *, MCPseudoProbeInlineTree *>> Vec;
 215   Vec.reserve(MCProbeDivisions.size());
 216   for (auto &ProbeSec : MCProbeDivisions)
 217     Vec.emplace_back(ProbeSec.first, &ProbeSec.second);
 218   for (auto I : llvm::enumerate(MCOS->getAssembler()))
 219     I.value().setOrdinal(I.index());
 220   llvm::sort(Vec, [](auto A, auto B) {
 221     return A.first->getSection().getOrdinal() <
 222            B.first->getSection().getOrdinal();
 223   });
 224   for (auto [FuncSym, RootPtr] : Vec) {
 225     const auto &Root = *RootPtr;
 226     if (auto *S = Ctx.getObjectFileInfo()->getPseudoProbeSection(
 227             FuncSym->getSection())) {
 228       // Switch to the .pseudoprobe section or a comdat group.
 229       MCOS->switchSection(S);
 230       // Emit probes grouped by GUID.
 231       // Emit sorted descendant. InlineSite is unique for each pair, so there
 232       // will be no ordering of Inlinee based on MCPseudoProbeInlineTree*
 233       using InlineeType = std::pair<InlineSite, MCPseudoProbeInlineTree *>;
 234       std::vector<InlineeType> Inlinees;
 235       for (const auto &Child : Root.getChildren())
 236         Inlinees.emplace_back(Child.first, Child.second.get());
 237       llvm::sort(Inlinees, llvm::less_first());
 238
 239       for (const auto &Inlinee : Inlinees) {
 240         // Emit the group guarded by a sentinel probe.
 241         MCPseudoProbe SentinelProbe(
 242             const_cast<MCSymbol *>(FuncSym), MD5Hash(FuncSym->getName()),
 243             (uint32_t)PseudoProbeReservedId::Invalid,
 244             (uint32_t)PseudoProbeType::Block,
 245             (uint32_t)PseudoProbeAttributes::Sentinel, 0);
 246         const MCPseudoProbe *Probe = &SentinelProbe;
 247         Inlinee.second->emit(MCOS, Probe);
 248       }
 249     }
 250   }
 251 }
 252
 253 //
 254 // This emits the pseudo probe tables.
 255 //
 256 void MCPseudoProbeTable::emit(MCObjectStreamer *MCOS) {
 257   MCContext &Ctx = MCOS->getContext();
 258   auto &ProbeTable = Ctx.getMCPseudoProbeTable();
 259
 260   // Bail out early so we don't switch to the pseudo_probe section needlessly
 261   // and in doing so create an unnecessary (if empty) section.
 262   auto &ProbeSections = ProbeTable.getProbeSections();
 263   if (ProbeSections.empty())
 264     return;
 265
 266   LLVM_DEBUG(MCPseudoProbeTable::DdgPrintIndent = 0);
 267
 268   // Put out the probe.
 269   ProbeSections.emit(MCOS);
 270 }
 271
 272 static StringRef getProbeFNameForGUID(const GUIDProbeFunctionMap &GUID2FuncMAP,
 273                                       uint64_t GUID) {
 274   auto It = GUID2FuncMAP.find(GUID);
 275   assert(It != GUID2FuncMAP.end() &&
 276          "Probe function must exist for a valid GUID");
 277   return It->FuncName;
 278 }
 279
 280 void MCPseudoProbeFuncDesc::print(raw_ostream &OS) {
 281   OS << "GUID: " << FuncGUID << " Name: " << FuncName << "\n";
 282   OS << "Hash: " << FuncHash << "\n";
 283 }
 284
 285 void MCDecodedPseudoProbe::getInlineContext(
 286     SmallVectorImpl<MCPseudoProbeFrameLocation> &ContextStack,
 287     const GUIDProbeFunctionMap &GUID2FuncMAP) const {
 288   uint32_t Begin = ContextStack.size();
 289   MCDecodedPseudoProbeInlineTree *Cur = InlineTree;
 290   // It will add the string of each node's inline site during iteration.
 291   // Note that it won't include the probe's belonging function(leaf location)
 292   while (Cur->hasInlineSite()) {
 293     StringRef FuncName = getProbeFNameForGUID(GUID2FuncMAP, Cur->Parent->Guid);
 294     ContextStack.emplace_back(MCPseudoProbeFrameLocation(
 295         FuncName, std::get<1>(Cur->getInlineSite())));
 296     Cur = static_cast<MCDecodedPseudoProbeInlineTree *>(Cur->Parent);
 297   }
 298   // Make the ContextStack in caller-callee order
 299   std::reverse(ContextStack.begin() + Begin, ContextStack.end());
 300 }
 301
 302 std::string MCDecodedPseudoProbe::getInlineContextStr(
 303     const GUIDProbeFunctionMap &GUID2FuncMAP) const {
 304   std::ostringstream OContextStr;
 305   SmallVector<MCPseudoProbeFrameLocation, 16> ContextStack;
 306   getInlineContext(ContextStack, GUID2FuncMAP);
 307   for (auto &Cxt : ContextStack) {
 308     if (OContextStr.str().size())
 309       OContextStr << " @ ";
 310     OContextStr << Cxt.first.str() << ":" << Cxt.second;
 311   }
 312   return OContextStr.str();
 313 }
 314
 315 static const char *PseudoProbeTypeStr[3] = {"Block", "IndirectCall",
 316                                             "DirectCall"};
 317
 318 void MCDecodedPseudoProbe::print(raw_ostream &OS,
 319                                  const GUIDProbeFunctionMap &GUID2FuncMAP,
 320                                  bool ShowName) const {
 321   OS << "FUNC: ";
 322   if (ShowName) {
 323     StringRef FuncName = getProbeFNameForGUID(GUID2FuncMAP, getGuid());
 324     OS << FuncName.str() << " ";
 325   } else {
 326     OS << getGuid() << " ";
 327   }
 328   OS << "Index: " << Index << "  ";
 329   if (Discriminator)
 330     OS << "Discriminator: " << Discriminator << "  ";
 331   OS << "Type: " << PseudoProbeTypeStr[static_cast<uint8_t>(Type)] << "  ";
 332   std::string InlineContextStr = getInlineContextStr(GUID2FuncMAP);
 333   if (InlineContextStr.size()) {
 334     OS << "Inlined: @ ";
 335     OS << InlineContextStr;
 336   }
 337   OS << "\n";
 338 }
 339
 340 template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readUnencodedNumber() {
 341   if (Data + sizeof(T) > End) {
 342     return std::error_code();
 343   }
 344   T Val = endian::readNext<T, llvm::endianness::little>(Data);
 345   return ErrorOr<T>(Val);
 346 }
 347
 348 template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readUnsignedNumber() {
 349   unsigned NumBytesRead = 0;
 350   uint64_t Val = decodeULEB128(Data, &NumBytesRead);
 351   if (Val > std::numeric_limits<T>::max() || (Data + NumBytesRead > End)) {
 352     return std::error_code();
 353   }
 354   Data += NumBytesRead;
 355   return ErrorOr<T>(static_cast<T>(Val));
 356 }
 357
 358 template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readSignedNumber() {
 359   unsigned NumBytesRead = 0;
 360   int64_t Val = decodeSLEB128(Data, &NumBytesRead);
 361   if (Val > std::numeric_limits<T>::max() || (Data + NumBytesRead > End)) {
 362     return std::error_code();
 363   }
 364   Data += NumBytesRead;
 365   return ErrorOr<T>(static_cast<T>(Val));
 366 }
 367
 368 ErrorOr<StringRef> MCPseudoProbeDecoder::readString(uint32_t Size) {
 369   StringRef Str(reinterpret_cast<const char *>(Data), Size);
 370   if (Data + Size > End) {
 371     return std::error_code();
 372   }
 373   Data += Size;
 374   return ErrorOr<StringRef>(Str);
 375 }
 376
 377 bool MCPseudoProbeDecoder::buildGUID2FuncDescMap(const uint8_t *Start,
 378                                                  std::size_t Size,
 379                                                  bool IsMMapped) {
 380   // The pseudo_probe_desc section has a format like:
 381   // .section .pseudo_probe_desc,"",@progbits
 382   // .quad -5182264717993193164   // GUID
 383   // .quad 4294967295             // Hash
 384   // .uleb 3                      // Name size
 385   // .ascii "foo"                 // Name
 386   // .quad -2624081020897602054
 387   // .quad 174696971957
 388   // .uleb 34
 389   // .ascii "main"
 390
 391   Data = Start;
 392   End = Data + Size;
 393
 394   uint32_t FuncDescCount = 0;
 395   while (Data < End) {
 396     // GUID
 397     if (!readUnencodedNumber<uint64_t>())
 398       return false;
 399     // Hash
 400     if (!readUnencodedNumber<uint64_t>())
 401       return false;
 402
 403     auto ErrorOrNameSize = readUnsignedNumber<uint32_t>();
 404     if (!ErrorOrNameSize)
 405       return false;
 406     // Function name
 407     if (!readString(*ErrorOrNameSize))
 408       return false;
 409     ++FuncDescCount;
 410   }
 411   assert(Data == End && "Have unprocessed data in pseudo_probe_desc section");
 412   GUID2FuncDescMap.reserve(FuncDescCount);
 413
 414   Data = Start;
 415   End = Data + Size;
 416   while (Data < End) {
 417     uint64_t GUID =
 418         cantFail(errorOrToExpected(readUnencodedNumber<uint64_t>()));
 419     uint64_t Hash =
 420         cantFail(errorOrToExpected(readUnencodedNumber<uint64_t>()));
 421     uint32_t NameSize =
 422         cantFail(errorOrToExpected(readUnsignedNumber<uint32_t>()));
 423     StringRef Name = cantFail(errorOrToExpected(readString(NameSize)));
 424
 425     // Initialize PseudoProbeFuncDesc and populate it into GUID2FuncDescMap
 426     GUID2FuncDescMap.emplace_back(
 427         GUID, Hash, IsMMapped ? Name : Name.copy(FuncNameAllocator));
 428   }
 429   assert(Data == End && "Have unprocessed data in pseudo_probe_desc section");
 430   assert(GUID2FuncDescMap.size() == FuncDescCount &&
 431          "Mismatching function description count pre- and post-parsing");
 432   llvm::sort(GUID2FuncDescMap, [](const auto &LHS, const auto &RHS) {
 433     return LHS.FuncGUID < RHS.FuncGUID;
 434   });
 435   return true;
 436 }
 437
 438 template <bool IsTopLevelFunc>
 439 bool MCPseudoProbeDecoder::buildAddress2ProbeMap(
 440     MCDecodedPseudoProbeInlineTree *Cur, uint64_t &LastAddr,
 441     const Uint64Set &GuidFilter, const Uint64Map &FuncStartAddrs,
 442     const uint32_t CurChildIndex) {
 443   // The pseudo_probe section encodes an inline forest and each tree has a
 444   // format defined in MCPseudoProbe.h
 445
 446   uint32_t Index = 0;
 447   if (IsTopLevelFunc) {
 448     // Use a sequential id for top level inliner.
 449     Index = CurChildIndex;
 450   } else {
 451     // Read inline site for inlinees
 452     Index = cantFail(errorOrToExpected(readUnsignedNumber<uint32_t>()));
 453   }
 454
 455   // Read guid
 456   uint64_t Guid = cantFail(errorOrToExpected(readUnencodedNumber<uint64_t>()));
 457
 458   // Decide if top-level node should be disgarded.
 459   if (IsTopLevelFunc && !GuidFilter.empty() && !GuidFilter.count(Guid))
 460     Cur = nullptr;
 461
 462   // If the incoming node is null, all its children nodes should be disgarded.
 463   if (Cur) {
 464     // Switch/add to a new tree node(inlinee)
 465     Cur->getChildren()[CurChildIndex] =
 466         MCDecodedPseudoProbeInlineTree(InlineSite(Guid, Index), Cur);
 467     Cur = &Cur->getChildren()[CurChildIndex];
 468     if (IsTopLevelFunc && !EncodingIsAddrBased) {
 469       if (auto V = FuncStartAddrs.lookup(Guid))
 470         LastAddr = V;
 471     }
 472   }
 473
 474   // Read number of probes in the current node.
 475   uint32_t NodeCount =
 476       cantFail(errorOrToExpected(readUnsignedNumber<uint32_t>()));
 477   uint32_t CurrentProbeCount = 0;
 478   // Read number of direct inlinees
 479   uint32_t ChildrenToProcess =
 480       cantFail(errorOrToExpected(readUnsignedNumber<uint32_t>()));
 481   // Read all probes in this node
 482   for (std::size_t I = 0; I < NodeCount; I++) {
 483     // Read index
 484     uint32_t Index =
 485         cantFail(errorOrToExpected(readUnsignedNumber<uint32_t>()));
 486     // Read type | flag.
 487     uint8_t Value = cantFail(errorOrToExpected(readUnencodedNumber<uint8_t>()));
 488     uint8_t Kind = Value & 0xf;
 489     uint8_t Attr = (Value & 0x70) >> 4;
 490     // Read address
 491     uint64_t Addr = 0;
 492     if (Value & 0x80) {
 493       int64_t Offset = cantFail(errorOrToExpected(readSignedNumber<int64_t>()));
 494       Addr = LastAddr + Offset;
 495     } else {
 496       Addr = cantFail(errorOrToExpected(readUnencodedNumber<int64_t>()));
 497       if (isSentinelProbe(Attr)) {
 498         // For sentinel probe, the addr field actually stores the GUID of the
 499         // split function. Convert it to the real address.
 500         if (auto V = FuncStartAddrs.lookup(Addr))
 501           Addr = V;
 502       } else {
 503         // For now we assume all probe encoding should be either based on
 504         // leading probe address or function start address.
 505         // The scheme is for downwards compatibility.
 506         // TODO: retire this scheme once compatibility is no longer an issue.
 507         EncodingIsAddrBased = true;
 508       }
 509     }
 510
 511     uint32_t Discriminator = 0;
 512     if (hasDiscriminator(Attr)) {
 513       Discriminator =
 514           cantFail(errorOrToExpected(readUnsignedNumber<uint32_t>()));
 515     }
 516
 517     if (Cur && !isSentinelProbe(Attr)) {
 518       PseudoProbeVec.emplace_back(Addr, Index, PseudoProbeType(Kind), Attr,
 519                                   Discriminator, Cur);
 520       ++CurrentProbeCount;
 521     }
 522     LastAddr = Addr;
 523   }
 524
 525   if (Cur) {
 526     Cur->setProbes(
 527         MutableArrayRef(PseudoProbeVec).take_back(CurrentProbeCount));
 528     InlineTreeVec.resize(InlineTreeVec.size() + ChildrenToProcess);
 529     Cur->getChildren() =
 530         MutableArrayRef(InlineTreeVec).take_back(ChildrenToProcess);
 531   }
 532   for (uint32_t I = 0; I < ChildrenToProcess; I++) {
 533     buildAddress2ProbeMap<false>(Cur, LastAddr, GuidFilter, FuncStartAddrs, I);
 534   }
 535   return Cur;
 536 }
 537
 538 template <bool IsTopLevelFunc>
 539 bool MCPseudoProbeDecoder::countRecords(bool &Discard, uint32_t &ProbeCount,
 540                                         uint32_t &InlinedCount,
 541                                         const Uint64Set &GuidFilter) {
 542   if (!IsTopLevelFunc)
 543     // Read inline site for inlinees
 544     if (!readUnsignedNumber<uint32_t>())
 545       return false;
 546
 547   // Read guid
 548   auto ErrorOrCurGuid = readUnencodedNumber<uint64_t>();
 549   if (!ErrorOrCurGuid)
 550     return false;
 551   uint64_t Guid = std::move(*ErrorOrCurGuid);
 552
 553   // Decide if top-level node should be disgarded.
 554   if (IsTopLevelFunc) {
 555     Discard = !GuidFilter.empty() && !GuidFilter.count(Guid);
 556     if (!Discard)
 557       // Allocate an entry for top-level function record.
 558       ++InlinedCount;
 559   }
 560
 561   // Read number of probes in the current node.
 562   auto ErrorOrNodeCount = readUnsignedNumber<uint32_t>();
 563   if (!ErrorOrNodeCount)
 564     return false;
 565   uint32_t NodeCount = std::move(*ErrorOrNodeCount);
 566   uint32_t CurrentProbeCount = 0;
 567
 568   // Read number of direct inlinees
 569   auto ErrorOrCurChildrenToProcess = readUnsignedNumber<uint32_t>();
 570   if (!ErrorOrCurChildrenToProcess)
 571     return false;
 572   uint32_t ChildrenToProcess = std::move(*ErrorOrCurChildrenToProcess);
 573
 574   // Read all probes in this node
 575   for (std::size_t I = 0; I < NodeCount; I++) {
 576     // Read index
 577     if (!readUnsignedNumber<uint32_t>())
 578       return false;
 579
 580     // Read type | flag.
 581     auto ErrorOrValue = readUnencodedNumber<uint8_t>();
 582     if (!ErrorOrValue)
 583       return false;
 584     uint8_t Value = std::move(*ErrorOrValue);
 585
 586     uint8_t Attr = (Value & 0x70) >> 4;
 587     if (Value & 0x80) {
 588       // Offset
 589       if (!readSignedNumber<int64_t>())
 590         return false;
 591     } else {
 592       // Addr
 593       if (!readUnencodedNumber<int64_t>())
 594         return false;
 595     }
 596
 597     if (hasDiscriminator(Attr))
 598       // Discriminator
 599       if (!readUnsignedNumber<uint32_t>())
 600         return false;
 601
 602     if (!Discard && !isSentinelProbe(Attr))
 603       ++CurrentProbeCount;
 604   }
 605
 606   if (!Discard) {
 607     ProbeCount += CurrentProbeCount;
 608     InlinedCount += ChildrenToProcess;
 609   }
 610
 611   for (uint32_t I = 0; I < ChildrenToProcess; I++)
 612     if (!countRecords<false>(Discard, ProbeCount, InlinedCount, GuidFilter))
 613       return false;
 614   return true;
 615 }
 616
 617 bool MCPseudoProbeDecoder::buildAddress2ProbeMap(
 618     const uint8_t *Start, std::size_t Size, const Uint64Set &GuidFilter,
 619     const Uint64Map &FuncStartAddrs) {
 620   // For function records in the order of their appearance in the encoded data
 621   // (DFS), count the number of contained probes and inlined function records.
 622   uint32_t ProbeCount = 0;
 623   uint32_t InlinedCount = 0;
 624   uint32_t TopLevelFuncs = 0;
 625   Data = Start;
 626   End = Data + Size;
 627   bool Discard = false;
 628   while (Data < End) {
 629     if (!countRecords<true>(Discard, ProbeCount, InlinedCount, GuidFilter))
 630       return false;
 631     TopLevelFuncs += !Discard;
 632   }
 633   assert(Data == End && "Have unprocessed data in pseudo_probe section");
 634   PseudoProbeVec.reserve(ProbeCount);
 635   InlineTreeVec.reserve(InlinedCount);
 636
 637   // Allocate top-level function records as children of DummyInlineRoot.
 638   InlineTreeVec.resize(TopLevelFuncs);
 639   DummyInlineRoot.getChildren() = MutableArrayRef(InlineTreeVec);
 640
 641   Data = Start;
 642   End = Data + Size;
 643   uint64_t LastAddr = 0;
 644   uint32_t CurChildIndex = 0;
 645   while (Data < End)
 646     CurChildIndex += buildAddress2ProbeMap<true>(
 647         &DummyInlineRoot, LastAddr, GuidFilter, FuncStartAddrs, CurChildIndex);
 648   assert(Data == End && "Have unprocessed data in pseudo_probe section");
 649   assert(PseudoProbeVec.size() == ProbeCount &&
 650          "Mismatching probe count pre- and post-parsing");
 651   assert(InlineTreeVec.size() == InlinedCount &&
 652          "Mismatching function records count pre- and post-parsing");
 653
 654   std::vector<std::pair<uint64_t, uint32_t>> SortedA2P(ProbeCount);
 655   for (const auto &[I, Probe] : llvm::enumerate(PseudoProbeVec))
 656     SortedA2P[I] = {Probe.getAddress(), I};
 657   llvm::sort(SortedA2P);
 658   Address2ProbesMap.reserve(ProbeCount);
 659   for (const uint32_t I : llvm::make_second_range(SortedA2P))
 660     Address2ProbesMap.emplace_back(PseudoProbeVec[I]);
 661   SortedA2P.clear();
 662   return true;
 663 }
 664
 665 void MCPseudoProbeDecoder::printGUID2FuncDescMap(raw_ostream &OS) {
 666   OS << "Pseudo Probe Desc:\n";
 667   for (auto &I : GUID2FuncDescMap)
 668     I.print(OS);
 669 }
 670
 671 void MCPseudoProbeDecoder::printProbeForAddress(raw_ostream &OS,
 672                                                 uint64_t Address) {
 673   for (const MCDecodedPseudoProbe &Probe : Address2ProbesMap.find(Address)) {
 674     OS << " [Probe]:\t";
 675     Probe.print(OS, GUID2FuncDescMap, true);
 676   }
 677 }
 678
 679 void MCPseudoProbeDecoder::printProbesForAllAddresses(raw_ostream &OS) {
 680   uint64_t PrevAddress = INT64_MAX;
 681   for (MCDecodedPseudoProbe &Probe : Address2ProbesMap) {
 682     uint64_t Address = Probe.getAddress();
 683     if (Address != PrevAddress) {
 684       PrevAddress = Address;
 685       OS << "Address:\t" << Address << '\n';
 686     }
 687     OS << " [Probe]:\t";
 688     Probe.print(OS, GUID2FuncDescMap, true);
 689   }
 690 }
 691
 692 const MCDecodedPseudoProbe *
 693 MCPseudoProbeDecoder::getCallProbeForAddr(uint64_t Address) const {
 694   const MCDecodedPseudoProbe *CallProbe = nullptr;
 695   for (const MCDecodedPseudoProbe &Probe : Address2ProbesMap.find(Address)) {
 696     if (Probe.isCall()) {
 697       // Disabling the assert and returning first call probe seen so far.
 698       // Subsequent call probes, if any, are ignored. Due to the the way
 699       // .pseudo_probe section is decoded, probes of the same-named independent
 700       // static functions are merged thus multiple call probes may be seen for a
 701       // callsite. This should only happen to compiler-generated statics, with
 702       // -funique-internal-linkage-names where user statics get unique names.
 703       //
 704       // TODO: re-enable or narrow down the assert to static functions only.
 705       //
 706       // assert(!CallProbe &&
 707       //        "There should be only one call probe corresponding to address "
 708       //        "which is a callsite.");
 709       CallProbe = &Probe;
 710       break;
 711     }
 712   }
 713   return CallProbe;
 714 }
 715
 716 const MCPseudoProbeFuncDesc *
 717 MCPseudoProbeDecoder::getFuncDescForGUID(uint64_t GUID) const {
 718   auto It = GUID2FuncDescMap.find(GUID);
 719   assert(It != GUID2FuncDescMap.end() && "Function descriptor doesn't exist");
 720   return &*It;
 721 }
 722
 723 void MCPseudoProbeDecoder::getInlineContextForProbe(
 724     const MCDecodedPseudoProbe *Probe,
 725     SmallVectorImpl<MCPseudoProbeFrameLocation> &InlineContextStack,
 726     bool IncludeLeaf) const {
 727   Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap);
 728   if (!IncludeLeaf)
 729     return;
 730   // Note that the context from probe doesn't include leaf frame,
 731   // hence we need to retrieve and prepend leaf if requested.
 732   const auto *FuncDesc = getFuncDescForGUID(Probe->getGuid());
 733   InlineContextStack.emplace_back(
 734       MCPseudoProbeFrameLocation(FuncDesc->FuncName, Probe->getIndex()));
 735 }
 736
 737 const MCPseudoProbeFuncDesc *MCPseudoProbeDecoder::getInlinerDescForProbe(
 738     const MCDecodedPseudoProbe *Probe) const {
 739   MCDecodedPseudoProbeInlineTree *InlinerNode = Probe->getInlineTreeNode();
 740   if (!InlinerNode->hasInlineSite())
 741     return nullptr;
 742   return getFuncDescForGUID(InlinerNode->Parent->Guid);
 743 }