1 files changed, 15 insertions, 3 deletions
diff --git a/utils/yaml-bench/YAMLBench.cpp b/utils/yaml-bench/YAMLBench.cpp
index 8bd1ea17ea8e2..634622a710c80 100644
--- a/utils/yaml-bench/YAMLBench.cpp
+++ b/utils/yaml-bench/YAMLBench.cpp
@@ -19,6 +19,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/Timer.h"
+#include "llvm/Support/Process.h"
 #include "llvm/Support/YAMLParser.h"
 #include "llvm/Support/raw_ostream.h"
 #include <system_error>
@@ -52,6 +53,10 @@ static cl::opt<unsigned>
                   "Do not use more megabytes of memory"),
                 cl::init(1000));
 
+cl::opt<cl::boolOrDefault>
+    UseColor("use-color", cl::desc("Emit colored output (default=autodetect)"),
+             cl::init(cl::BOU_UNSET));
+
 struct indent {
   unsigned distance;
   indent(unsigned d) : distance(d) {}
@@ -69,7 +74,7 @@ static std::string prettyTag(yaml::Node *N) {
   if (StringRef(Tag).startswith("tag:yaml.org,2002:")) {
     std::string Ret = "!!";
     Ret += StringRef(Tag).substr(18);
-    return std::move(Ret);
+    return Ret;
   }
   std::string Ret = "!<";
   Ret += Tag;
@@ -91,6 +96,8 @@ static void dumpNode( yaml::Node *n
     SmallString<32> Storage;
     StringRef Val = sn->getValue(Storage);
     outs() << prettyTag(n) << " \"" << yaml::escape(Val) << "\"";
+  } else if (yaml::BlockScalarNode *BN = dyn_cast<yaml::BlockScalarNode>(n)) {
+    outs() << prettyTag(n) << " \"" << yaml::escape(BN->getValue()) << "\"";
   } else if (yaml::SequenceNode *sn = dyn_cast<yaml::SequenceNode>(n)) {
     outs() << prettyTag(n) << " [\n";
     ++Indent;
@@ -117,7 +124,7 @@ static void dumpNode( yaml::Node *n
     outs() << indent(Indent) << "}";
   } else if (yaml::AliasNode *an = dyn_cast<yaml::AliasNode>(n)){
     outs() << "*" << an->getName();
-  } else if (dyn_cast<yaml::NullNode>(n)) {
+  } else if (isa<yaml::NullNode>(n)) {
     outs() << prettyTag(n) << " null";
   }
 }
@@ -187,6 +194,9 @@ static std::string createJSONText(size_t MemoryMB, unsigned ValueSize) {
 
 int main(int argc, char **argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv);
+  bool ShowColors = UseColor == cl::BOU_UNSET
+                        ? sys::Process::StandardOutHasColors()
+                        : UseColor == cl::BOU_TRUE;
   if (Input.getNumOccurrences()) {
     ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
         MemoryBuffer::getFileOrSTDIN(Input);
@@ -200,8 +210,10 @@ int main(int argc, char **argv) {
     }
 
     if (DumpCanonical) {
-      yaml::Stream stream(Buf.getBuffer(), sm);
+      yaml::Stream stream(Buf.getBuffer(), sm, ShowColors);
       dumpStream(stream);
+      if (stream.failed())
+        return 1;
     }
   }