[Frontend] Support color output in DiagnosticVerifier

Move the color stream utilities into ColorUtils.h,
and use ColorStream to print diagnostic messages
if `-color-diagnostics` is passed.

Author: hamishknight

include/swift/Basic/ColorUtils.h

@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines an 'OSColor' class for helping printing colorful outputs
+// This file defines several utilities for helping print colorful outputs
 // to the terminal.
 //
 //===----------------------------------------------------------------------===//
@@ -42,6 +42,59 @@ class OSColor {
   OSColor &operator<<(llvm::StringRef Str) { OS << Str; return *this; }
 };
 
+/// A stream which forces color output.
+class ColoredStream : public raw_ostream {
+  raw_ostream &Underlying;
+
+public:
+  explicit ColoredStream(raw_ostream &underlying) : Underlying(underlying) {}
+  ~ColoredStream() override { flush(); }
+
+  raw_ostream &changeColor(Colors color, bool bold = false,
+                           bool bg = false) override {
+    Underlying.changeColor(color, bold, bg);
+    return *this;
+  }
+  raw_ostream &resetColor() override {
+    Underlying.resetColor();
+    return *this;
+  }
+  raw_ostream &reverseColor() override {
+    Underlying.reverseColor();
+    return *this;
+  }
+  bool has_colors() const override { return true; }
+
+  void write_impl(const char *ptr, size_t size) override {
+    Underlying.write(ptr, size);
+  }
+  uint64_t current_pos() const override {
+    return Underlying.tell() - GetNumBytesInBuffer();
+  }
+
+  size_t preferred_buffer_size() const override { return 0; }
+};
+
+/// A stream which drops all color settings.
+class NoColorStream : public raw_ostream {
+  raw_ostream &Underlying;
+
+public:
+  explicit NoColorStream(raw_ostream &underlying) : Underlying(underlying) {}
+  ~NoColorStream() override { flush(); }
+
+  bool has_colors() const override { return false; }
+
+  void write_impl(const char *ptr, size_t size) override {
+    Underlying.write(ptr, size);
+  }
+  uint64_t current_pos() const override {
+    return Underlying.tell() - GetNumBytesInBuffer();
+  }
+
+  size_t preferred_buffer_size() const override { return 0; }
+};
+
 } // namespace swift
 
 #endif // SWIFT_BASIC_COLORUTILS_H

include/swift/Frontend/DiagnosticVerifier.h

@@ -94,14 +94,16 @@ class DiagnosticVerifier : public DiagnosticConsumer {
   SmallVector<unsigned, 4> AdditionalBufferIDs;
   bool AutoApplyFixes;
   bool IgnoreUnknown;
+  bool UseColor;
   ArrayRef<std::string> AdditionalExpectedPrefixes;
 
 public:
   explicit DiagnosticVerifier(SourceManager &SM, ArrayRef<unsigned> BufferIDs,
                               bool AutoApplyFixes, bool IgnoreUnknown,
+                              bool UseColor,
                               ArrayRef<std::string> AdditionalExpectedPrefixes)
       : SM(SM), BufferIDs(BufferIDs), AutoApplyFixes(AutoApplyFixes),
-        IgnoreUnknown(IgnoreUnknown),
+        IgnoreUnknown(IgnoreUnknown), UseColor(UseColor),
         AdditionalExpectedPrefixes(AdditionalExpectedPrefixes) {}
 
   void appendAdditionalBufferID(unsigned bufferID) {
@@ -124,6 +126,11 @@ class DiagnosticVerifier : public DiagnosticConsumer {
     bool HadUnexpectedDiag;
   };
 
+  void printDiagnostic(const llvm::SMDiagnostic &Diag) const;
+
+  bool
+  verifyUnknown(std::vector<CapturedDiagnosticInfo> &CapturedDiagnostics) const;
+
   /// verifyFile - After the file has been processed, check to see if we
   /// got all of the expected diagnostics and check to see if there were any
   /// unexpected ones.

lib/Frontend/DiagnosticVerifier.cpp

@@ -15,6 +15,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/Frontend/DiagnosticVerifier.h"
+#include "swift/Basic/ColorUtils.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Parse/Lexer.h"
 #include "llvm/ADT/STLExtras.h"
@@ -328,11 +329,11 @@ static void autoApplyFixes(SourceManager &SM, unsigned BufferID,
   if (!error)
     outs << Result;
 }
+} // end anonymous namespace
 
 /// diagnostics for '<unknown>:0' should be considered as unexpected.
-static bool
-verifyUnknown(SourceManager &SM,
-              std::vector<CapturedDiagnosticInfo> &CapturedDiagnostics) {
+bool DiagnosticVerifier::verifyUnknown(
+    std::vector<CapturedDiagnosticInfo> &CapturedDiagnostics) const {
   bool HadError = false;
   for (unsigned i = 0, e = CapturedDiagnostics.size(); i != e; ++i) {
     if (CapturedDiagnostics[i].Loc.isValid())
@@ -346,11 +347,10 @@ verifyUnknown(SourceManager &SM,
             .str();
 
     auto diag = SM.GetMessage({}, llvm::SourceMgr::DK_Error, Message, {}, {});
-    SM.getLLVMSourceMgr().PrintMessage(llvm::errs(), diag);
+    printDiagnostic(diag);
   }
   return HadError;
 }
-} // end anonymous namespace
 
 /// Return true if the given \p ExpectedFixIt is in the fix-its emitted by
 /// diagnostic \p D.
@@ -377,6 +377,13 @@ bool DiagnosticVerifier::checkForFixIt(
   return false;
 }
 
+void DiagnosticVerifier::printDiagnostic(const llvm::SMDiagnostic &Diag) const {
+  raw_ostream &stream = llvm::errs();
+  ColoredStream coloredStream{stream};
+  raw_ostream &out = UseColor ? coloredStream : stream;
+  SM.getLLVMSourceMgr().PrintMessage(out, Diag);
+}
+
 std::string
 DiagnosticVerifier::renderFixits(ArrayRef<CapturedFixItInfo> ActualFixIts,
                                  unsigned BufferID,
@@ -1184,7 +1191,7 @@ DiagnosticVerifier::Result DiagnosticVerifier::verifyFile(unsigned BufferID) {
 
   // Emit all of the queue'd up errors.
   for (auto Err : Errors)
-    SM.getLLVMSourceMgr().PrintMessage(llvm::errs(), Err);
+    printDiagnostic(Err);
 
   // If auto-apply fixits is on, rewrite the original source file.
   if (AutoApplyFixes)
@@ -1214,10 +1221,11 @@ void DiagnosticVerifier::printRemainingDiagnostics() const {
       break;
     }
 
-    SM.getLLVMSourceMgr().PrintMessage(
-        llvm::errs(), getRawLoc(diag.Loc), SMKind,
-        "diagnostic produced elsewhere: " + diag.Message.str(),
-        /*Ranges=*/{}, {});
+    auto message =
+        SM.GetMessage(diag.Loc, SMKind,
+                      "diagnostic produced elsewhere: " + diag.Message.str(),
+                      /*Ranges=*/{}, {});
+    printDiagnostic(message);
   }
 }
 
@@ -1288,7 +1296,7 @@ bool DiagnosticVerifier::finishProcessing() {
       Result.HadUnexpectedDiag |= FileResult.HadUnexpectedDiag;
     }
   if (!IgnoreUnknown) {
-    bool HadError = verifyUnknown(SM, CapturedDiagnostics);
+    bool HadError = verifyUnknown(CapturedDiagnostics);
     Result.HadError |= HadError;
     // For <unknown>, all errors are unexpected.
     Result.HadUnexpectedDiag |= HadError;

lib/Frontend/Frontend.cpp

@@ -387,7 +387,7 @@ bool CompilerInstance::setupDiagnosticVerifierIfNeeded() {
     DiagVerifier = std::make_unique<DiagnosticVerifier>(
         SourceMgr, InputSourceCodeBufferIDs,
         diagOpts.VerifyMode == DiagnosticOptions::VerifyAndApplyFixes,
-        diagOpts.VerifyIgnoreUnknown,
+        diagOpts.VerifyIgnoreUnknown, diagOpts.UseColor,
         diagOpts.AdditionalDiagnosticVerifierPrefixes);
     for (const auto &filename : diagOpts.AdditionalVerifierFiles) {
       auto result = getFileSystem().getBufferForFile(filename);

lib/Frontend/PrintingDiagnosticConsumer.cpp

@@ -18,6 +18,7 @@
 #include "swift/AST/ASTBridging.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsCommon.h"
+#include "swift/Basic/ColorUtils.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Bridging/ASTGen.h"
@@ -36,61 +37,6 @@ using namespace swift;
 using namespace swift::markup;
 
 namespace {
-  class ColoredStream : public raw_ostream {
-    raw_ostream &Underlying;
-  public:
-    explicit ColoredStream(raw_ostream &underlying) : Underlying(underlying) {}
-    ~ColoredStream() override { flush(); }
-
-    raw_ostream &changeColor(Colors color, bool bold = false,
-                             bool bg = false) override {
-      Underlying.changeColor(color, bold, bg);
-      return *this;
-    }
-    raw_ostream &resetColor() override {
-      Underlying.resetColor();
-      return *this;
-    }
-    raw_ostream &reverseColor() override {
-      Underlying.reverseColor();
-      return *this;
-    }
-    bool has_colors() const override {
-      return true;
-    }
-
-    void write_impl(const char *ptr, size_t size) override {
-      Underlying.write(ptr, size);
-    }
-    uint64_t current_pos() const override {
-      return Underlying.tell() - GetNumBytesInBuffer();
-    }
-
-    size_t preferred_buffer_size() const override {
-      return 0;
-    }
-  };
-
-  /// A stream which drops all color settings.
-  class NoColorStream : public raw_ostream {
-    raw_ostream &Underlying;
-
-  public:
-    explicit NoColorStream(raw_ostream &underlying) : Underlying(underlying) {}
-    ~NoColorStream() override { flush(); }
-
-    bool has_colors() const override { return false; }
-
-    void write_impl(const char *ptr, size_t size) override {
-      Underlying.write(ptr, size);
-    }
-    uint64_t current_pos() const override {
-      return Underlying.tell() - GetNumBytesInBuffer();
-    }
-
-    size_t preferred_buffer_size() const override { return 0; }
-  };
-
 // MARK: Markdown Printing
     class TerminalMarkupPrinter : public MarkupASTVisitor<TerminalMarkupPrinter> {
       llvm::raw_ostream &OS;