Re-submit r367649: Improve raw_ostream so that you can "write" colors using operator<<

The original patch broke buildbots, perhaps because it changed the
default setting whether colors are enabled or not.

llvm-svn: 368131

13 files changed, 130 insertions, 92 deletions

diff --git a/clang/tools/diagtool/TreeView.cpp b/clang/tools/diagtool/TreeView.cpp
index 154c52a485a..5160afc64ab 100644
--- a/clang/tools/diagtool/TreeView.cpp
+++ b/clang/tools/diagtool/TreeView.cpp
@@ -37,7 +37,8 @@ public:
 
   void setColor(llvm::raw_ostream::Colors Color) {
     if (ShowColors)
-      out << llvm::sys::Process::OutputColor(Color, false, false);
+      out << llvm::sys::Process::OutputColor(static_cast<char>(Color), false,
+                                             false);
   }
 
   void resetColor() {
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 9fb2c51ff68..601c2e8de89 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -62,16 +62,16 @@ LinkerDriver *driver;
 bool link(ArrayRef<const char *> args, bool canExitEarly, raw_ostream &diag) {
   errorHandler().logName = args::getFilenameWithoutExe(args[0]);
   errorHandler().errorOS = &diag;
-  errorHandler().colorDiagnostics = diag.has_colors();
   errorHandler().errorLimitExceededMsg =
       "too many errors emitted, stopping now"
       " (use /errorlimit:0 to see all errors)";
   errorHandler().exitEarly = canExitEarly;
-  config = make<Configuration>();
+  enableColors(diag.has_colors());
 
+  config = make<Configuration>();
   symtab = make<SymbolTable>();
-
   driver = make<LinkerDriver>();
+
   driver->link(args);
 
   // Call exit() if we can to avoid calling destructors.
diff --git a/lld/COFF/DriverUtils.cpp b/lld/COFF/DriverUtils.cpp
index 6167520e589..49aff059a09 100644
--- a/lld/COFF/DriverUtils.cpp
+++ b/lld/COFF/DriverUtils.cpp
@@ -757,15 +757,15 @@ static void handleColorDiagnostics(opt::InputArgList &args) {
   if (!arg)
     return;
   if (arg->getOption().getID() == OPT_color_diagnostics) {
-    errorHandler().colorDiagnostics = true;
+    enableColors(true);
   } else if (arg->getOption().getID() == OPT_no_color_diagnostics) {
-    errorHandler().colorDiagnostics = false;
+    enableColors(false);
   } else {
     StringRef s = arg->getValue();
     if (s == "always")
-      errorHandler().colorDiagnostics = true;
+      enableColors(true);
     else if (s == "never")
-      errorHandler().colorDiagnostics = false;
+      enableColors(false);
     else if (s != "auto")
       error("unknown option: --color-diagnostics=" + s);
   }
diff --git a/lld/Common/ErrorHandler.cpp b/lld/Common/ErrorHandler.cpp
index 4c6a5b382eb..5934a10af58 100644
--- a/lld/Common/ErrorHandler.cpp
+++ b/lld/Common/ErrorHandler.cpp
@@ -46,6 +46,10 @@ ErrorHandler &lld::errorHandler() {
   return handler;
 }
 
+void lld::enableColors(bool enable) {
+  errorHandler().errorOS->enable_colors(enable);
+}
+
 void lld::exitLld(int val) {
   // Delete any temporary file, while keeping the memory mapping open.
   if (errorHandler().outputBuffer)
@@ -106,54 +110,41 @@ void lld::checkError(Error e) {
 //
 // This function returns an error location string. An error location is
 // extracted from an error message using regexps.
-static std::string getLocation(std::string msg, std::string defaultMsg) {
-  static std::vector<std::regex> Regexes{
-      std::regex(R"(^undefined (?:\S+ )?symbol:.*\n>>> referenced by (\S+):(\d+)\n.*)"),
+std::string ErrorHandler::getLocation(const Twine &msg) {
+  if (!vsDiagnostics)
+    return logName;
+
+  static std::regex regexes[] = {
+      std::regex(
+          R"(^undefined (?:\S+ )?symbol:.*\n>>> referenced by (\S+):(\d+)\n.*)"),
       std::regex(R"(^undefined symbol:.*\n>>> referenced by (.*):)"),
       std::regex(
           R"(^duplicate symbol: .*\n>>> defined in (\S+)\n>>> defined in.*)"),
-      std::regex(
-          R"(^duplicate symbol: .*\n>>> defined at (\S+):(\d+).*)"),
-      std::regex(
-          R"(.*\n>>> defined in .*\n>>> referenced by (\S+):(\d+))"),
+      std::regex(R"(^duplicate symbol: .*\n>>> defined at (\S+):(\d+).*)"),
+      std::regex(R"(.*\n>>> defined in .*\n>>> referenced by (\S+):(\d+))"),
       std::regex(R"((\S+):(\d+): unclosed quote)"),
   };
 
-  std::smatch Match;
-  for (std::regex &Re : Regexes) {
-    if (std::regex_search(msg, Match, Re)) {
-      return Match.size() > 2 ? Match.str(1) + "(" + Match.str(2) + ")"
-                              : Match.str(1);
-    }
-  }
-  return defaultMsg;
-}
+  std::string str = msg.str();
+  for (std::regex &re : regexes) {
+    std::smatch m;
+    if (!std::regex_search(str, m, re))
+      continue;
 
-void ErrorHandler::printHeader(StringRef s, raw_ostream::Colors c,
-                               const Twine &msg) {
-
-  if (vsDiagnostics) {
-    // A Visual Studio-style error message starts with an error location.
-    // If a location cannot be extracted then we default to LogName.
-    *errorOS << getLocation(msg.str(), logName) << ": ";
-  } else {
-    *errorOS << logName << ": ";
+    assert(m.size() == 2 || m.size() == 3);
+    if (m.size() == 2)
+      return m.str(1);
+    return m.str(1) + "(" + m.str(2) + ")";
   }
 
-  if (colorDiagnostics) {
-    errorOS->changeColor(c, true);
-    *errorOS << s;
-    errorOS->resetColor();
-  } else {
-    *errorOS << s;
-  }
+  return logName;
 }
 
 void ErrorHandler::log(const Twine &msg) {
-  if (verbose) {
-    std::lock_guard<std::mutex> lock(mu);
-    *errorOS << logName << ": " << msg << "\n";
-  }
+  if (!verbose)
+    return;
+  std::lock_guard<std::mutex> lock(mu);
+  *errorOS << logName << ": " << msg << "\n";
 }
 
 void ErrorHandler::message(const Twine &msg) {
@@ -170,42 +161,37 @@ void ErrorHandler::warn(const Twine &msg) {
 
   std::lock_guard<std::mutex> lock(mu);
   newline(errorOS, msg);
-  printHeader("warning: ", raw_ostream::MAGENTA, msg);
-  *errorOS << msg << "\n";
+  *errorOS << getLocation(msg) << ": " << Colors::MAGENTA
+           << "warning: " << Colors::RESET << msg << "\n";
 }
 
-void ErrorHandler::printErrorMsg(const Twine &msg) {
-  newline(errorOS, msg);
-  printHeader("error: ", raw_ostream::RED, msg);
-  *errorOS << msg << "\n";
-}
-
-void ErrorHandler::printError(const Twine &msg) {
+void ErrorHandler::error(const Twine &msg) {
+  // If Visual Studio-style error message mode is enabled,
+  // this particular error is printed out as two errors.
   if (vsDiagnostics) {
-    static std::regex reDuplicateSymbol(
-        R"(^(duplicate symbol: .*))"
-        R"((\n>>> defined at \S+:\d+\n>>>.*))"
-        R"((\n>>> defined at \S+:\d+\n>>>.*))");
-    std::string msgStr = msg.str();
-    std::smatch match;
-    if (std::regex_match(msgStr, match, reDuplicateSymbol)) {
-      printErrorMsg(match.str(1) + match.str(2));
-      printErrorMsg(match.str(1) + match.str(3));
+    static std::regex re(R"(^(duplicate symbol: .*))"
+                         R"((\n>>> defined at \S+:\d+\n>>>.*))"
+                         R"((\n>>> defined at \S+:\d+\n>>>.*))");
+    std::string str = msg.str();
+    std::smatch m;
+
+    if (std::regex_match(str, m, re)) {
+      error(m.str(1) + m.str(2));
+      error(m.str(1) + m.str(3));
       return;
     }
   }
-  printErrorMsg(msg);
-}
 
-void ErrorHandler::error(const Twine &msg) {
   std::lock_guard<std::mutex> lock(mu);
 
   if (errorLimit == 0 || errorCount < errorLimit) {
-    printError(msg);
+    newline(errorOS, msg);
+    *errorOS << getLocation(msg) << ": " << Colors::RED
+             << "error: " << Colors::RESET << msg << "\n";
   } else if (errorCount == errorLimit) {
     newline(errorOS, msg);
-    printHeader("error: ", raw_ostream::RED, msg);
-    *errorOS << errorLimitExceededMsg << "\n";
+    *errorOS << getLocation(msg) << ": " << Colors::RED
+             << "error: " << Colors::RESET << errorLimitExceededMsg << "\n";
     if (exitEarly)
       exitLld(1);
   }
diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index 0f0b567a7c2..9bcc91b6c8b 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -82,7 +82,7 @@ bool elf::link(ArrayRef<const char *> args, bool canExitEarly,
       "-error-limit=0 to see all errors)";
   errorHandler().errorOS = &error;
   errorHandler().exitEarly = canExitEarly;
-  errorHandler().colorDiagnostics = error.has_colors();
+  enableColors(error.has_colors());
 
   inputSections.clear();
   outputSections.clear();
diff --git a/lld/ELF/DriverUtils.cpp b/lld/ELF/DriverUtils.cpp
index 87f0aa2a982..5ed37e0171b 100644
--- a/lld/ELF/DriverUtils.cpp
+++ b/lld/ELF/DriverUtils.cpp
@@ -59,15 +59,15 @@ static void handleColorDiagnostics(opt::InputArgList &args) {
   if (!arg)
     return;
   if (arg->getOption().getID() == OPT_color_diagnostics) {
-    errorHandler().colorDiagnostics = true;
+    enableColors(true);
   } else if (arg->getOption().getID() == OPT_no_color_diagnostics) {
-    errorHandler().colorDiagnostics = false;
+    enableColors(false);
   } else {
     StringRef s = arg->getValue();
     if (s == "always")
-      errorHandler().colorDiagnostics = true;
+      enableColors(true);
     else if (s == "never")
-      errorHandler().colorDiagnostics = false;
+      enableColors(false);
     else if (s != "auto")
       error("unknown option: --color-diagnostics=" + s);
   }
diff --git a/lld/include/lld/Common/ErrorHandler.h b/lld/include/lld/Common/ErrorHandler.h
index 52e22001f05..5086fa9862d 100644
--- a/lld/include/lld/Common/ErrorHandler.h
+++ b/lld/include/lld/Common/ErrorHandler.h
@@ -87,7 +87,6 @@ public:
   StringRef errorLimitExceededMsg = "too many errors emitted, stopping now";
   StringRef logName = "lld";
   llvm::raw_ostream *errorOS = &llvm::errs();
-  bool colorDiagnostics = llvm::errs().has_colors();
   bool exitEarly = true;
   bool fatalWarnings = false;
   bool verbose = false;
@@ -102,14 +101,16 @@ public:
   std::unique_ptr<llvm::FileOutputBuffer> outputBuffer;
 
 private:
-  void printHeader(StringRef s, raw_ostream::Colors c, const Twine &msg);
-  void printErrorMsg(const Twine &msg);
-  void printError(const Twine &msg);
+  using Colors = raw_ostream::Colors;
+
+  std::string getLocation(const Twine &msg);
 };
 
 /// Returns the default error handler.
 ErrorHandler &errorHandler();
 
+void enableColors(bool enable);
+
 inline void error(const Twine &msg) { errorHandler().error(msg); }
 inline LLVM_ATTRIBUTE_NORETURN void fatal(const Twine &msg) {
   errorHandler().fatal(msg);
diff --git a/lld/lib/Driver/DarwinLdDriver.cpp b/lld/lib/Driver/DarwinLdDriver.cpp
index 8646d86c08e..68dc789479f 100644
--- a/lld/lib/Driver/DarwinLdDriver.cpp
+++ b/lld/lib/Driver/DarwinLdDriver.cpp
@@ -1151,7 +1151,7 @@ bool link(llvm::ArrayRef<const char *> args, bool CanExitEarly,
       "'-error-limit 0' to see all errors)";
   errorHandler().errorOS = &Error;
   errorHandler().exitEarly = CanExitEarly;
-  errorHandler().colorDiagnostics = Error.has_colors();
+  enableColors(Error.has_colors());
 
   MachOLinkingContext ctx;
   if (!parse(args, ctx))
diff --git a/lld/test/COFF/color-diagnostics.test b/lld/test/COFF/color-diagnostics.test
index 11cf8aff7a0..210965d5267 100644
--- a/lld/test/COFF/color-diagnostics.test
+++ b/lld/test/COFF/color-diagnostics.test
@@ -6,8 +6,8 @@
 # RUN: not lld-link -xyz --color-diagnostics=always /nosuchfile 2>&1 \
 # RUN:   | FileCheck -check-prefix=COLOR %s
 
-# COLOR: {{lld-link: .\[0;1;35mwarning: .\[0mignoring unknown argument '-xyz'}}
-# COLOR: {{lld-link: .\[0;1;31merror: .\[0mcould not open '/nosuchfile'}}
+# COLOR: {{lld-link: .\[0;35mwarning: .\[0mignoring unknown argument '-xyz'}}
+# COLOR: {{lld-link: .\[0;31merror: .\[0mcould not open '/nosuchfile'}}
 
 # RUN: not lld-link /nosuchfile 2>&1 | FileCheck -check-prefix=NOCOLOR %s
 # RUN: not lld-link -color-diagnostics=never /nosuchfile 2>&1 \
diff --git a/lld/test/ELF/color-diagnostics.test b/lld/test/ELF/color-diagnostics.test
index 8613074dc56..7f1e46c1349 100644
--- a/lld/test/ELF/color-diagnostics.test
+++ b/lld/test/ELF/color-diagnostics.test
@@ -6,8 +6,8 @@
 # RUN: not ld.lld -xyz -color-diagnostics=always /nosuchfile 2>&1 \
 # RUN:   | FileCheck -check-prefix=COLOR %s
 
-# COLOR: {{ld.lld: .\[0;1;31merror: .\[0munknown argument '-xyz'}}
-# COLOR: {{ld.lld: .\[0;1;31merror: .\[0mcannot open /nosuchfile}}
+# COLOR: {{ld.lld: .\[0;31merror: .\[0munknown argument '-xyz'}}
+# COLOR: {{ld.lld: .\[0;31merror: .\[0mcannot open /nosuchfile}}
 
 # RUN: not ld.lld -color-diagnostics=foobar 2>&1 | FileCheck -check-prefix=ERR %s
 # ERR: unknown option: --color-diagnostics=foobar
diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index a1d22bf662c..069d16dea5a 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -83,10 +83,10 @@ bool lld::wasm::link(ArrayRef<const char *> args, bool canExitEarly,
                      raw_ostream &error) {
   errorHandler().logName = args::getFilenameWithoutExe(args[0]);
   errorHandler().errorOS = &error;
-  errorHandler().colorDiagnostics = error.has_colors();
   errorHandler().errorLimitExceededMsg =
       "too many errors emitted, stopping now (use "
       "-error-limit=0 to see all errors)";
+  enableColors(error.has_colors());
 
   config = make<Configuration>();
   symtab = make<SymbolTable>();
@@ -134,15 +134,15 @@ static void handleColorDiagnostics(opt::InputArgList &args) {
   if (!arg)
     return;
   if (arg->getOption().getID() == OPT_color_diagnostics) {
-    errorHandler().colorDiagnostics = true;
+    enableColors(true);
   } else if (arg->getOption().getID() == OPT_no_color_diagnostics) {
-    errorHandler().colorDiagnostics = false;
+    enableColors(false);
   } else {
     StringRef s = arg->getValue();
     if (s == "always")
-      errorHandler().colorDiagnostics = true;
+      enableColors(true);
     else if (s == "never")
-      errorHandler().colorDiagnostics = false;
+      enableColors(false);
     else if (s != "auto")
       error("unknown option: --color-diagnostics=" + s);
   }
diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 48bb623b063..4a6b51ac609 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -72,7 +72,7 @@ private:
 
 public:
   // color order matches ANSI escape sequence, don't change
-  enum Colors {
+  enum class Colors {
     BLACK = 0,
     RED,
     GREEN,
@@ -81,9 +81,21 @@ public:
     MAGENTA,
     CYAN,
     WHITE,
-    SAVEDCOLOR
+    SAVEDCOLOR,
+    RESET,
   };
 
+  static const Colors BLACK = Colors::BLACK;
+  static const Colors RED = Colors::RED;
+  static const Colors GREEN = Colors::GREEN;
+  static const Colors YELLOW = Colors::YELLOW;
+  static const Colors BLUE = Colors::BLUE;
+  static const Colors MAGENTA = Colors::MAGENTA;
+  static const Colors CYAN = Colors::CYAN;
+  static const Colors WHITE = Colors::WHITE;
+  static const Colors SAVEDCOLOR = Colors::SAVEDCOLOR;
+  static const Colors RESET = Colors::RESET;
+
   explicit raw_ostream(bool unbuffered = false)
       : BufferMode(unbuffered ? Unbuffered : InternalBuffer) {
     // Start out ready to flush.
@@ -214,6 +226,9 @@ public:
   /// Output \p N in hexadecimal, without any prefix or padding.
   raw_ostream &write_hex(unsigned long long N);
 
+  // Change the foreground color of text.
+  raw_ostream &operator<<(Colors C);
+
   /// Output a formatted UUID with dash separators.
   using uuid_t = uint8_t[16];
   raw_ostream &write_uuid(const uuid_t UUID);
@@ -277,6 +292,10 @@ public:
   /// This function determines if this stream is displayed and supports colors.
   virtual bool has_colors() const { return is_displayed(); }
 
+  // Enable or disable colors. Once disable_colors() is called,
+  // changeColor() has no effect until enable_colors() is called.
+  virtual void enable_colors(bool enable) {}
+
   //===--------------------------------------------------------------------===//
   // Subclass Interface
   //===--------------------------------------------------------------------===//
@@ -365,8 +384,8 @@ public:
 class raw_fd_ostream : public raw_pwrite_stream {
   int FD;
   bool ShouldClose;
-
   bool SupportsSeeking;
+  bool ColorEnabled = true;
 
 #ifdef _WIN32
   /// True if this fd refers to a Windows console device. Mintty and other
@@ -442,6 +461,8 @@ public:
 
   bool has_colors() const override;
 
+  void enable_colors(bool enable) override { ColorEnabled = enable; }
+
   std::error_code error() const { return EC; }
 
   /// Return the value of the flag in this raw_fd_ostream indicating whether an
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 16e56ed7835..b9989371f5e 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -65,6 +65,17 @@
 
 using namespace llvm;
 
+const raw_ostream::Colors raw_ostream::BLACK;
+const raw_ostream::Colors raw_ostream::RED;
+const raw_ostream::Colors raw_ostream::GREEN;
+const raw_ostream::Colors raw_ostream::YELLOW;
+const raw_ostream::Colors raw_ostream::BLUE;
+const raw_ostream::Colors raw_ostream::MAGENTA;
+const raw_ostream::Colors raw_ostream::CYAN;
+const raw_ostream::Colors raw_ostream::WHITE;
+const raw_ostream::Colors raw_ostream::SAVEDCOLOR;
+const raw_ostream::Colors raw_ostream::RESET;
+
 raw_ostream::~raw_ostream() {
   // raw_ostream's subclasses should take care to flush the buffer
   // in their destructors.
@@ -133,6 +144,14 @@ raw_ostream &raw_ostream::write_hex(unsigned long long N) {
   return *this;
 }
 
+raw_ostream &raw_ostream::operator<<(Colors C) {
+  if (C == Colors::RESET)
+    resetColor();
+  else
+    changeColor(C);
+  return *this;
+}
+
 raw_ostream &raw_ostream::write_uuid(const uuid_t UUID) {
   for (int Idx = 0; Idx < 16; ++Idx) {
     *this << format("%02" PRIX32, UUID[Idx]);
@@ -784,11 +803,15 @@ size_t raw_fd_ostream::preferred_buffer_size() const {
 
 raw_ostream &raw_fd_ostream::changeColor(enum Colors colors, bool bold,
                                          bool bg) {
+  if (!ColorEnabled)
+    return *this;
+
   if (sys::Process::ColorNeedsFlush())
     flush();
   const char *colorcode =
-    (colors == SAVEDCOLOR) ? sys::Process::OutputBold(bg)
-    : sys::Process::OutputColor(colors, bold, bg);
+      (colors == SAVEDCOLOR)
+          ? sys::Process::OutputBold(bg)
+          : sys::Process::OutputColor(static_cast<char>(colors), bold, bg);
   if (colorcode) {
     size_t len = strlen(colorcode);
     write(colorcode, len);
@@ -799,6 +822,9 @@ raw_ostream &raw_fd_ostream::changeColor(enum Colors colors, bool bold,
 }
 
 raw_ostream &raw_fd_ostream::resetColor() {
+  if (!ColorEnabled)
+    return *this;
+
   if (sys::Process::ColorNeedsFlush())
     flush();
   const char *colorcode = sys::Process::ResetColor();
@@ -812,6 +838,9 @@ raw_ostream &raw_fd_ostream::resetColor() {
 }
 
 raw_ostream &raw_fd_ostream::reverseColor() {
+  if (!ColorEnabled)
+    return *this;
+
   if (sys::Process::ColorNeedsFlush())
     flush();
   const char *colorcode = sys::Process::OutputReverse();