[flang] Support the color diagnostics on scanning, parsing, and semantics

The options -f{no-}color-diagnostics have been supported in driver. This
supports the behaviors in scanning, parsing, and semantics, and the
behaviors are exactly the same as the driver.

To illustrate the added behaviour, consider the following input file:
```! file.f90
program m
  integer :: i = k
end
```
In the following invocations, "error: Must be a constant value" _will be_
formatted:
```
$ flang-new file.f90
error: Semantic errors in file.f90
./file.f90:2:18: error: Must be a constant value
    integer :: i = k
```
Note that "error: Semantic errors in file.f90" is also formatted, which
is supported in https://reviews.llvm.org/D126164.

Also note that only "error", "warning" and "portability" are formatted.
Check the following input file:
```! file2.f90
program m
  integer :: i =
end
```
```
$ flang-new test2.f90
error: Could not parse test2.f90
./test2.f90:2:11: error: expected '('
    integer :: i =
            ^
./test2.f90:2:3: in the context: statement function definition
    integer :: i =
    ^
...
```
The "error: Could not parse test2.f90" and "error: expected '('" are
formatted. Others such as "in the context" are not formatted yet, which
may or may not be supported.

Reviewed By: awarzynski

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

12 files changed, 126 insertions, 18 deletions

diff --git a/flang/docs/ReleaseNotes.md b/flang/docs/ReleaseNotes.md
index 38d5c3f..198ff4b 100644
--- a/flang/docs/ReleaseNotes.md
+++ b/flang/docs/ReleaseNotes.md
@@ -41,9 +41,7 @@ page](https://llvm.org/releases/).
   In particular, both `-fcolor-diagnostics` and `-fno-color-diagnostics` are
   now available in `flang-new` (the diagnostics are formatted by default). In
   the frontend driver, `flang-new -fc1`, only `-fcolor-diagnostics` is
-  available (by default, the diagnostics are not formatted). Note that this
-  will only affect the diagnostics printed by driver (scanning, parsing and
-  semantic diagnostics are not affected).
+  available (by default, the diagnostics are not formatted).
 
 ## Windows Support
 
diff --git a/flang/include/flang/Parser/parsing.h b/flang/include/flang/Parser/parsing.h
index f33c1ce..e80d8f7 100644
--- a/flang/include/flang/Parser/parsing.h
+++ b/flang/include/flang/Parser/parsing.h
@@ -39,6 +39,7 @@ struct Options {
   bool needProvenanceRangeToCharBlockMappings{false};
   Fortran::parser::Encoding encoding{Fortran::parser::Encoding::UTF_8};
   bool prescanAndReformat{false}; // -E
+  bool showColors{false};
 };
 
 class Parsing {
@@ -65,9 +66,12 @@ public:
   void ClearLog();
 
   void EmitMessage(llvm::raw_ostream &o, const char *at,
-      const std::string &message, bool echoSourceLine = false) const {
+      const std::string &message, const std::string &prefix,
+      llvm::raw_ostream::Colors color = llvm::raw_ostream::SAVEDCOLOR,
+      bool echoSourceLine = false) const {
     allCooked_.allSources().EmitMessage(o,
-        allCooked_.GetProvenanceRange(CharBlock(at)), message, echoSourceLine);
+        allCooked_.GetProvenanceRange(CharBlock(at)), message, prefix, color,
+        echoSourceLine);
   }
 
 private:
diff --git a/flang/include/flang/Parser/provenance.h b/flang/include/flang/Parser/provenance.h
index eec0b9b..1f18ff6 100644
--- a/flang/include/flang/Parser/provenance.h
+++ b/flang/include/flang/Parser/provenance.h
@@ -165,8 +165,11 @@ public:
   bool IsValid(ProvenanceRange range) const {
     return range.size() > 0 && range_.Contains(range);
   }
+  void setShowColors(bool showColors) { showColors_ = showColors; }
+  bool getShowColors() const { return showColors_; }
   void EmitMessage(llvm::raw_ostream &, const std::optional<ProvenanceRange> &,
-      const std::string &message, bool echoSourceLine = false) const;
+      const std::string &message, const std::string &prefix,
+      llvm::raw_ostream::Colors color, bool echoSourceLine = false) const;
   const SourceFile *GetSourceFile(
       Provenance, std::size_t *offset = nullptr) const;
   const char *GetSource(ProvenanceRange) const;
@@ -214,6 +217,7 @@ private:
   std::vector<std::unique_ptr<SourceFile>> ownedSourceFiles_;
   std::list<std::string> searchPath_;
   Encoding encoding_{Encoding::UTF_8};
+  bool showColors_{false};
 };
 
 // Represents the result of preprocessing and prescanning a single source
diff --git a/flang/lib/Frontend/CompilerInvocation.cpp b/flang/lib/Frontend/CompilerInvocation.cpp
index af6deb2..f794eb7 100644
--- a/flang/lib/Frontend/CompilerInvocation.cpp
+++ b/flang/lib/Frontend/CompilerInvocation.cpp
@@ -769,6 +769,9 @@ void CompilerInvocation::setFortranOpts() {
   if (frontendOptions.instrumentedParse)
     fortranOptions.instrumentedParse = true;
 
+  if (frontendOptions.showColors)
+    fortranOptions.showColors = true;
+
   if (frontendOptions.needProvenanceRangeToCharBlockMappings)
     fortranOptions.needProvenanceRangeToCharBlockMappings = true;
 
diff --git a/flang/lib/Parser/message.cpp b/flang/lib/Parser/message.cpp
index 081dd5c..d729c84 100644
--- a/flang/lib/Parser/message.cpp
+++ b/flang/lib/Parser/message.cpp
@@ -234,18 +234,33 @@ static std::string Prefix(Severity severity) {
   return "";
 }
 
+static llvm::raw_ostream::Colors PrefixColor(Severity severity) {
+  switch (severity) {
+  case Severity::Error:
+  case Severity::Todo:
+    return llvm::raw_ostream::RED;
+  case Severity::Warning:
+  case Severity::Portability:
+    return llvm::raw_ostream::MAGENTA;
+  default:
+    // TODO: Set the color.
+    break;
+  }
+  return llvm::raw_ostream::SAVEDCOLOR;
+}
+
 void Message::Emit(llvm::raw_ostream &o, const AllCookedSources &allCooked,
     bool echoSourceLine) const {
   std::optional<ProvenanceRange> provenanceRange{GetProvenanceRange(allCooked)};
   const AllSources &sources{allCooked.allSources()};
-  sources.EmitMessage(
-      o, provenanceRange, Prefix(severity()) + ToString(), echoSourceLine);
+  sources.EmitMessage(o, provenanceRange, ToString(), Prefix(severity()),
+      PrefixColor(severity()), echoSourceLine);
   bool isContext{attachmentIsContext_};
   for (const Message *attachment{attachment_.get()}; attachment;
        attachment = attachment->attachment_.get()) {
+    Severity severity = isContext ? Severity::Context : attachment->severity();
     sources.EmitMessage(o, attachment->GetProvenanceRange(allCooked),
-        Prefix(isContext ? Severity::Context : attachment->severity()) +
-            attachment->ToString(),
+        attachment->ToString(), Prefix(severity), PrefixColor(severity),
         echoSourceLine);
   }
 }
diff --git a/flang/lib/Parser/parsing.cpp b/flang/lib/Parser/parsing.cpp
index 3ea14e9..1af8afe 100644
--- a/flang/lib/Parser/parsing.cpp
+++ b/flang/lib/Parser/parsing.cpp
@@ -96,6 +96,9 @@ const SourceFile *Parsing::Prescan(const std::string &path, Options options) {
   if (options.needProvenanceRangeToCharBlockMappings) {
     currentCooked_->CompileProvenanceRangeToOffsetMappings(allSources);
   }
+  if (options.showColors) {
+    allSources.setShowColors(/*showColors=*/true);
+  }
   return sourceFile;
 }
 
diff --git a/flang/lib/Parser/provenance.cpp b/flang/lib/Parser/provenance.cpp
index a46111c..355d280 100644
--- a/flang/lib/Parser/provenance.cpp
+++ b/flang/lib/Parser/provenance.cpp
@@ -223,10 +223,26 @@ ProvenanceRange AllSources::AddCompilerInsertion(std::string text) {
   return covers;
 }
 
+static void EmitPrefix(llvm::raw_ostream &o, llvm::raw_ostream::Colors color,
+    const std::string &prefix, bool showColors) {
+  if (prefix.empty()) {
+    return;
+  }
+  if (showColors) {
+    o.changeColor(color, true);
+  }
+  o << prefix;
+  if (showColors) {
+    o.resetColor();
+  }
+}
+
 void AllSources::EmitMessage(llvm::raw_ostream &o,
     const std::optional<ProvenanceRange> &range, const std::string &message,
+    const std::string &prefix, llvm::raw_ostream::Colors color,
     bool echoSourceLine) const {
   if (!range) {
+    EmitPrefix(o, color, prefix, this->getShowColors());
     o << message << '\n';
     return;
   }
@@ -238,8 +254,9 @@ void AllSources::EmitMessage(llvm::raw_ostream &o,
             o << inc.source.path();
             std::size_t offset{origin.covers.MemberOffset(range->start())};
             SourcePosition pos{inc.source.FindOffsetLineAndColumn(offset)};
-            o << ':' << pos.line << ':' << pos.column;
-            o << ": " << message << '\n';
+            o << ':' << pos.line << ':' << pos.column << ": ";
+            EmitPrefix(o, color, prefix, this->getShowColors());
+            o << message << '\n';
             if (echoSourceLine) {
               const char *text{inc.source.content().data() +
                   inc.source.GetLineStartOffset(pos.line)};
@@ -269,14 +286,15 @@ void AllSources::EmitMessage(llvm::raw_ostream &o,
             }
             if (IsValid(origin.replaces)) {
               EmitMessage(o, origin.replaces,
-                  inc.isModule ? "used here"s : "included here"s,
+                  inc.isModule ? "used here"s : "included here"s, prefix, color,
                   echoSourceLine);
             }
           },
           [&](const Macro &mac) {
-            EmitMessage(o, origin.replaces, message, echoSourceLine);
             EmitMessage(
-                o, mac.definition, "in a macro defined here", echoSourceLine);
+                o, origin.replaces, message, prefix, color, echoSourceLine);
+            EmitMessage(o, mac.definition, "in a macro defined here", prefix,
+                color, echoSourceLine);
             if (echoSourceLine) {
               o << "that expanded to:\n  " << mac.expansion << "\n  ";
               for (std::size_t j{0};
@@ -286,7 +304,10 @@ void AllSources::EmitMessage(llvm::raw_ostream &o,
               o << "^\n";
             }
           },
-          [&](const CompilerInsertion &) { o << message << '\n'; },
+          [&](const CompilerInsertion &) {
+            EmitPrefix(o, color, prefix, this->getShowColors());
+            o << message << '\n';
+          },
       },
       origin.u);
 }
diff --git a/flang/test/Driver/color-diagnostics-parse.f90 b/flang/test/Driver/color-diagnostics-parse.f90
new file mode 100644
index 0000000..3d8d3a3
--- /dev/null
+++ b/flang/test/Driver/color-diagnostics-parse.f90
@@ -0,0 +1,20 @@
+! Test the behaviors of -f{no-}color-diagnostics when emitting parsing
+! diagnostics.
+! Windows command prompt doesn't support ANSI escape sequences.
+! REQUIRES: shell
+
+! RUN: not %flang %s -fcolor-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_CD
+! RUN: not %flang %s -fno-color-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_NCD
+! RUN: not %flang_fc1 %s -fcolor-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_CD
+! RUN: not %flang_fc1 %s 2>&1 | FileCheck %s --check-prefix=CHECK_NCD
+
+! CHECK_CD: {{.*}}[0;1;31merror: {{.*}}[0mexpected '('
+
+! CHECK_NCD: error: expected '('
+
+program m
+  integer :: i =
+end
diff --git a/flang/test/Driver/color-diagnostics-scan.f b/flang/test/Driver/color-diagnostics-scan.f
new file mode 100644
index 0000000..d901d77
--- /dev/null
+++ b/flang/test/Driver/color-diagnostics-scan.f
@@ -0,0 +1,19 @@
+! Test the behaviors of -f{no-}color-diagnostics when emitting scanning
+! diagnostics.
+! Windows command prompt doesn't support ANSI escape sequences.
+! REQUIRES: shell
+
+! RUN: not %flang %s -E -Werror -fcolor-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_CD
+! RUN: not %flang %s -E -Werror -fno-color-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_NCD
+! RUN: not %flang_fc1 -E -Werror %s -fcolor-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_CD
+! RUN: not %flang_fc1 -E -Werror %s 2>&1 | FileCheck %s --check-prefix=CHECK_NCD
+
+! CHECK_CD: {{.*}}[0;1;35mwarning: {{.*}}[0mCharacter in fixed-form label field must be a digit
+
+! CHECK_NCD: warning: Character in fixed-form label field must be a digit
+
+1 continue
+end
diff --git a/flang/test/Driver/color-diagnostics-sema.f90 b/flang/test/Driver/color-diagnostics-sema.f90
new file mode 100644
index 0000000..df7a69f
--- /dev/null
+++ b/flang/test/Driver/color-diagnostics-sema.f90
@@ -0,0 +1,20 @@
+! Test the behaviors of -f{no-}color-diagnostics when emitting semantic
+! diagnostics.
+! Windows command prompt doesn't support ANSI escape sequences.
+! REQUIRES: shell
+
+! RUN: not %flang %s -fcolor-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_CD
+! RUN: not %flang %s -fno-color-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_NCD
+! RUN: not %flang_fc1 %s -fcolor-diagnostics 2>&1 \
+! RUN:     | FileCheck %s --check-prefix=CHECK_CD
+! RUN: not %flang_fc1 %s 2>&1 | FileCheck %s --check-prefix=CHECK_NCD
+
+! CHECK_CD: {{.*}}[0;1;31merror: {{.*}}[0mMust be a constant value
+
+! CHECK_NCD: error: Must be a constant value
+
+program m
+  integer :: i = k
+end
diff --git a/flang/tools/bbc/bbc.cpp b/flang/tools/bbc/bbc.cpp
index 56ca8f7..2ad437a 100644
--- a/flang/tools/bbc/bbc.cpp
+++ b/flang/tools/bbc/bbc.cpp
@@ -167,7 +167,8 @@ static mlir::LogicalResult convertFortranSourceToMLIR(
   parsing.messages().Emit(llvm::errs(), parsing.allCooked());
   if (!parsing.consumedWholeFile()) {
     parsing.EmitMessage(llvm::errs(), parsing.finalRestingPlace(),
-                        "parser FAIL (final position)");
+                        "parser FAIL (final position)",
+                        "error: ", llvm::raw_ostream::RED);
     return mlir::failure();
   }
   if ((!parsing.messages().empty() && (parsing.messages().AnyFatalError())) ||
diff --git a/flang/tools/f18-parse-demo/f18-parse-demo.cpp b/flang/tools/f18-parse-demo/f18-parse-demo.cpp
index 2c44b90..d251154 100644
--- a/flang/tools/f18-parse-demo/f18-parse-demo.cpp
+++ b/flang/tools/f18-parse-demo/f18-parse-demo.cpp
@@ -201,7 +201,7 @@ std::string CompileFortran(
   parsing.messages().Emit(llvm::errs(), parsing.allCooked());
   if (!parsing.consumedWholeFile()) {
     parsing.EmitMessage(llvm::errs(), parsing.finalRestingPlace(),
-        "parser FAIL (final position)");
+        "parser FAIL (final position)", "error: ", llvm::raw_ostream::RED);
     exitStatus = EXIT_FAILURE;
     return {};
   }