//===-- instr.cpp -----------------------------------------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
// This file contains code that is linked to the final binary with a function
// that is called at program exit to dump instrumented data collected during
// execution.
//
//===----------------------------------------------------------------------===//
//
// BOLT runtime instrumentation library for x86 Linux.
//
//===----------------------------------------------------------------------===//

#include <cstdint>
#include <elf.h>

//#define ENABLE_DEBUG

#ifdef ENABLE_DEBUG
#define DEBUG(X)                                                               \
  { X; }
#else
#define DEBUG(X)                                                               \
  {}
#endif

// All extern declarations here need to be defined by BOLT itself

// Counters inserted by instrumentation, incremented during runtime when
// points of interest (locations) in the program are reached.
extern uint64_t __bolt_instr_locations[];
// Number of call descriptions
extern uint32_t __bolt_instr_num_calls;
// Number of function descriptions
extern uint32_t __bolt_instr_num_funcs;
// Filename to dump data to
extern char __bolt_instr_filename[];

// Anonymous namespace covering everything but our library entry point
namespace {

// We use a stack-allocated buffer for string manipulation in some
// places with this size.
constexpr uint32_t BufSize = 10240;

// A location is a function name plus offset. Function name needs to be
// retrieved from the string table and is stored as an index to this table.
struct Location {
  uint32_t FunctionName;
  uint32_t Offset;
};

struct CallDescription {
  Location From;
  Location To;
  uint32_t Counter;
};

// An edge description defines an instrumented edge in the program, fully
// identified by where the jump is located and its destination.
struct EdgeDescription {
  Location From;
  uint32_t FromNode;
  Location To;
  uint32_t ToNode;
  uint32_t Counter;
};

// This description is used for instrumented basic blocks. We only instrument
// blocks that are leaves of the spanning tree associated with the CFG. We use
// their execution count as a starting point to determine the frequency of the
// incoming edge in the spanning tree (the one that is not explictly
// instrumented). All other edges (not in the spanning tree) have an associated
// counter themselves. We never write a node's counter to the profile, it is
// only used for inference of other edges.
struct InstrumentedNode {
  uint32_t Node;
  uint32_t Counter;
};

// General metadata describing a function - number of CFG nodes, leaf nodes and
// edges.
struct FunctionDescription {
  uint32_t NumNodes;
  InstrumentedNode *LeafNodes;
  uint32_t NumEdges;
  EdgeDescription *Edges;
};

// These need to be read from disk. They are generated by BOLT and written to
// an ELF note section in the binary itself.
struct InstrumentationInfo {
  CallDescription *CallDescriptions;
  uint8_t *FuncDescriptions;
  char *Strings;  // String table with function names used in this binary
  int FileDesc;   // File descriptor for the file on disk backing this
                  // information in memory via mmap
  void *MMapPtr;  // The mmap ptr
  int MMapSize;   // The mmap size
};

// Declare some syscall wrappers we use throughout this code to avoid linking
// against system libc.
uint64_t __open(const char *pathname, uint64_t flags, uint64_t mode) {
  uint64_t ret;
  __asm__ __volatile__ (
          "movq $2, %%rax\n"
          "syscall"
          : "=a"(ret)
          : "D"(pathname), "S"(flags), "d"(mode)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

uint64_t __write(uint64_t fd, const void *buf, uint64_t count) {
  uint64_t ret;
  __asm__ __volatile__ (
          "movq $1, %%rax\n"
          "syscall\n"
          : "=a"(ret)
          : "D"(fd), "S"(buf), "d"(count)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

uint64_t __lseek(uint64_t fd, uint64_t pos, uint64_t whence) {
  uint64_t ret;
  __asm__ __volatile__ (
          "movq $8, %%rax\n"
          "syscall\n"
          : "=a"(ret)
          : "D"(fd), "S"(pos), "d"(whence)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

int __close(uint64_t fd) {
  uint64_t ret;
  __asm__ __volatile__ (
          "movq $3, %%rax\n"
          "syscall\n"
          : "=a"(ret)
          : "D"(fd)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

void *__mmap(uint64_t addr, uint64_t size, uint64_t prot, uint64_t flags,
             uint64_t fd, uint64_t offset) {
  void *ret;
  register uint64_t r8 asm("r8") = fd;
  register uint64_t r9 asm("r9") = offset;
  register uint64_t r10 asm("r10") = flags;
  __asm__ __volatile__ (
          "movq $9, %%rax\n"
          "syscall\n"
          : "=a"(ret)
          : "D"(addr), "S"(size), "d"(prot), "r"(r10), "r"(r8), "r"(r9)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

uint64_t __munmap(void *addr, uint64_t size) {
  uint64_t ret;
  __asm__ __volatile__ (
          "movq $11, %%rax\n"
          "syscall\n"
          : "=a"(ret)
          : "D"(addr), "S"(size)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

uint64_t __exit(uint64_t code) {
  uint64_t ret;
  __asm__ __volatile__ (
          "movq $231, %%rax\n"
          "syscall\n"
          : "=a"(ret)
          : "D"(code)
          : "cc", "rcx", "r11", "memory");
  return ret;
}

// Helper functions for writing strings to the .fdata file

// Write number Num using Base to the buffer in OutBuf, returns a pointer to
// the end of the string.
char *intToStr(char *OutBuf, uint64_t Num, uint32_t Base) {
  const char *Chars = "0123456789abcdef";
  char Buf[21];
  char *Ptr = Buf;
  while (Num) {
    *Ptr++ = *(Chars + (Num % Base));
    Num /= Base;
  }
  if (Ptr == Buf) {
    *OutBuf++ = '0';
    return OutBuf;
  }
  while (Ptr != Buf) {
    *OutBuf++ = *--Ptr;
  }
  return OutBuf;
}

// Copy Str to OutBuf, returns a pointer to the end of the copied string
char *strCopy(char *OutBuf, const char *Str, int32_t Size = BufSize) {
  while (*Str) {
    *OutBuf++ = *Str++;
    if (--Size <= 0)
      return OutBuf;
  }
  return OutBuf;
}

void memSet(char *Buf, char C, uint32_t Size) {
  for (int I = 0; I < Size; ++I)
    *Buf++ = C;
}

uint32_t strLen(const char *Str) {
  uint32_t Size = 0;
  while (*Str++)
    ++Size;
  return Size;
}

// Print Msg to STDERR and quits with error code 1
void reportError(const char *Msg, uint64_t Size) {
  __write(2, Msg, Size);
  __exit(1);
}

void assert(bool Assertion, const char *Msg) {
  if (Assertion)
    return;
  char Buf[BufSize];
  char *Ptr = Buf;
  Ptr = strCopy(Ptr, "Assertion failed: ");
  Ptr = strCopy(Ptr, Msg, BufSize - 40);
  Ptr = strCopy(Ptr, "\n");
  reportError(Buf, Ptr - Buf);
}

void reportNumber(const char *Msg, uint64_t Num, uint32_t Base) {
  char Buf[BufSize];
  char *Ptr = Buf;
  Ptr = strCopy(Ptr, Msg, BufSize - 23);
  Ptr = intToStr(Ptr, Num, Base);
  Ptr = strCopy(Ptr, "\n");
  __write(2, Buf, Ptr - Buf);
}

void report(const char *Msg) {
  __write(2, Msg, strLen(Msg));
}

inline uint64_t alignTo(uint64_t Value, uint64_t Align) {
  return (Value + Align - 1) / Align * Align;
}

/// A simple allocator that mmaps a fixed size region and manages this space
/// in a stack fashion, meaning you always deallocate the last element that
/// was allocated.
class BumpPtrAllocator {
  struct EntryMetadata {
    uint64_t Magic;
    uint64_t AllocSize;
  };
public:
  void *allocate(uintptr_t Size) {
    if (StackBase == nullptr) {
      StackBase = reinterpret_cast<uint8_t *>(__mmap(
          0, MAXSIZE, 0x3 /* PROT_READ | PROT_WRITE*/,
          0x22 /* MAP_PRIVATE | MAP_ANONYMOUS*/, -1, 0));
      StackSize = 0;
    }
    Size = alignTo(Size + sizeof(EntryMetadata), 16);
    uint8_t * AllocAddress = StackBase + StackSize + sizeof(EntryMetadata);
    auto *M = reinterpret_cast<EntryMetadata *>(StackBase + StackSize);
    M->Magic = MAGIC;
    M->AllocSize = Size;
    StackSize += Size;
    assert(StackSize < MAXSIZE, "allocator ran out of memory");
    return AllocAddress;
  }

  void deallocate(void *Ptr) {
    uint8_t MetadataOffset = sizeof(EntryMetadata);
    auto *M = reinterpret_cast<EntryMetadata *>(
        reinterpret_cast<uint8_t *>(Ptr) - MetadataOffset);
    const uint8_t *StackTop = StackBase + StackSize + MetadataOffset;
    // Validate size
    if (Ptr != StackTop - M->AllocSize) {
      // This could be a pointer returned by operator new []
      MetadataOffset +=
          sizeof(uint64_t); // Space for number of elements alloc'ed
      M = reinterpret_cast<EntryMetadata *>(reinterpret_cast<uint8_t *>(Ptr) -
                                            MetadataOffset);
      assert(Ptr == StackTop - M->AllocSize,
             "must deallocate the last element alloc'ed");
    }
    assert(M->Magic == MAGIC, "allocator magic is corrupt");
    StackSize -= M->AllocSize;
  }

private:
  static constexpr uint64_t MAGIC = 0x1122334455667788ull;
  static constexpr uint64_t MAXSIZE = 0xa00000;
  uint8_t *StackBase{nullptr};
  uint64_t StackSize{0};
  uint8_t *LastAlloc{nullptr};
};

} // anonymous namespace

// User-defined placement new operators. We only use those (as opposed to
// overriding the regular operator new) so we can keep our allocator in the
// stack instead of in a data section (global).
void *operator new(uintptr_t Sz, BumpPtrAllocator &A) {
  return A.allocate(Sz);
}
void *operator new (uintptr_t Sz, BumpPtrAllocator &A, char C) {
  auto *Ptr = reinterpret_cast<char *>(A.allocate(Sz));
  memSet(Ptr, C, Sz);
  return Ptr;
}
void *operator new [] (uintptr_t Sz, BumpPtrAllocator &A) {
  return A.allocate(Sz);
}
void *operator new [] (uintptr_t Sz, BumpPtrAllocator &A, char C) {
  auto *Ptr = reinterpret_cast<char *>(A.allocate(Sz));
  memSet(Ptr, C, Sz);
  return Ptr;
}
// Only called during exception unwinding (useless). We must manually dealloc.
// C++ language weirdness
void operator delete(void *Ptr, BumpPtrAllocator &A) {
  A.deallocate(Ptr);
}

namespace {

// Perform a string comparison and returns zero if Str1 matches Str2. Compares
// at most Size characters.
int compareStr(const char *Str1, const char *Str2, int Size) {
  while (*Str1 == *Str2) {
    if (*Str1 == '\0' || --Size == 0)
      return 0;
    ++Str1;
    ++Str2;
  }
  return 1;
}

// Write as a string in OutBuf an identifier for the program point at function
// whose name is in the string table index FuncStrIndex plus Offset.
char *serializeLoc(const InstrumentationInfo &Info, char *OutBuf,
                   const Location Loc, uint32_t BufSize) {
  // fdata location format: Type Name Offset
  // Type 1 - regular symbol
  OutBuf = strCopy(OutBuf, "1 ");
  const char *Str = Info.Strings + Loc.FunctionName;
  uint32_t Size = 25;
  while (*Str) {
    *OutBuf++ = *Str++;
    if (++Size >= BufSize)
      break;
  }
  assert(!*Str, "buffer overflow, function name too large");
  *OutBuf++ = ' ';
  OutBuf = intToStr(OutBuf, Loc.Offset, 16);
  *OutBuf++ = ' ';
  return OutBuf;
}

// Read and map to memory the descriptions written by BOLT into the executable's
// notes section
InstrumentationInfo readDescriptions() {
  InstrumentationInfo Result;
  uint64_t FD = __open("/proc/self/exe",
                       /*flags=*/0 /*O_RDONLY*/,
                       /*mode=*/0666);
  assert(static_cast<int64_t>(FD) > 0, "Failed to open /proc/self/exe");
  Result.FileDesc = FD;

  // mmap our binary to memory
  uint64_t Size = __lseek(FD, 0, 2 /*SEEK_END*/);
  uint8_t *BinContents = reinterpret_cast<uint8_t *>(
      __mmap(0, Size, 0x1 /* PROT_READ*/, 0x2 /* MAP_PRIVATE*/, FD, 0));
  Result.MMapPtr = BinContents;
  Result.MMapSize = Size;
  Elf64_Ehdr *Hdr = reinterpret_cast<Elf64_Ehdr *>(BinContents);
  Elf64_Shdr *Shdr = reinterpret_cast<Elf64_Shdr *>(BinContents + Hdr->e_shoff);
  Elf64_Shdr *StringTblHeader = reinterpret_cast<Elf64_Shdr *>(
      BinContents + Hdr->e_shoff + Hdr->e_shstrndx * Hdr->e_shentsize);

  // Find .bolt.instr.tables with the data we need and set pointers to it
  for (int I = 0; I < Hdr->e_shnum; ++I) {
    char *SecName = reinterpret_cast<char *>(
        BinContents + StringTblHeader->sh_offset + Shdr->sh_name);
    if (compareStr(SecName, ".bolt.instr.tables", 64) != 0) {
      Shdr = reinterpret_cast<Elf64_Shdr *>(BinContents + Hdr->e_shoff +
                                            (I + 1) * Hdr->e_shentsize);
      continue;
    }
    // Actual contents of the ELF note start after offset 20 decimal:
    // Offset 0: Producer name size (4 bytes)
    // Offset 4: Contents size (4 bytes)
    // Offset 8: Note type (4 bytes)
    // Offset 12: Producer name (BOLT\0) (5 bytes + align to 4-byte boundary)
    // Offset 20: Contents
    uint32_t CallDescSize =
        *reinterpret_cast<uint32_t *>(BinContents + Shdr->sh_offset + 20);
    uint32_t FuncDescSize = *reinterpret_cast<uint32_t *>(
        BinContents + Shdr->sh_offset + 24 + CallDescSize);
    Result.CallDescriptions =
        reinterpret_cast<CallDescription *>(BinContents + Shdr->sh_offset + 24);
    Result.FuncDescriptions = BinContents + Shdr->sh_offset + 28 + CallDescSize;
    Result.Strings = reinterpret_cast<char *>(BinContents + Shdr->sh_offset +
                                              28 + CallDescSize + FuncDescSize);
    return Result;
  }
  const char ErrMsg[] =
      "BOLT instrumentation runtime error: could not find section "
      ".bolt.instr.tables\n";
  reportError(ErrMsg, sizeof(ErrMsg));
  return Result;
}

void printStats(const InstrumentationInfo &Info) {
  char StatMsg[BufSize];
  char *StatPtr = StatMsg;
  StatPtr = strCopy(
      StatPtr, "\nBOLT INSTRUMENTATION RUNTIME STATISTICS\n\nCallDescSize: ");
  StatPtr = intToStr(StatPtr,
                     Info.FuncDescriptions -
                         reinterpret_cast<uint8_t *>(Info.CallDescriptions),
                     10);
  StatPtr = strCopy(StatPtr, "\nFuncDescSize: ");
  StatPtr = intToStr(
      StatPtr,
      reinterpret_cast<uint8_t *>(Info.Strings) - Info.FuncDescriptions, 10);
  StatPtr = strCopy(StatPtr, "\n__bolt_instr_num_calls: ");
  StatPtr = intToStr(StatPtr, __bolt_instr_num_calls, 10);
  StatPtr = strCopy(StatPtr, "\n__bolt_instr_num_funcs: ");
  StatPtr = intToStr(StatPtr, __bolt_instr_num_funcs, 10);
  StatPtr = strCopy(StatPtr, "\n");
  __write(2, StatMsg, StatPtr - StatMsg);
}

/// This is part of a simple CFG representation in memory, where we store
/// a dynamically sized array of input and output edges per node, and store
/// a dynamically sized array of nodes per graph. We also store the spanning
/// tree edges for that CFG in a separate array of nodes in
/// \p SpanningTreeNodes, while the regular nodes live in \p CFGNodes.
struct Edge {
  uint32_t Node; // Index in nodes array regarding the destination of this edge
  uint32_t ID;   // Edge index in an array comprising all edges of the graph
};

/// A regular graph node or a spanning tree node
struct Node {
  uint32_t NumInEdges{0};  // Input edge count used to size InEdge
  uint32_t NumOutEdges{0}; // Output edge count used to size OutEdges
  Edge *InEdges{nullptr};  // Created and managed by \p Graph
  Edge *OutEdges{nullptr}; // ditto
};

/// Main class for CFG representation in memory. Manages object creation and
/// destruction, populates an array of CFG nodes as well as corresponding
/// spanning tree nodes.
struct Graph {
  uint32_t NumNodes;
  Node *CFGNodes;
  Node *SpanningTreeNodes;
  BumpPtrAllocator &Alloc;

  /// Reads a list of \p NumEdgeDescs descriptions in \p EdgeDescs and builds
  /// the graph from it. Allocates several internal dynamic structures that are
  /// later destroyed by ~Graph() and uses \p Alloc. \p LeafNodes contain all
  /// spanning tree leaf nodes descriptions (their counters). They are the seed
  /// used to compute the rest of the missing edge counts in a bottom-up
  /// traversal of the spanning tree.
  Graph(BumpPtrAllocator &Alloc, const EdgeDescription *EdgeDescs,
        uint32_t NumEdgeDescs, const InstrumentedNode *LeafNodes,
        uint32_t NumLeafNodes);

  ~Graph();
  void dump() const;
};

Graph::Graph(BumpPtrAllocator &Alloc, const EdgeDescription *EdgeDescs,
             uint32_t NumEdgeDescs, const InstrumentedNode *LeafNodes,
             uint32_t NumLeafNodes) : Alloc(Alloc) {
  DEBUG(reportNumber("G = 0x", (uint64_t)this, 16));
  // First pass to determine number of nodes
  uint32_t MaxNodes = 0;
  for (int I = 0; I < NumEdgeDescs; ++I) {
    if (EdgeDescs[I].FromNode > MaxNodes)
      MaxNodes = EdgeDescs[I].FromNode;
    if (EdgeDescs[I].ToNode > MaxNodes)
      MaxNodes = EdgeDescs[I].ToNode;
  }
  for (int I = 0; I < NumLeafNodes; ++I) {
    if (LeafNodes[I].Node > MaxNodes)
      MaxNodes = LeafNodes[I].Node;
  }
  // No edges? Nothing to do
  if (!MaxNodes) {
    CFGNodes = nullptr;
    SpanningTreeNodes = nullptr;
    NumNodes = 0;
    return;
  }
  ++MaxNodes;
  DEBUG(reportNumber("NumNodes = ", MaxNodes, 10));
  NumNodes = MaxNodes;

  // Initial allocations
  CFGNodes = new (Alloc) Node[MaxNodes];
  DEBUG(reportNumber("G->CFGNodes = 0x", (uint64_t)CFGNodes, 16));
  SpanningTreeNodes = new (Alloc) Node[MaxNodes];
  DEBUG(reportNumber("G->SpanningTreeNodes = 0x",
                     (uint64_t)SpanningTreeNodes, 16));

  // Figure out how much to allocate to each vector (in/out edge sets)
  for (int I = 0; I < NumEdgeDescs; ++I) {
    CFGNodes[EdgeDescs[I].FromNode].NumOutEdges++;
    CFGNodes[EdgeDescs[I].ToNode].NumInEdges++;
    if (EdgeDescs[I].Counter != 0xffffffff)
      continue;

    SpanningTreeNodes[EdgeDescs[I].FromNode].NumOutEdges++;
    SpanningTreeNodes[EdgeDescs[I].ToNode].NumInEdges++;
  }

  // Allocate in/out edge sets
  for (int I = 0; I < MaxNodes; ++I) {
    if (CFGNodes[I].NumInEdges > 0)
      CFGNodes[I].InEdges = new (Alloc) Edge[CFGNodes[I].NumInEdges];
    if (CFGNodes[I].NumOutEdges > 0)
      CFGNodes[I].OutEdges = new (Alloc) Edge[CFGNodes[I].NumOutEdges];
    if (SpanningTreeNodes[I].NumInEdges > 0)
      SpanningTreeNodes[I].InEdges =
          new (Alloc) Edge[SpanningTreeNodes[I].NumInEdges];
    if (SpanningTreeNodes[I].NumOutEdges > 0)
      SpanningTreeNodes[I].OutEdges =
          new (Alloc) Edge[SpanningTreeNodes[I].NumOutEdges];
    CFGNodes[I].NumInEdges = 0;
    CFGNodes[I].NumOutEdges = 0;
    SpanningTreeNodes[I].NumInEdges = 0;
    SpanningTreeNodes[I].NumOutEdges = 0;
  }

  // Fill in/out edge sets
  for (int I = 0; I < NumEdgeDescs; ++I) {
    const uint32_t Src = EdgeDescs[I].FromNode;
    const uint32_t Dst = EdgeDescs[I].ToNode;
    Edge *E = &CFGNodes[Src].OutEdges[CFGNodes[Src].NumOutEdges++];
    E->Node = Dst;
    E->ID = I;

    E = &CFGNodes[Dst].InEdges[CFGNodes[Dst].NumInEdges++];
    E->Node = Src;
    E->ID = I;

    if (EdgeDescs[I].Counter != 0xffffffff)
      continue;

    E = &SpanningTreeNodes[Src]
             .OutEdges[SpanningTreeNodes[Src].NumOutEdges++];
    E->Node = Dst;
    E->ID = I;

    E = &SpanningTreeNodes[Dst]
             .InEdges[SpanningTreeNodes[Dst].NumInEdges++];
    E->Node = Src;
    E->ID = I;
  }
}

Graph::~Graph() {
  for (int I = NumNodes - 1; I >= 0; --I) {
    if (SpanningTreeNodes[I].OutEdges)
      Alloc.deallocate(SpanningTreeNodes[I].OutEdges);
    if (SpanningTreeNodes[I].InEdges)
      Alloc.deallocate(SpanningTreeNodes[I].InEdges);
    if (CFGNodes[I].OutEdges)
      Alloc.deallocate(CFGNodes[I].OutEdges);
    if (CFGNodes[I].InEdges)
      Alloc.deallocate(CFGNodes[I].InEdges);
  }
  if (SpanningTreeNodes)
    Alloc.deallocate(SpanningTreeNodes);
  if (CFGNodes)
    Alloc.deallocate(CFGNodes);
}

void Graph::dump() const {
  reportNumber("Dumping graph with number of nodes: ", NumNodes, 10);
  report("  Full graph:\n");
  for (int I = 0; I < NumNodes; ++I) {
    const Node *N = &CFGNodes[I];
    reportNumber("    Node #", I, 10);
    reportNumber("      InEdges total ", N->NumInEdges, 10);
    for (int J = 0; J < N->NumInEdges; ++J)
      reportNumber("        ", N->InEdges[J].Node, 10);
    reportNumber("      OutEdges total ", N->NumOutEdges, 10);
    for (int J = 0; J < N->NumOutEdges; ++J)
      reportNumber("        ", N->OutEdges[J].Node, 10);
    report("\n");
  }
  report("  Spanning tree:\n");
  for (int I = 0; I < NumNodes; ++I) {
    const Node *N = &SpanningTreeNodes[I];
    reportNumber("    Node #", I, 10);
    reportNumber("      InEdges total ", N->NumInEdges, 10);
    for (int J = 0; J < N->NumInEdges; ++J)
      reportNumber("        ", N->InEdges[J].Node, 10);
    reportNumber("      OutEdges total ", N->NumOutEdges, 10);
    for (int J = 0; J < N->NumOutEdges; ++J)
      reportNumber("        ", N->OutEdges[J].Node, 10);
    report("\n");
  }
}

void dumpEdgeFreqs(const uint64_t *EdgeFreqs, const EdgeDescription *EdgeDescs,
                   uint32_t NumEdges) {
  reportNumber("Dumping edge frequencies for graph with num edges: ", NumEdges,
               10);
  for (int I = 0; I < NumEdges; ++I) {
    reportNumber("* Src: ", EdgeDescs[I].FromNode, 10);
    reportNumber("  Dst: ", EdgeDescs[I].ToNode, 10);
    reportNumber("    Cnt: ", EdgeFreqs[I], 10);
  }
}

// Return an array with the frequency of each edge in the function represented
// by G.
uint64_t *computeEdgeFrequencies(BumpPtrAllocator &Alloc, Graph *G,
                                 const EdgeDescription *EdgeDescs,
                                 uint32_t NumEdges,
                                 const InstrumentedNode *LeafNodes,
                                 uint32_t NumLeafNodes,
                                 const uint64_t *Counters) {
  if (G->NumNodes == 0 || NumEdges == 0)
    return 0;
  assert(NumLeafNodes > 0, "no leaf node frequency");

  uint64_t *EdgeFrequency = new (Alloc, 0) uint64_t [NumEdges];

  // Perform a bottom-up, BFS traversal of the spanning tree in G. Edges in the
  // spanning tree don't have explicit counters. We must infer their value using
  // a linear combination of other counters (sum of counters of the outgoing
  // edges minus sum of counters of the incoming edges).
  uint32_t *Stack = new (Alloc) uint32_t [G->NumNodes];
  uint32_t StackTop = 0;
  enum Status : uint8_t { S_NEW = 0, S_VISITING, S_VISITED };
  Status *Visited = new (Alloc, 0) Status[G->NumNodes];
  uint64_t *LeafFrequency = new (Alloc, 0) uint64_t[G->NumNodes];

  // Setup a fast lookup for frequency of leaf nodes, which have special
  // basic block frequency instrumentation (they are not edge profiled).
  uint64_t TotalFreq = 0;
  for (int I = 0; I < NumLeafNodes; ++I) {
    LeafFrequency[LeafNodes[I].Node] = Counters[LeafNodes[I].Counter];
    DEBUG({
      if (Counters[LeafNodes[I].Counter] > 0) {
        reportNumber("Leaf Node# ", LeafNodes[I].Node, 10);
        reportNumber("     Counter: ", Counters[LeafNodes[I].Counter], 10);
      }
    });
    TotalFreq += Counters[LeafNodes[I].Counter];
  }
  // Add all root nodes to the stack
  for (int I = 0; I < G->NumNodes; ++I) {
    if (G->SpanningTreeNodes[I].NumInEdges == 0)
      Stack[StackTop++] = I;
  }
  // Empty stack?
  if (StackTop == 0) {
    Alloc.deallocate(LeafFrequency);
    Alloc.deallocate(Visited);
    Alloc.deallocate(Stack);
    Alloc.deallocate(EdgeFrequency);
    return 0;
  }
  // Add all known edge counts, will infer the rest
  for (int I = 0; I < NumEdges; ++I) {
    const uint32_t C = EdgeDescs[I].Counter;
    if (C == 0xffffffff) // inferred counter - we will compute its value
      continue;
    EdgeFrequency[I] = Counters[C];
    TotalFreq += Counters[C];
  }

  // This function is completely cold, there is no point in computing anything
  // since inferred edges will be zero too.
#ifndef ENABLE_DEBUG
  if (TotalFreq == 0) {
    Alloc.deallocate(LeafFrequency);
    Alloc.deallocate(Visited);
    Alloc.deallocate(Stack);
    return EdgeFrequency;
  }
#endif

  while (StackTop > 0) {
    const uint32_t Cur = Stack[--StackTop];
    DEBUG({
      if (Visited[Cur] == S_VISITING)
        report("(visiting) ");
      else
        report("(new) ");
      reportNumber("Cur: ", Cur, 10);
    });

    // This shouldn't happen in a tree
    assert(Visited[Cur] != S_VISITED, "should not have visited nodes in stack");
    if (Visited[Cur] == S_NEW) {
      Visited[Cur] = S_VISITING;
      Stack[StackTop++] = Cur;
      assert(StackTop <= G->NumNodes, "stack grew too large");
      for (int I = 0, E = G->SpanningTreeNodes[Cur].NumOutEdges; I < E; ++I) {
        const uint32_t Succ = G->SpanningTreeNodes[Cur].OutEdges[I].Node;
        Stack[StackTop++] = Succ;
        assert(StackTop <= G->NumNodes, "stack grew too large");
      }
      continue;
    }
    Visited[Cur] = S_VISITED;

    // No parent? Reached a tree root, nothing to do.
    if (G->SpanningTreeNodes[Cur].NumInEdges == 0)
      continue;

    assert(G->SpanningTreeNodes[Cur].NumInEdges == 1, "must have 1 parent");
    const uint32_t Parent = G->SpanningTreeNodes[Cur].InEdges[0].Node;
    const uint32_t ParentEdge = G->SpanningTreeNodes[Cur].InEdges[0].ID;

    // Establish our node frequency based on outgoing edges, which should all be
    // resolved by now.
    int64_t CurNodeFreq = LeafFrequency[Cur];
    // Not a leaf?
    if (!CurNodeFreq) {
      for (int I = 0, E = G->CFGNodes[Cur].NumOutEdges; I != E; ++I) {
        const uint32_t SuccEdge = G->CFGNodes[Cur].OutEdges[I].ID;
        CurNodeFreq += EdgeFrequency[SuccEdge];
      }
    }
    int64_t ParentEdgeFreq = CurNodeFreq;
    // Calculate parent edge freq.
    for (int I = 0, E = G->CFGNodes[Cur].NumInEdges; I != E; ++I) {
      const uint32_t PredEdge = G->CFGNodes[Cur].InEdges[I].ID;
      ParentEdgeFreq -= EdgeFrequency[PredEdge];
    }
    // Sometimes the conservative CFG that BOLT builds will lead to incorrect
    // flow computation. For example, in a BB that transitively calls the exit
    // syscall, BOLT will add a fall-through successor even though it should not
    // have any successors. So this block execution will likely be wrong. We
    // tolerate this imperfection since this case should be quite infrequent.
    if (ParentEdgeFreq < 0) {
      DEBUG(dumpEdgeFreqs(EdgeFrequency, EdgeDescs, NumEdges));
      DEBUG(report("WARNING: incorrect flow"));
      ParentEdgeFreq = 0;
    }
    DEBUG(reportNumber("  Setting freq for ParentEdge: ", ParentEdge, 10));
    DEBUG(reportNumber("  with ParentEdgeFreq: ", ParentEdgeFreq, 10));
    EdgeFrequency[ParentEdge] = ParentEdgeFreq;
  }

  Alloc.deallocate(LeafFrequency);
  Alloc.deallocate(Visited);
  Alloc.deallocate(Stack);
  return EdgeFrequency;
}

// Write to \p FD all of the edge profiles for function \p FuncDesc. Uses
// \p Alloc to allocate helper dynamic structures used to compute profile for
// edges that we do not explictly instrument.
uint8_t *writeFunctionProfile(int FD, const InstrumentationInfo &Info,
                              uint8_t *FuncDesc, BumpPtrAllocator &Alloc) {
  uint32_t NumLeafNodes = *reinterpret_cast<uint32_t *>(FuncDesc);
  DEBUG(reportNumber("NumLeafNodes = ", NumLeafNodes, 10));
  InstrumentedNode *LeafNodes =
      reinterpret_cast<InstrumentedNode *>(FuncDesc + 4);

  uint32_t NumEdges = *reinterpret_cast<uint32_t *>(
      FuncDesc + 4 + NumLeafNodes * sizeof(InstrumentedNode));
  DEBUG(reportNumber("NumEdges = ", NumEdges, 10));
  EdgeDescription *EdgeDescs = reinterpret_cast<EdgeDescription *>(
      FuncDesc + 8 + NumLeafNodes * sizeof(InstrumentedNode));

  uint8_t *next = (FuncDesc + 8 + NumLeafNodes * sizeof(InstrumentedNode) +
                   NumEdges * sizeof(EdgeDescription));

  // Skip funcs we know are cold
#ifndef ENABLE_DEBUG
  uint64_t LeafFreq = 0;
  for (int I = 0; I < NumLeafNodes; ++I) {
    LeafFreq += __bolt_instr_locations[LeafNodes[I].Counter];
  }
  if (LeafFreq == 0)
    return next;
#endif

  Graph *G =
      new (Alloc) Graph(Alloc, EdgeDescs, NumEdges, LeafNodes, NumLeafNodes);
  DEBUG(G->dump());
  uint64_t *Freqs =
      computeEdgeFrequencies(Alloc, G, EdgeDescs, NumEdges, LeafNodes,
                             NumLeafNodes, __bolt_instr_locations);
  if (!Freqs) {
    G->~Graph();
    Alloc.deallocate(G);
    return next;
  }

  DEBUG(dumpEdgeFreqs(Freqs, EdgeDescs, NumEdges));
  for (int I = 0; I < NumEdges; ++I) {
    const uint64_t Freq = Freqs[I];
    if (Freq == 0)
      continue;
    const EdgeDescription *Desc = &EdgeDescs[I];
    char LineBuf[BufSize];
    char *Ptr = LineBuf;
    Ptr = serializeLoc(Info, Ptr, Desc->From, BufSize);
    Ptr = serializeLoc(Info, Ptr, Desc->To, BufSize - (Ptr - LineBuf));
    Ptr = strCopy(Ptr, "0 ", BufSize - (Ptr - LineBuf) - 22);
    Ptr = intToStr(Ptr, Freq, 10);
    *Ptr++ = '\n';
    __write(FD, LineBuf, Ptr - LineBuf);
  }

  Alloc.deallocate(Freqs);
  G->~Graph();
  Alloc.deallocate(G);
  return next;
}

} // anonymous namespace

// This is the entry point called at program exit. BOLT patches the executable's
// FINI entry in the .dynamic section with the address of this function. Our
// goal here is to flush to disk all instrumentation data in memory, using
// BOLT's fdata format.
extern "C" void __bolt_instr_data_dump() {
  const InstrumentationInfo Info = readDescriptions();

  DEBUG(printStats(Info));

  uint64_t FD = __open(__bolt_instr_filename,
                       /*flags=*/0x241 /*O_WRONLY|O_TRUNC|O_CREAT*/,
                       /*mode=*/0666);
  if (static_cast<int64_t>(FD) < 0) {
    reportNumber("Assertion error: failed to open profile file for writing. "
                 "Error number: 0x",
                 0 - static_cast<int64_t>(FD), 16);
    __exit(1);
  }

  for (int I = 0, E = __bolt_instr_num_calls; I < E; ++I) {
    char LineBuf[BufSize];
    char *Ptr = LineBuf;

    CallDescription *Desc = &Info.CallDescriptions[I];
    uint64_t HitCount = __bolt_instr_locations[Desc->Counter];
    if (!HitCount)
      continue;
    Ptr = serializeLoc(Info, Ptr, Desc->From, BufSize);
    Ptr = serializeLoc(Info, Ptr, Desc->To, BufSize - (Ptr - LineBuf));
    Ptr = strCopy(Ptr, "0 ", BufSize - (Ptr - LineBuf) - 25);
    Ptr = intToStr(Ptr, HitCount, 10);
    *Ptr++ = '\n';
    __write(FD, LineBuf, Ptr - LineBuf);
  }

  BumpPtrAllocator Alloc;
  uint8_t *FuncDesc = Info.FuncDescriptions;
  for (int I = 0, E = __bolt_instr_num_funcs; I < E; ++I) {
    FuncDesc = writeFunctionProfile(FD, Info, FuncDesc, Alloc);
    DEBUG(reportNumber("FuncDesc now: ", (uint64_t)FuncDesc, 16));
  }
  assert(FuncDesc == (void *)Info.Strings,
         "FuncDesc ptr must be equal to stringtable");

  __close(FD);
  __munmap(Info.MMapPtr, Info.MMapSize);
  __close(Info.FileDesc);
}