From 140b7b6f09ca01f89701ccc86ddd8553cf42c311 Mon Sep 17 00:00:00 2001
From: Sam McCall <sam.mccall@gmail.com>
Date: Wed, 23 Sep 2020 23:33:19 +0200
Subject: [PATCH] [JSON] Allow emitting comments in json::OStream

This isn't standard JSON, but is a popular extension.
It will be used to show errors in context, rendering pseudo-json for humans.

Differential Revision: https://reviews.llvm.org/D88103
---
 llvm/include/llvm/Support/JSON.h    |  7 ++++++
 llvm/lib/Support/JSON.cpp           | 36 +++++++++++++++++++++++++++++
 llvm/unittests/Support/JSONTest.cpp | 14 ++++++++---
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Support/JSON.h b/llvm/include/llvm/Support/JSON.h
index 8b1c66234fe8..80fc1ee238e7 100644
--- a/llvm/include/llvm/Support/JSON.h
+++ b/llvm/include/llvm/Support/JSON.h
@@ -706,6 +706,7 @@ public:
 /// json::OStream allows writing well-formed JSON without materializing
 /// all structures as json::Value ahead of time.
 /// It's faster, lower-level, and less safe than OS << json::Value.
+/// It also allows emitting more constructs, such as comments.
 ///
 /// Only one "top-level" object can be written to a stream.
 /// Simplest usage involves passing lambdas (Blocks) to fill in containers:
@@ -791,6 +792,10 @@ class OStream {
     Contents();
     objectEnd();
   }
+  /// Emit a JavaScript comment associated with the next printed value.
+  /// The string must be valid until the next attribute or value is emitted.
+  /// Comments are not part of standard JSON, and many parsers reject them!
+  void comment(llvm::StringRef);
 
   // High level functions to output object attributes.
   // Valid only within an object (any number of times).
@@ -826,6 +831,7 @@ class OStream {
   }
 
   void valueBegin();
+  void flushComment();
   void newline();
 
   enum Context {
@@ -838,6 +844,7 @@ class OStream {
     bool HasValue = false;
   };
   llvm::SmallVector<State, 16> Stack; // Never empty.
+  llvm::StringRef PendingComment;
   llvm::raw_ostream &OS;
   unsigned IndentSize;
   unsigned Indent = 0;
diff --git a/llvm/lib/Support/JSON.cpp b/llvm/lib/Support/JSON.cpp
index 16b1d11efd08..db4121cf82bc 100644
--- a/llvm/lib/Support/JSON.cpp
+++ b/llvm/lib/Support/JSON.cpp
@@ -9,6 +9,7 @@
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Format.h"
+#include "llvm/Support/raw_ostream.h"
 #include <cctype>
 
 namespace llvm {
@@ -633,9 +634,40 @@ void llvm::json::OStream::valueBegin() {
   }
   if (Stack.back().Ctx == Array)
     newline();
+  flushComment();
   Stack.back().HasValue = true;
 }
 
+void OStream::comment(llvm::StringRef Comment) {
+  assert(PendingComment.empty() && "Only one comment per value!");
+  PendingComment = Comment;
+}
+
+void OStream::flushComment() {
+  if (PendingComment.empty())
+    return;
+  OS << (IndentSize ? "/* " : "/*");
+  // Be sure not to accidentally emit "*/". Transform to "* /".
+  while (!PendingComment.empty()) {
+    auto Pos = PendingComment.find("*/");
+    if (Pos == StringRef::npos) {
+      OS << PendingComment;
+      PendingComment = "";
+    } else {
+      OS << PendingComment.take_front(Pos) << "* /";
+      PendingComment = PendingComment.drop_front(Pos + 2);
+    }
+  }
+  OS << (IndentSize ? " */" : "*/");
+  // Comments are on their own line unless attached to an attribute value.
+  if (Stack.size() > 1 && Stack.back().Ctx == Singleton) {
+    if (IndentSize)
+      OS << ' ';
+  } else {
+    newline();
+  }
+}
+
 void llvm::json::OStream::newline() {
   if (IndentSize) {
     OS.write('\n');
@@ -657,6 +689,7 @@ void llvm::json::OStream::arrayEnd() {
   if (Stack.back().HasValue)
     newline();
   OS << ']';
+  assert(PendingComment.empty());
   Stack.pop_back();
   assert(!Stack.empty());
 }
@@ -675,6 +708,7 @@ void llvm::json::OStream::objectEnd() {
   if (Stack.back().HasValue)
     newline();
   OS << '}';
+  assert(PendingComment.empty());
   Stack.pop_back();
   assert(!Stack.empty());
 }
@@ -684,6 +718,7 @@ void llvm::json::OStream::attributeBegin(llvm::StringRef Key) {
   if (Stack.back().HasValue)
     OS << ',';
   newline();
+  flushComment();
   Stack.back().HasValue = true;
   Stack.emplace_back();
   Stack.back().Ctx = Singleton;
@@ -701,6 +736,7 @@ void llvm::json::OStream::attributeBegin(llvm::StringRef Key) {
 void llvm::json::OStream::attributeEnd() {
   assert(Stack.back().Ctx == Singleton);
   assert(Stack.back().HasValue && "Attribute must have a value");
+  assert(PendingComment.empty());
   Stack.pop_back();
   assert(Stack.back().Ctx == Object);
 }
diff --git a/llvm/unittests/Support/JSONTest.cpp b/llvm/unittests/Support/JSONTest.cpp
index 73fc626af8cb..986cf5e73a37 100644
--- a/llvm/unittests/Support/JSONTest.cpp
+++ b/llvm/unittests/Support/JSONTest.cpp
@@ -420,15 +420,19 @@ TEST(JSONTest, Stream) {
     std::string S;
     llvm::raw_string_ostream OS(S);
     OStream J(OS, Indent);
+    J.comment("top*/level");
     J.object([&] {
       J.attributeArray("foo", [&] {
         J.value(nullptr);
+        J.comment("element");
         J.value(42.5);
         J.arrayBegin();
         J.value(43);
         J.arrayEnd();
       });
+      J.comment("attribute");
       J.attributeBegin("bar");
+      J.comment("attribute value");
       J.objectBegin();
       J.objectEnd();
       J.attributeEnd();
@@ -437,17 +441,21 @@ TEST(JSONTest, Stream) {
     return OS.str();
   };
 
-  const char *Plain = R"({"foo":[null,42.5,[43]],"bar":{},"baz":"xyz"})";
+  const char *Plain =
+      R"(/*top* /level*/{"foo":[null,/*element*/42.5,[43]],/*attribute*/"bar":/*attribute value*/{},"baz":"xyz"})";
   EXPECT_EQ(Plain, StreamStuff(0));
-  const char *Pretty = R"({
+  const char *Pretty = R"(/* top* /level */
+{
   "foo": [
     null,
+    /* element */
     42.5,
     [
       43
     ]
   ],
-  "bar": {},
+  /* attribute */
+  "bar": /* attribute value */ {},
   "baz": "xyz"
 })";
   EXPECT_EQ(Pretty, StreamStuff(2));