[Support] Add JSON streaming output API, faster where the heavy value types aren't needed.

Summary:
There's still a little bit of constant factor that could be trimmed (e.g.
more overloads to avoid round-tripping primitives through json::Value).
But this solves the memory scaling problem, and greatly improves the performance
constant factor, and the API should leave room for optimization if needed.

Adapt TimeProfiler to use it, eliminating almost all the performance regression
from r358476.

Performance test on my machine:
perf stat -r 5 ~/llvmbuild-opt/bin/clang++ -w -S -ftime-trace -mllvm -time-trace-granularity=0 spirit.cpp

Handcrafted JSON (HEAD=r358532 with r358476 reverted): 2480ms
json::Value (HEAD): 2757ms (+11%)
After this patch: 2520 ms (+1.6%)

Reviewers: anton-afanasyev, lebedev.ri

Subscribers: kristina, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D60804

llvm-svn: 359186

2 files changed, 148 insertions, 129 deletions

diff --git a/llvm/lib/Support/JSON.cpp b/llvm/lib/Support/JSON.cpp
index 790b28f6e31..95e5ed65427 100644
--- a/llvm/lib/Support/JSON.cpp
+++ b/llvm/lib/Support/JSON.cpp
@@ -560,9 +560,6 @@ std::string fixUTF8(llvm::StringRef S) {
   return Res;
 }
 
-} // namespace json
-} // namespace llvm
-
 static void quote(llvm::raw_ostream &OS, llvm::StringRef S) {
   OS << '\"';
   for (unsigned char C : S) {
@@ -593,106 +590,129 @@ static void quote(llvm::raw_ostream &OS, llvm::StringRef S) {
   OS << '\"';
 }
 
-enum IndenterAction {
-  Indent,
-  Outdent,
-  Newline,
-  Space,
-};
-
-// Prints JSON. The indenter can be used to control formatting.
-template <typename Indenter>
-void llvm::json::Value::print(raw_ostream &OS, const Indenter &I) const {
-  switch (Type) {
-  case T_Null:
+void llvm::json::OStream::value(const Value &V) {
+  switch (V.kind()) {
+  case Value::Null:
+    valueBegin();
     OS << "null";
-    break;
-  case T_Boolean:
-    OS << (as<bool>() ? "true" : "false");
-    break;
-  case T_Double:
-    OS << format("%.*g", std::numeric_limits<double>::max_digits10,
-                 as<double>());
-    break;
-  case T_Integer:
-    OS << as<int64_t>();
-    break;
-  case T_StringRef:
-    quote(OS, as<StringRef>());
-    break;
-  case T_String:
-    quote(OS, as<std::string>());
-    break;
-  case T_Object: {
-    bool Comma = false;
-    OS << '{';
-    I(Indent);
-    for (const auto *P : sortedElements(as<json::Object>())) {
-      if (Comma)
-        OS << ',';
-      Comma = true;
-      I(Newline);
-      quote(OS, P->first);
-      OS << ':';
-      I(Space);
-      P->second.print(OS, I);
-    }
-    I(Outdent);
-    if (Comma)
-      I(Newline);
-    OS << '}';
-    break;
+    return;
+  case Value::Boolean:
+    valueBegin();
+    OS << (*V.getAsBoolean() ? "true" : "false");
+    return;
+  case Value::Number:
+    valueBegin();
+    if (V.Type == Value::T_Integer)
+      OS << *V.getAsInteger();
+    else
+      OS << format("%.*g", std::numeric_limits<double>::max_digits10,
+                   *V.getAsNumber());
+    return;
+  case Value::String:
+    valueBegin();
+    quote(OS, *V.getAsString());
+    return;
+  case Value::Array:
+    return array([&] {
+      for (const Value &E : *V.getAsArray())
+        value(E);
+    });
+  case Value::Object:
+    return object([&] {
+      for (const Object::value_type *E : sortedElements(*V.getAsObject()))
+        attribute(E->first, E->second);
+    });
   }
-  case T_Array: {
-    bool Comma = false;
-    OS << '[';
-    I(Indent);
-    for (const auto &E : as<json::Array>()) {
-      if (Comma)
-        OS << ',';
-      Comma = true;
-      I(Newline);
-      E.print(OS, I);
-    }
-    I(Outdent);
-    if (Comma)
-      I(Newline);
-    OS << ']';
-    break;
+}
+
+void llvm::json::OStream::valueBegin() {
+  assert(Stack.back().Ctx != Object && "Only attributes allowed here");
+  if (Stack.back().HasValue) {
+    assert(Stack.back().Ctx != Singleton && "Only one value allowed here");
+    OS << ',';
+  }
+  if (Stack.back().Ctx == Array)
+    newline();
+  Stack.back().HasValue = true;
+}
+
+void llvm::json::OStream::newline() {
+  if (IndentSize) {
+    OS.write('\n');
+    OS.indent(Indent);
   }
+}
+
+void llvm::json::OStream::arrayBegin() {
+  valueBegin();
+  Stack.emplace_back();
+  Stack.back().Ctx = Array;
+  Indent += IndentSize;
+  OS << '[';
+}
+
+void llvm::json::OStream::arrayEnd() {
+  assert(Stack.back().Ctx == Array);
+  Indent -= IndentSize;
+  if (Stack.back().HasValue)
+    newline();
+  OS << ']';
+  Stack.pop_back();
+  assert(!Stack.empty());
+}
+
+void llvm::json::OStream::objectBegin() {
+  valueBegin();
+  Stack.emplace_back();
+  Stack.back().Ctx = Object;
+  Indent += IndentSize;
+  OS << '{';
+}
+
+void llvm::json::OStream::objectEnd() {
+  assert(Stack.back().Ctx == Object);
+  Indent -= IndentSize;
+  if (Stack.back().HasValue)
+    newline();
+  OS << '}';
+  Stack.pop_back();
+  assert(!Stack.empty());
+}
+
+void llvm::json::OStream::attributeBegin(llvm::StringRef Key) {
+  assert(Stack.back().Ctx == Object);
+  if (Stack.back().HasValue)
+    OS << ',';
+  newline();
+  Stack.back().HasValue = true;
+  Stack.emplace_back();
+  Stack.back().Ctx = Singleton;
+  if (LLVM_LIKELY(isUTF8(Key))) {
+    quote(OS, Key);
+  } else {
+    assert(false && "Invalid UTF-8 in attribute key");
+    quote(OS, fixUTF8(Key));
   }
+  OS.write(':');
+  if (IndentSize)
+    OS.write(' ');
+}
+
+void llvm::json::OStream::attributeEnd() {
+  assert(Stack.back().Ctx == Singleton);
+  assert(Stack.back().HasValue && "Attribute must have a value");
+  Stack.pop_back();
+  assert(Stack.back().Ctx == Object);
 }
 
+} // namespace json
+} // namespace llvm
+
 void llvm::format_provider<llvm::json::Value>::format(
     const llvm::json::Value &E, raw_ostream &OS, StringRef Options) {
-  if (Options.empty()) {
-    OS << E;
-    return;
-  }
   unsigned IndentAmount = 0;
-  if (Options.getAsInteger(/*Radix=*/10, IndentAmount))
+  if (!Options.empty() && Options.getAsInteger(/*Radix=*/10, IndentAmount))
     llvm_unreachable("json::Value format options should be an integer");
-  unsigned IndentLevel = 0;
-  E.print(OS, [&](IndenterAction A) {
-    switch (A) {
-    case Newline:
-      OS << '\n';
-      OS.indent(IndentLevel);
-      break;
-    case Space:
-      OS << ' ';
-      break;
-    case Indent:
-      IndentLevel += IndentAmount;
-      break;
-    case Outdent:
-      IndentLevel -= IndentAmount;
-      break;
-    };
-  });
+  json::OStream(OS, IndentAmount).value(E);
 }
 
-llvm::raw_ostream &llvm::json::operator<<(raw_ostream &OS, const Value &E) {
-  E.print(OS, [](IndenterAction A) { /*ignore*/ });
-  return OS;
-}
diff --git a/llvm/lib/Support/TimeProfiler.cpp b/llvm/lib/Support/TimeProfiler.cpp
index 447ddf5ab3c..bc234081564 100644
--- a/llvm/lib/Support/TimeProfiler.cpp
+++ b/llvm/lib/Support/TimeProfiler.cpp
@@ -87,25 +87,24 @@ struct TimeTraceProfiler {
   void Write(raw_pwrite_stream &OS) {
     assert(Stack.empty() &&
            "All profiler sections should be ended when calling Write");
-
-    json::Array Events;
-    const size_t ExpectedEntryCount =
-        Entries.size() + CountAndTotalPerName.size() + 1;
-    Events.reserve(ExpectedEntryCount);
+    json::OStream J(OS);
+    J.objectBegin();
+    J.attributeBegin("traceEvents");
+    J.arrayBegin();
 
     // Emit all events for the main flame graph.
     for (const auto &E : Entries) {
       auto StartUs = duration_cast<microseconds>(E.Start - StartTime).count();
       auto DurUs = duration_cast<microseconds>(E.Duration).count();
 
-      Events.emplace_back(json::Object{
-          {"pid", 1},
-          {"tid", 0},
-          {"ph", "X"},
-          {"ts", StartUs},
-          {"dur", DurUs},
-          {"name", E.Name},
-          {"args", json::Object{{"detail", E.Detail}}},
+      J.object([&]{
+        J.attribute("pid", 1);
+        J.attribute("tid", 0);
+        J.attribute("ph", "X");
+        J.attribute("ts", StartUs);
+        J.attribute("dur", DurUs);
+        J.attribute("name", E.Name);
+        J.attributeObject("args", [&] { J.attribute("detail", E.Detail); });
       });
     }
 
@@ -126,36 +125,36 @@ struct TimeTraceProfiler {
       auto DurUs = duration_cast<microseconds>(E.second.second).count();
       auto Count = CountAndTotalPerName[E.first].first;
 
-      Events.emplace_back(json::Object{
-          {"pid", 1},
-          {"tid", Tid},
-          {"ph", "X"},
-          {"ts", 0},
-          {"dur", DurUs},
-          {"name", "Total " + E.first},
-          {"args", json::Object{{"count", static_cast<int64_t>(Count)},
-                                {"avg ms",
-                                 static_cast<int64_t>(DurUs / Count / 1000)}}},
+      J.object([&]{
+        J.attribute("pid", 1);
+        J.attribute("tid", Tid);
+        J.attribute("ph", "X");
+        J.attribute("ts", 0);
+        J.attribute("dur", DurUs);
+        J.attribute("name", "Total " + E.first);
+        J.attributeObject("args", [&] {
+          J.attribute("count", int64_t(Count));
+          J.attribute("avg ms", int64_t(DurUs / Count / 1000));
+        });
       });
 
       ++Tid;
     }
 
     // Emit metadata event with process name.
-    Events.emplace_back(json::Object{
-        {"cat", ""},
-        {"pid", 1},
-        {"tid", 0},
-        {"ts", 0},
-        {"ph", "M"},
-        {"name", "process_name"},
-        {"args", json::Object{{"name", "clang"}}},
+    J.object([&] {
+      J.attribute("cat", "");
+      J.attribute("pid", 1);
+      J.attribute("tid", 0);
+      J.attribute("ts", 0);
+      J.attribute("ph", "M");
+      J.attribute("name", "process_name");
+      J.attributeObject("args", [&] { J.attribute("name", "clang"); });
     });
 
-    assert(Events.size() == ExpectedEntryCount && "Size prediction failed!");
-
-    OS << formatv("{0:2}", json::Value(json::Object(
-                               {{"traceEvents", std::move(Events)}})));
+    J.arrayEnd();
+    J.attributeEnd();
+    J.objectEnd();
   }
 
   SmallVector<Entry, 16> Stack;