Improve raw_ostream so that you can "write" colors using operator<<

1. raw_ostream supports ANSI colors so that you can write messages to
the termina with colors. Previously, in order to change and reset
color, you had to call `changeColor` and `resetColor` functions,
respectively.

So, if you print out "error: " in red, for example, you had to do
something like this:

  OS.changeColor(raw_ostream::RED);
  OS << "error: ";
  OS.resetColor();

With this patch, you can write the same code as follows:

  OS << raw_ostream::RED << "error: " << raw_ostream::RESET;

2. Add a boolean flag to raw_ostream so that you can disable colored
output. If you disable colors, changeColor, operator<<(Color),
resetColor and other color-related functions have no effect.

Most LLVM tools automatically prints out messages using colors, and
you can disable it by passing a flag such as `--disable-colors`.
This new flag makes it easy to write code that works that way.

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

llvm-svn: 367649

clang/include/clang/AST/ASTDumperUtils.h

@@ -27,7 +27,7 @@ enum ASTDumpOutputFormat {
 // Do not use bold yellow for any text.  It is hard to read on white screens.
 
 struct TerminalColor {
-  llvm::raw_ostream::Colors Color;
+  llvm::raw_ostream::Color Color;
   bool Bold;
 };
 

clang/lib/Analysis/CFG.cpp

@@ -5509,7 +5509,7 @@ static void print_block(raw_ostream &OS, const CFG* cfg,
   if (print_edges) {
     // Print the predecessors of this block.
     if (!B.pred_empty()) {
-      const raw_ostream::Colors Color = raw_ostream::BLUE;
+      const raw_ostream::Color Color = raw_ostream::BLUE;
       if (ShowColors)
         OS.changeColor(Color);
       OS << "   Preds " ;
@@ -5546,7 +5546,7 @@ static void print_block(raw_ostream &OS, const CFG* cfg,
 
     // Print the successors of this block.
     if (!B.succ_empty()) {
-      const raw_ostream::Colors Color = raw_ostream::MAGENTA;
+      const raw_ostream::Color Color = raw_ostream::MAGENTA;
       if (ShowColors)
         OS.changeColor(Color);
       OS << "   Succs ";

clang/lib/Frontend/CompilerInvocation.cpp

@@ -1491,6 +1491,10 @@ bool clang::ParseDiagnosticArgs(DiagnosticOptions &Opts, ArgList &Args,
                    OPT_fno_diagnostics_show_option, DefaultShowOpt);
 
   llvm::sys::Process::UseANSIEscapeCodes(Args.hasArg(OPT_fansi_escape_codes));
+  if (Opts.ShowColors) {
+    llvm::outs().enable_colors();
+    llvm::errs().enable_colors();
+  }
 
   // Default behavior is to not to show note include stacks.
   Opts.ShowNoteIncludeStack = false;

clang/lib/Frontend/TextDiagnostic.cpp

@@ -23,23 +23,16 @@
 
 using namespace clang;
 
-static const enum raw_ostream::Colors noteColor =
-  raw_ostream::BLACK;
-static const enum raw_ostream::Colors remarkColor =
-  raw_ostream::BLUE;
-static const enum raw_ostream::Colors fixitColor =
-  raw_ostream::GREEN;
-static const enum raw_ostream::Colors caretColor =
-  raw_ostream::GREEN;
-static const enum raw_ostream::Colors warningColor =
-  raw_ostream::MAGENTA;
-static const enum raw_ostream::Colors templateColor =
-  raw_ostream::CYAN;
-static const enum raw_ostream::Colors errorColor = raw_ostream::RED;
-static const enum raw_ostream::Colors fatalColor = raw_ostream::RED;
+static const raw_ostream::Color noteColor = raw_ostream::BLACK;
+static const raw_ostream::Color remarkColor = raw_ostream::BLUE;
+static const raw_ostream::Color fixitColor = raw_ostream::GREEN;
+static const raw_ostream::Color caretColor = raw_ostream::GREEN;
+static const raw_ostream::Color warningColor = raw_ostream::MAGENTA;
+static const raw_ostream::Color templateColor = raw_ostream::CYAN;
+static const raw_ostream::Color errorColor = raw_ostream::RED;
+static const raw_ostream::Color fatalColor = raw_ostream::RED;
 // Used for changing only the bold attribute.
-static const enum raw_ostream::Colors savedColor =
-  raw_ostream::SAVEDCOLOR;
+static const raw_ostream::Color savedColor = raw_ostream::SAVEDCOLOR;
 
 /// Add highlights to differences in template strings.
 static void applyTemplateHighlighting(raw_ostream &OS, StringRef Str,

clang/tools/diagtool/TreeView.cpp

@@ -20,29 +20,16 @@ DEF_DIAGTOOL("tree", "Show warning flags in a tree view", TreeView)
 using namespace clang;
 using namespace diagtool;
 
-static bool hasColors(const llvm::raw_ostream &out) {
-  if (&out != &llvm::errs() && &out != &llvm::outs())
-    return false;
-  return llvm::errs().is_displayed() && llvm::outs().is_displayed();
-}
-
 class TreePrinter {
+  using Color = llvm::raw_ostream::Color;
+
 public:
   llvm::raw_ostream &out;
-  const bool ShowColors;
   bool Internal;
 
-  TreePrinter(llvm::raw_ostream &out)
-      : out(out), ShowColors(hasColors(out)), Internal(false) {}
-
-  void setColor(llvm::raw_ostream::Colors Color) {
-    if (ShowColors)
-      out << llvm::sys::Process::OutputColor(Color, false, false);
-  }
-
-  void resetColor() {
-    if (ShowColors)
-      out << llvm::sys::Process::ResetColor();
+  TreePrinter(llvm::raw_ostream &out) : out(out), Internal(false) {
+    if (&out != &llvm::errs() && &out != &llvm::outs())
+      out.disable_colors();
   }
 
   static bool isIgnored(unsigned DiagID) {
@@ -70,12 +57,11 @@ class TreePrinter {
     out.indent(Indent * 2);
 
     if (enabledByDefault(Group))
-      setColor(llvm::raw_ostream::GREEN);
+      out << Color::GREEN;
     else
-      setColor(llvm::raw_ostream::YELLOW);
+      out << Color::YELLOW;
 
-    out << "-W" << Group.getName() << "\n";
-    resetColor();
+    out << "-W" << Group.getName() << "\n" << Color::RESET;
 
     ++Indent;
     for (const GroupRecord &GR : Group.subgroups()) {
@@ -84,12 +70,10 @@ class TreePrinter {
 
     if (Internal) {
       for (const DiagnosticRecord &DR : Group.diagnostics()) {
-        if (ShowColors && !isIgnored(DR.DiagID))
-          setColor(llvm::raw_ostream::GREEN);
+        if (!isIgnored(DR.DiagID))
+          out << Color::GREEN;
         out.indent(Indent * 2);
-        out << DR.getName();
-        resetColor();
-        out << "\n";
+        out << DR.getName() << Color::RESET << "\n";
       }
     }
   }
@@ -135,13 +119,9 @@ class TreePrinter {
   }
 
   void showKey() {
-    if (ShowColors) {
-      out << '\n';
-      setColor(llvm::raw_ostream::GREEN);
-      out << "GREEN";
-      resetColor();
-      out << " = enabled by default\n\n";
-    }
+    out << '\n'
+        << Color::GREEN << "GREEN" << Color::RESET
+        << " = enabled by default\n\n";
   }
 };
 

lld/COFF/Driver.cpp

@@ -62,7 +62,6 @@ 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)";

lld/COFF/DriverUtils.cpp

@@ -756,16 +756,17 @@ static void handleColorDiagnostics(opt::InputArgList &args) {
                               OPT_no_color_diagnostics);
   if (!arg)
     return;
+
   if (arg->getOption().getID() == OPT_color_diagnostics) {
-    errorHandler().colorDiagnostics = true;
+    errorHandler().errorOS->enable_colors();
   } else if (arg->getOption().getID() == OPT_no_color_diagnostics) {
-    errorHandler().colorDiagnostics = false;
+    errorHandler().errorOS->disable_colors();
   } else {
     StringRef s = arg->getValue();
     if (s == "always")
-      errorHandler().colorDiagnostics = true;
+      errorHandler().errorOS->enable_colors();
     else if (s == "never")
-      errorHandler().colorDiagnostics = false;
+      errorHandler().errorOS->disable_colors();
     else if (s != "auto")
       error("unknown option: --color-diagnostics=" + s);
   }

lld/Common/ErrorHandler.cpp

@@ -85,54 +85,41 @@ void lld::checkError(Error e) {
                   [&](ErrorInfoBase &eib) { error(eib.message()); });
 }
 
-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();
 
-void ErrorHandler::printHeader(StringRef s, raw_ostream::Colors c,
-                               const Twine &msg) {
+  for (std::regex &re : regexes) {
+    std::smatch match;
+    if (!std::regex_search(str, match, re))
+      continue;
 
-  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 << ": ";
+    if (match.size() > 2)
+      return match.str(1) + "(" + match.str(2) + ")";
+    return match.str(1);
   }
 
-  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) {
@@ -149,42 +136,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";
-}
-
-void ErrorHandler::printErrorMsg(const Twine &msg) {
-  newline(errorOS, msg);
-  printHeader("error: ", raw_ostream::RED, msg);
-  *errorOS << msg << "\n";
+  *errorOS << getLocation(msg) << ": " << Color::MAGENTA
+           << "warning: " << Color::RESET << msg << "\n";
 }
 
-void ErrorHandler::printError(const Twine &msg) {
+void ErrorHandler::error(const Twine &msg) {
+  // If Microsoft 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) << ": " << Color::RED
+             << "error: " << Color::RESET << msg << "\n";
   } else if (errorCount == errorLimit) {
     newline(errorOS, msg);
-    printHeader("error: ", raw_ostream::RED, msg);
-    *errorOS << errorLimitExceededMsg << "\n";
+    *errorOS << getLocation(msg) << ": " << Color::RED
+             << "error: " << Color::RESET << errorLimitExceededMsg << "\n";
     if (exitEarly)
       exitLld(1);
   }

lld/ELF/Driver.cpp

@@ -82,7 +82,6 @@ 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();
 
   inputSections.clear();
   outputSections.clear();

lld/ELF/DriverUtils.cpp

@@ -58,16 +58,17 @@ static void handleColorDiagnostics(opt::InputArgList &args) {
                               OPT_no_color_diagnostics);
   if (!arg)
     return;
+
   if (arg->getOption().getID() == OPT_color_diagnostics) {
-    errorHandler().colorDiagnostics = true;
+    errorHandler().errorOS->enable_colors();
   } else if (arg->getOption().getID() == OPT_no_color_diagnostics) {
-    errorHandler().colorDiagnostics = false;
+    errorHandler().errorOS->disable_colors();
   } else {
     StringRef s = arg->getValue();
     if (s == "always")
-      errorHandler().colorDiagnostics = true;
+      errorHandler().errorOS->enable_colors();
     else if (s == "never")
-      errorHandler().colorDiagnostics = false;
+      errorHandler().errorOS->disable_colors();
     else if (s != "auto")
       error("unknown option: --color-diagnostics=" + s);
   }

lld/include/lld/Common/ErrorHandler.h

@@ -87,7 +87,6 @@ class ErrorHandler {
   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,9 +101,9 @@ class ErrorHandler {
   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 Color = raw_ostream::Color;
+
+  std::string getLocation(const Twine &msg);
 };
 
 /// Returns the default error handler.

lld/lib/Driver/DarwinLdDriver.cpp

@@ -1151,7 +1151,6 @@ 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();
 
   MachOLinkingContext ctx;
   if (!parse(args, ctx))

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 \