Modernize to range-based loop

[oota-llvm.git] / utils / yaml-bench / YAMLBench.cpp

diff --git a/utils/yaml-bench/YAMLBench.cpp b/utils/yaml-bench/YAMLBench.cpp
index eef4a725a1a3a4ba26e296f21ca1255957167927..634622a710c80551f91f5ec11f8393101b25ca48 100644 (file)
--- a/utils/yaml-bench/YAMLBench.cpp
+++ b/utils/yaml-bench/YAMLBench.cpp
@@ -7,7 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This program executes the YAMLParser on differntly sized YAML texts and
+// This program executes the YAMLParser on differently sized YAML texts and
 // outputs the run time.
 //
 //===----------------------------------------------------------------------===//
@@ -19,9 +19,10 @@
 #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 "llvm/Support/system_error.h"
+#include <system_error>
 
 using namespace llvm;
 
@@ -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) {}
@@ -63,6 +68,20 @@ static raw_ostream &operator <<(raw_ostream &os, const indent &in) {
   return os;
 }
 
+/// \brief Pretty print a tag by replacing tag:yaml.org,2002: with !!.
+static std::string prettyTag(yaml::Node *N) {
+  std::string Tag = N->getVerbatimTag();
+  if (StringRef(Tag).startswith("tag:yaml.org,2002:")) {
+    std::string Ret = "!!";
+    Ret += StringRef(Tag).substr(18);
+    return Ret;
+  }
+  std::string Ret = "!<";
+  Ret += Tag;
+  Ret += ">";
+  return Ret;
+}
+
 static void dumpNode( yaml::Node *n
                     , unsigned Indent = 0
                     , bool SuppressFirstIndent = false) {
@@ -76,9 +95,11 @@ static void dumpNode( yaml::Node *n
   if (yaml::ScalarNode *sn = dyn_cast<yaml::ScalarNode>(n)) {
     SmallString<32> Storage;
     StringRef Val = sn->getValue(Storage);
-    outs() << "!!str \"" << yaml::escape(Val) << "\"";
+    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() << "!!seq [\n";
+    outs() << prettyTag(n) << " [\n";
     ++Indent;
     for (yaml::SequenceNode::iterator i = sn->begin(), e = sn->end();
                                       i != e; ++i) {
@@ -88,7 +109,7 @@ static void dumpNode( yaml::Node *n
     --Indent;
     outs() << indent(Indent) << "]";
   } else if (yaml::MappingNode *mn = dyn_cast<yaml::MappingNode>(n)) {
-    outs() << "!!map {\n";
+    outs() << prettyTag(n) << " {\n";
     ++Indent;
     for (yaml::MappingNode::iterator i = mn->begin(), e = mn->end();
                                      i != e; ++i) {
@@ -103,8 +124,8 @@ 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)) {
-    outs() << "!!null null";
+  } else if (isa<yaml::NullNode>(n)) {
+    outs() << prettyTag(n) << " null";
   }
 }
 
@@ -173,19 +194,26 @@ 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()) {
-    OwningPtr<MemoryBuffer> Buf;
-    if (MemoryBuffer::getFileOrSTDIN(Input, Buf))
+    ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
+        MemoryBuffer::getFileOrSTDIN(Input);
+    if (!BufOrErr)
       return 1;
+    MemoryBuffer &Buf = *BufOrErr.get();
 
     llvm::SourceMgr sm;
     if (DumpTokens) {
-      yaml::dumpTokens(Buf->getBuffer(), outs());
+      yaml::dumpTokens(Buf.getBuffer(), outs());
     }
 
     if (DumpCanonical) {
-      yaml::Stream stream(Buf->getBuffer(), sm);
+      yaml::Stream stream(Buf.getBuffer(), sm, ShowColors);
       dumpStream(stream);
+      if (stream.failed())
+        return 1;
     }
   }