[XRay][tools] Use Support/JSON.h in llvm-xray convert

[dberris]

Summary

This change removes the ad-hoc implementation used by llvm-xray's
convert subcommand to generate JSON encoded catapult (AKA Chrome
Trace Viewer) trace output, to instead use the JSON encoder now in the
Support library.

[kpw]

This is way easier to read now. Did you sanity check loading a trace into chrome trace viewer that it outputs?

llvm/tools/llvm-xray/xray-converter.cpp

282 ↗

(On Diff #158462)

Strange that the code you're replacing changed the precision between 4 and 3 depending on the version. Four seems like a better choice and matches the behavior for version >= 3.

324 ↗

(On Diff #158462)

Is it meaningful to have the trailing comma just before the end of the initializer list?

I am kind of surprised to see that this works.

337 ↗

(On Diff #158462)

What does {0:2} mean in this context? Is there a Json particular format provider?

[dberris]

Add comment on the format line for JSON.

[dberris]

In D50129#1186823, @kpw wrote:

This is way easier to read now. Did you sanity check loading a trace into chrome trace viewer that it outputs?

Yep!

llvm/tools/llvm-xray/xray-converter.cpp

324 ↗

(On Diff #158462)

Not really, but that's valid C++, and is not creating an initializer list (but rather is initializing a KV).

337 ↗

(On Diff #158462)

Yep, this pretty-prints with indents of two characters.

llvm/trunk/tools/llvm-xray/xray-converter.cpp

 #include "xray-converter.h"
 
 #include "trie-node.h"
 #include "xray-registry.h"
 #include "llvm/DebugInfo/Symbolize/Symbolize.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/JSON.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/XRay/InstrumentationMap.h"
 #include "llvm/XRay/Trace.h"
 #include "llvm/XRay/YAMLXRayRecord.h"
 
 using namespace llvm;
StackTrieNode *findOrCreateStackNode(
   NodeStore.push_front({FuncId, Parent, {}, {stack_id, std::move(siblings)}});
   StackTrieNode *CurrentStack = &NodeStore.front();
   for (auto *sibling : CurrentStack->ExtraData.siblings)
     sibling->ExtraData.siblings.push_back(CurrentStack);
   ParentCallees.push_back(CurrentStack);
   return CurrentStack;
 }
 
-void writeTraceViewerRecord(uint16_t Version, raw_ostream &OS, int32_t FuncId,
-                            uint32_t TId, uint32_t PId, bool Symbolize,
-                            const FuncIdConversionHelper &FuncIdHelper,
-                            double EventTimestampUs,
-                            const StackTrieNode &StackCursor,
-                            StringRef FunctionPhenotype) {
-  OS << "    ";
-  if (Version >= 3) {
-    OS << llvm::formatv(
-        R"({ "name" : "{0}", "ph" : "{1}", "tid" : "{2}", "pid" : "{3}", )"
-        R"("ts" : "{4:f4}", "sf" : "{5}" })",
-        (Symbolize ? FuncIdHelper.SymbolOrNumber(FuncId)
-                   : llvm::to_string(FuncId)),
-        FunctionPhenotype, TId, PId, EventTimestampUs,
-        StackCursor.ExtraData.id);
-  } else {
-    OS << llvm::formatv(
-        R"({ "name" : "{0}", "ph" : "{1}", "tid" : "{2}", "pid" : "1", )"
-        R"("ts" : "{3:f3}", "sf" : "{4}" })",
-        (Symbolize ? FuncIdHelper.SymbolOrNumber(FuncId)
-                   : llvm::to_string(FuncId)),
-        FunctionPhenotype, TId, EventTimestampUs, StackCursor.ExtraData.id);
-  }
-}
-
 } // namespace
 
 void TraceConverter::exportAsChromeTraceEventFormat(const Trace &Records,
                                                     raw_ostream &OS) {
   const auto &FH = Records.getFileHeader();
   auto Version = FH.Version;
   auto CycleFreq = FH.CycleFrequency;
 
   unsigned id_counter = 0;
 
-  OS << "{\n  \"traceEvents\": [";
   DenseMap<uint32_t, StackTrieNode *> StackCursorByThreadId{};
   DenseMap<uint32_t, SmallVector<StackTrieNode *, 4>> StackRootsByThreadId{};
   DenseMap<unsigned, StackTrieNode *> StacksByStackId{};
   std::forward_list<StackTrieNode> NodeStore{};
-  int loop_count = 0;
-  for (const auto &R : Records) {
-    if (loop_count++ == 0)
-      OS << "\n";
-    else
-      OS << ",\n";
 
+  // Create a JSON Array which will hold all trace events.
+  json::Array TraceEvents;
+  for (const auto &R : Records) {
     // Chrome trace event format always wants data in micros.
     // CyclesPerMicro = CycleHertz / 10^6
     // TSC / CyclesPerMicro == TSC * 10^6 / CycleHertz == MicroTimestamp
     // Could lose some precision here by converting the TSC to a double to
     // multiply by the period in micros. 52 bit mantissa is a good start though.
     // TODO: Make feature request to Chrome Trace viewer to accept ticks and a
     // frequency or do some more involved calculation to avoid dangers of
     // conversion.
     double EventTimestampUs = double(1000000) / CycleFreq * double(R.TSC);
     StackTrieNode *&StackCursor = StackCursorByThreadId[R.TId];
     switch (R.Type) {
     case RecordTypes::ENTER:
     case RecordTypes::ENTER_ARG:
       StackCursor = findOrCreateStackNode(StackCursor, R.FuncId, R.TId,
                                           StackRootsByThreadId, StacksByStackId,
                                           &id_counter, NodeStore);
       // Each record is represented as a json dictionary with function name,
       // type of B for begin or E for end, thread id, process id,
       // timestamp in microseconds, and a stack frame id. The ids are logged
       // in an id dictionary after the events.
-      writeTraceViewerRecord(Version, OS, R.FuncId, R.TId, R.PId, Symbolize,
-                             FuncIdHelper, EventTimestampUs, *StackCursor, "B");
+      TraceEvents.push_back(json::Object({
+          {"name", Symbolize ? FuncIdHelper.SymbolOrNumber(R.FuncId)
+                             : llvm::to_string(R.FuncId)},
+          {"ph", "B"},
+          {"tid", llvm::to_string(R.TId)},
+          {"pid", llvm::to_string(Version >= 3 ? R.PId : 1)},
+          {"ts", llvm::formatv("{0:f4}", EventTimestampUs)},
+          {"sf", llvm::to_string(StackCursor->ExtraData.id)},
+      }));
       break;
     case RecordTypes::EXIT:
     case RecordTypes::TAIL_EXIT:
       // No entries to record end for.
       if (StackCursor == nullptr)
         break;
       // Should we emit an END record anyway or account this condition?
       // (And/Or in loop termination below)
       StackTrieNode *PreviousCursor = nullptr;
       do {
-        if (PreviousCursor != nullptr) {
-          OS << ",\n";
-        }
-        writeTraceViewerRecord(Version, OS, StackCursor->FuncId, R.TId, R.PId,
-                               Symbolize, FuncIdHelper, EventTimestampUs,
-                               *StackCursor, "E");
+        TraceEvents.push_back(json::Object({
+            {"name", Symbolize
+                         ? FuncIdHelper.SymbolOrNumber(StackCursor->FuncId)
+                         : llvm::to_string(StackCursor->FuncId)},
+            {"ph", "E"},
+            {"tid", llvm::to_string(R.TId)},
+            {"pid", llvm::to_string(Version >= 3 ? R.PId : 1)},
+            {"ts", llvm::formatv("{0:f4}", EventTimestampUs)},
+            {"sf", llvm::to_string(StackCursor->ExtraData.id)},
+        }));
         PreviousCursor = StackCursor;
         StackCursor = StackCursor->Parent;
       } while (PreviousCursor->FuncId != R.FuncId && StackCursor != nullptr);
       break;
     }
   }
-  OS << "\n  ],\n"; // Close the Trace Events array.
-  OS << "  "
-     << "\"displayTimeUnit\": \"ns\",\n";
 
   // The stackFrames dictionary substantially reduces size of the output file by
   // avoiding repeating the entire call stack of function names for each entry.
-  OS << R"(  "stackFrames": {)";
-  int stack_frame_count = 0;
-  for (auto map_iter : StacksByStackId) {
-    if (stack_frame_count++ == 0)
-      OS << "\n";
-    else
-      OS << ",\n";
-    OS << "    ";
-    OS << llvm::formatv(
-        R"("{0}" : { "name" : "{1}")", map_iter.first,
-        (Symbolize ? FuncIdHelper.SymbolOrNumber(map_iter.second->FuncId)
-                   : llvm::to_string(map_iter.second->FuncId)));
-    if (map_iter.second->Parent != nullptr)
-      OS << llvm::formatv(R"(, "parent": "{0}")",
-                          map_iter.second->Parent->ExtraData.id);
-    OS << " }";
+  json::Object StackFrames;
+  for (const auto &Stack : StacksByStackId) {
+    const auto &StackId = Stack.first;
+    const auto &StackFunctionNode = Stack.second;
+    json::Object::iterator It;
+    std::tie(It, std::ignore) = StackFrames.insert({
+        llvm::to_string(StackId),
+        json::Object{
+            {"name",
+             Symbolize ? FuncIdHelper.SymbolOrNumber(StackFunctionNode->FuncId)
+                       : llvm::to_string(StackFunctionNode->FuncId)}},
+    });
+
+    if (StackFunctionNode->Parent != nullptr)
+      It->second.getAsObject()->insert(
+          {"parent", llvm::to_string(StackFunctionNode->Parent->ExtraData.id)});
   }
-  OS << "\n  }\n"; // Close the stack frames map.
-  OS << "}\n";     // Close the JSON entry.
+
+  json::Object TraceJSON{
+      {"displayTimeUnit", "ns"},
+      {"traceEvents", std::move(TraceEvents)},
+      {"stackFrames", std::move(StackFrames)},
+  };
+
+  // Pretty-print the JSON using two spaces for indentations.
+  OS << formatv("{0:2}", json::Value(std::move(TraceJSON)));
 }
 
 namespace llvm {
 namespace xray {
 
 static CommandRegistration Unused(&Convert, []() -> Error {
   // FIXME: Support conversion to BINARY when upgrading XRay trace versions.
   InstrumentationMap Map;