[SourceMgr][FileCheck] Obey -color by extending WithColor

While this change specifically targets FileCheck, it affects any tool
using the same SourceMgr facilities.

Previously, -color was documented in FileCheck's -help output, but
-color had no effect.  Now, -color obeys its documentation: it forces
colors to be used in FileCheck diagnostics even when stderr is not a
terminal.

-color is especially helpful when combined with FileCheck's -v, which
can produce a long series of diagnostics that you might wish to pipe
to a pager, such as less -R.  The WithColor extensions here will also
help to clean up color usage in FileCheck's annotated dump of input,
which is proposed in D52999.

Reviewed By: JDevlieghere

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

llvm-svn: 344930

2 files changed, 92 insertions, 79 deletions

diff --git a/llvm/lib/Support/SourceMgr.cpp b/llvm/lib/Support/SourceMgr.cpp
index 582e2cf6c11..a55ad881d01 100644
--- a/llvm/lib/Support/SourceMgr.cpp
+++ b/llvm/lib/Support/SourceMgr.cpp
@@ -24,6 +24,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SMLoc.h"
+#include "llvm/Support/WithColor.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cassert>
@@ -370,65 +371,48 @@ static bool isNonASCII(char c) {
   return c & 0x80;
 }
 
-void SMDiagnostic::print(const char *ProgName, raw_ostream &S, bool ShowColors,
-                         bool ShowKindLabel) const {
-  // Display colors only if OS supports colors.
-  ShowColors &= S.has_colors();
+void SMDiagnostic::print(const char *ProgName, raw_ostream &OS,
+                         bool ShowColors, bool ShowKindLabel) const {
+  {
+    WithColor S(OS, raw_ostream::SAVEDCOLOR, true, false, !ShowColors);
 
-  if (ShowColors)
-    S.changeColor(raw_ostream::SAVEDCOLOR, true);
+    if (ProgName && ProgName[0])
+      S << ProgName << ": ";
 
-  if (ProgName && ProgName[0])
-    S << ProgName << ": ";
+    if (!Filename.empty()) {
+      if (Filename == "-")
+        S << "<stdin>";
+      else
+        S << Filename;
 
-  if (!Filename.empty()) {
-    if (Filename == "-")
-      S << "<stdin>";
-    else
-      S << Filename;
-
-    if (LineNo != -1) {
-      S << ':' << LineNo;
-      if (ColumnNo != -1)
-        S << ':' << (ColumnNo+1);
+      if (LineNo != -1) {
+        S << ':' << LineNo;
+        if (ColumnNo != -1)
+          S << ':' << (ColumnNo + 1);
+      }
+      S << ": ";
     }
-    S << ": ";
   }
 
   if (ShowKindLabel) {
     switch (Kind) {
     case SourceMgr::DK_Error:
-      if (ShowColors)
-        S.changeColor(raw_ostream::RED, true);
-      S << "error: ";
+      WithColor::error(OS, "", !ShowColors);
       break;
     case SourceMgr::DK_Warning:
-      if (ShowColors)
-        S.changeColor(raw_ostream::MAGENTA, true);
-      S << "warning: ";
+      WithColor::warning(OS, "", !ShowColors);
       break;
     case SourceMgr::DK_Note:
-      if (ShowColors)
-        S.changeColor(raw_ostream::BLACK, true);
-      S << "note: ";
+      WithColor::note(OS, "", !ShowColors);
       break;
     case SourceMgr::DK_Remark:
-      if (ShowColors)
-        S.changeColor(raw_ostream::BLUE, true);
-      S << "remark: ";
+      WithColor::remark(OS, "", !ShowColors);
       break;
     }
-
-    if (ShowColors) {
-      S.resetColor();
-      S.changeColor(raw_ostream::SAVEDCOLOR, true);
-    }
   }
 
-  S << Message << '\n';
-
-  if (ShowColors)
-    S.resetColor();
+  WithColor(OS, raw_ostream::SAVEDCOLOR, true, false, !ShowColors)
+      << Message << '\n';
 
   if (LineNo == -1 || ColumnNo == -1)
     return;
@@ -439,7 +423,7 @@ void SMDiagnostic::print(const char *ProgName, raw_ostream &S, bool ShowColors,
   // expanding them later, and bail out rather than show incorrect ranges and
   // misaligned fixits for any other odd characters.
   if (find_if(LineContents, isNonASCII) != LineContents.end()) {
-    printSourceLine(S, LineContents);
+    printSourceLine(OS, LineContents);
     return;
   }
   size_t NumColumns = LineContents.size();
@@ -473,29 +457,27 @@ void SMDiagnostic::print(const char *ProgName, raw_ostream &S, bool ShowColors,
   // least.
   CaretLine.erase(CaretLine.find_last_not_of(' ')+1);
 
-  printSourceLine(S, LineContents);
+  printSourceLine(OS, LineContents);
 
-  if (ShowColors)
-    S.changeColor(raw_ostream::GREEN, true);
+  {
+    WithColor S(OS, raw_ostream::GREEN, true, false, !ShowColors);
 
-  // Print out the caret line, matching tabs in the source line.
-  for (unsigned i = 0, e = CaretLine.size(), OutCol = 0; i != e; ++i) {
-    if (i >= LineContents.size() || LineContents[i] != '\t') {
-      S << CaretLine[i];
-      ++OutCol;
-      continue;
-    }
+    // Print out the caret line, matching tabs in the source line.
+    for (unsigned i = 0, e = CaretLine.size(), OutCol = 0; i != e; ++i) {
+      if (i >= LineContents.size() || LineContents[i] != '\t') {
+        S << CaretLine[i];
+        ++OutCol;
+        continue;
+      }
 
-    // Okay, we have a tab.  Insert the appropriate number of characters.
-    do {
-      S << CaretLine[i];
-      ++OutCol;
-    } while ((OutCol % TabStop) != 0);
+      // Okay, we have a tab.  Insert the appropriate number of characters.
+      do {
+        S << CaretLine[i];
+        ++OutCol;
+      } while ((OutCol % TabStop) != 0);
+    }
+    S << '\n';
   }
-  S << '\n';
-
-  if (ShowColors)
-    S.resetColor();
 
   // Print out the replacement line, matching tabs in the source line.
   if (FixItInsertionLine.empty())
@@ -503,14 +485,14 @@ void SMDiagnostic::print(const char *ProgName, raw_ostream &S, bool ShowColors,
 
   for (size_t i = 0, e = FixItInsertionLine.size(), OutCol = 0; i < e; ++i) {
     if (i >= LineContents.size() || LineContents[i] != '\t') {
-      S << FixItInsertionLine[i];
+      OS << FixItInsertionLine[i];
       ++OutCol;
       continue;
     }
 
     // Okay, we have a tab.  Insert the appropriate number of characters.
     do {
-      S << FixItInsertionLine[i];
+      OS << FixItInsertionLine[i];
       // FIXME: This is trying not to break up replacements, but then to re-sync
       // with the tabs between replacements. This will fail, though, if two
       // fix-it replacements are exactly adjacent, or if a fix-it contains a
@@ -521,5 +503,5 @@ void SMDiagnostic::print(const char *ProgName, raw_ostream &S, bool ShowColors,
       ++OutCol;
     } while (((OutCol % TabStop) != 0) && i != e);
   }
-  S << '\n';
+  OS << '\n';
 }
diff --git a/llvm/lib/Support/WithColor.cpp b/llvm/lib/Support/WithColor.cpp
index d2e13f0e86d..cf4c10956f2 100644
--- a/llvm/lib/Support/WithColor.cpp
+++ b/llvm/lib/Support/WithColor.cpp
@@ -19,15 +19,10 @@ static cl::opt<cl::boolOrDefault>
              cl::desc("Use colors in output (default=autodetect)"),
              cl::init(cl::BOU_UNSET));
 
-bool WithColor::colorsEnabled(raw_ostream &OS) {
-  if (UseColor == cl::BOU_UNSET)
-    return OS.has_colors();
-  return UseColor == cl::BOU_TRUE;
-}
-
-WithColor::WithColor(raw_ostream &OS, HighlightColor Color) : OS(OS) {
+WithColor::WithColor(raw_ostream &OS, HighlightColor Color, bool DisableColors)
+    : OS(OS), DisableColors(DisableColors) {
   // Detect color from terminal type unless the user passed the --color option.
-  if (colorsEnabled(OS)) {
+  if (colorsEnabled()) {
     switch (Color) {
     case HighlightColor::Address:
       OS.changeColor(raw_ostream::YELLOW);
@@ -56,6 +51,9 @@ WithColor::WithColor(raw_ostream &OS, HighlightColor Color) : OS(OS) {
     case HighlightColor::Note:
       OS.changeColor(raw_ostream::BLACK, true);
       break;
+    case HighlightColor::Remark:
+      OS.changeColor(raw_ostream::BLUE, true);
+      break;
     }
   }
 }
@@ -66,25 +64,58 @@ raw_ostream &WithColor::warning() { return warning(errs()); }
 
 raw_ostream &WithColor::note() { return note(errs()); }
 
-raw_ostream &WithColor::error(raw_ostream &OS, StringRef Prefix) {
+raw_ostream &WithColor::remark() { return remark(errs()); }
+
+raw_ostream &WithColor::error(raw_ostream &OS, StringRef Prefix,
+                              bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Error).get() << "error: ";
+  return WithColor(OS, HighlightColor::Error, DisableColors).get()
+         << "error: ";
 }
 
-raw_ostream &WithColor::warning(raw_ostream &OS, StringRef Prefix) {
+raw_ostream &WithColor::warning(raw_ostream &OS, StringRef Prefix,
+                                bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Warning).get() << "warning: ";
+  return WithColor(OS, HighlightColor::Warning, DisableColors).get()
+         << "warning: ";
 }
 
-raw_ostream &WithColor::note(raw_ostream &OS, StringRef Prefix) {
+raw_ostream &WithColor::note(raw_ostream &OS, StringRef Prefix,
+                             bool DisableColors) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Note).get() << "note: ";
+  return WithColor(OS, HighlightColor::Note, DisableColors).get() << "note: ";
 }
 
-WithColor::~WithColor() {
-  if (colorsEnabled(OS))
+raw_ostream &WithColor::remark(raw_ostream &OS, StringRef Prefix,
+                               bool DisableColors) {
+  if (!Prefix.empty())
+    OS << Prefix << ": ";
+  return WithColor(OS, HighlightColor::Remark, DisableColors).get()
+         << "remark: ";
+}
+
+bool WithColor::colorsEnabled() {
+  if (DisableColors)
+    return false;
+  if (UseColor == cl::BOU_UNSET)
+    return OS.has_colors();
+  return UseColor == cl::BOU_TRUE;
+}
+
+WithColor &WithColor::changeColor(raw_ostream::Colors Color, bool Bold,
+                                  bool BG) {
+  if (colorsEnabled())
+    OS.changeColor(Color, Bold, BG);
+  return *this;
+}
+
+WithColor &WithColor::resetColor() {
+  if (colorsEnabled())
     OS.resetColor();
+  return *this;
 }
+
+WithColor::~WithColor() { resetColor(); }