[Support] Support restoring colors in WithColor

[seiya]

Summary

WithColor provides a way to change text color temporarily but currently it resets colors
in the destructor regardless of the previous color. Consider this example:

{
  WithColor Red(outs(), raw_ostream::RED);
  Red << "AAA\n";
  WithColor(outs(), raw_ostream::BLUE) << "BBB\n";
  Red << "CCC\n";
}

In this example "CCC" is not colored with red because the WithColor instance for
"BBB" resets the color in its destructor.

This patch addresses the issue by saving the previous color state and restoring
the color settings in the destructor.

[seiya]

This patch does not include tests but it will be used in D65191.

[seiya]

Added reviewers who may be familiar with this part.

[ruiu]

Instead of adding a new class, do you think you can change the behavior of WithColor so that it restores the original colors instead of resetting it?

[jhenderson]

In D67060#1654259, @ruiu wrote:

Instead of adding a new class, do you think you can change the behavior of WithColor so that it restores the original colors instead of resetting it?

This would be my preference too, if it's viable. I don't think WithColor has been around that long, so I suspect it may be straightforward to identify existing usage patterns, and I'd be surprised if any expect restoration to default.

[MaskRay]

Reusing WithColor is also my preference. I think you can change the third argument DisableColors to WithColorContext *Context = nullptr in the following overload:

static raw_ostream &error(raw_ostream &OS, StringRef Prefix = "",
                          bool DisableColors = false);

And make DisableColors a parameter with a default argument of the constructor of WithColorContext

[seiya]

In D67060#1654353, @jhenderson wrote:

In D67060#1654259, @ruiu wrote:

Instead of adding a new class, do you think you can change the behavior of WithColor so that it restores the original colors instead of resetting it?

This would be my preference too, if it's viable. I don't think WithColor has been around that long, so I suspect it may be straightforward to identify existing usage patterns, and I'd be surprised if any expect restoration to default.

I didn't came up with a good idea to add this feature to WithColor with the compatibility with existing usage but @MaskRay's idea looks good to me. I'll try the idea.

[ruiu]

Did you find any existing usage of WithColor that depends on the behavior of resetting the color? If not, you don't have to care too much about keeping the compatibility, as that is not really an intended feature as far as I understand. You probably can just change the behavior of WithColor so that it restores to the previous color instead of resseting.

[seiya]

Moved NestableWithColor features into WithColor.

[seiya]

In D67060#1655157, @ruiu wrote:

Did you find any existing usage of WithColor that depends on the behavior of resetting the color? If not, you don't have to care too much about keeping the compatibility, as that is not really an intended feature as far as I understand. You probably can just change the behavior of WithColor so that it restores to the previous color instead of resseting.

I grep'ed the existing usage and ran tests in llvm/test/tools but it looks that no tests relies on the behavior. As you say it's safe to change the behavior I think.

[seiya]

llvm/include/llvm/Support/WithColor.h

163

In D67060#1655078, @MaskRay suggested to replace bool DisableColors with WithColorContext *Context. IIUC, the usage will look like:

WithColorContext Ctx(/*DisableColors=*/ false);
{
  WithColor Red(outs(), raw_ostream::RED, false, false, &Ctx);
  Red << "AAA\n";
  WithColor(outs(), raw_ostream::BLUE, false, false, &Ctx) << "BBB\n";
  Red << "CCC\n";
}

However, with this approach we need to Bold and BG every time we use WithColorContext. How about adding a new constructor (b):

// (a) Use DisableColors.
WithColor(raw_ostream &OS,
          raw_ostream::Colors Color = raw_ostream::SAVEDCOLOR,
          bool Bold = false, bool BG = false, bool DisableColors = false,
          WithColorContext *Context = nullptr);

// (b) Use Context->DisableColors.
WithColor(raw_ostream &OS, WithColorContext *Context,
          raw_ostream::Colors Color = raw_ostream::SAVEDCOLOR,
          bool Bold = false, bool BG = false)

[ruiu]

This looks better, but I don't think we need the concept of context. If a function takes a context object, it is usually expected that you can interlace function calls with different context objects freely. Say, you have context object A and B, then you literally have two contexts, and they don't interfere with each other. However, that's not the case for the terminal color. The current color of text output on console is naturally and fundamentally a global state associated to each terminal, and therefore there is only one context. So, I don't think that allowing multiple contexts captures the real concept correctly.

I believe the following scheme models captures the reality better:

1. Add a new method, getColor, to raw_fd_stream, which returns the current color set by changeColor before
2. Make a change to WithColor so that it remembers the current color (obtained by calling getColor) before changing color
3. In the destructor, WithColor restores the previous color

[seiya]

In D67060#1656821, @ruiu wrote:

This looks better, but I don't think we need the concept of context. If a function takes a context object, it is usually expected that you can interlace function calls with different context objects freely. Say, you have context object A and B, then you literally have two contexts, and they don't interfere with each other. However, that's not the case for the terminal color. The current color of text output on console is naturally and fundamentally a global state associated to each terminal, and therefore there is only one context. So, I don't think that allowing multiple contexts captures the real concept correctly.

I believe the following scheme models captures the reality better:

1. Add a new method, getColor, to raw_fd_stream, which returns the current color set by changeColor before
2. Make a change to WithColor so that it remembers the current color (obtained by calling getColor) before changing color
3. In the destructor, WithColor restores the previous color

That sounds good idea to me too but how can we obtain the current color from the terminal? I skimmed xterm control sequences but couldn't find out the feature (note: I'm not familiar with Windows).

Also, getting colors from the terminal may degrades the performance since we need to invoke a system call every time we use WithColor. We should discuss about the performance after measuring it, btw.

[ruiu]

I don't think there's a way to get the current color from the terminal, so you should modify changeColor so that the function remembers the color given to the function.

[seiya]

In D67060#1656873, @ruiu wrote:

I don't think there's a way to get the current color from the terminal, so you should modify changeColor so that the function remembers the color given to the function.

Oh I get it it's good idea! Will try it.

[seiya]

• Abandon WithColorContext: save the previous color settings in WithColor instead.
• Update the description.

[MaskRay]

The new approach looks good.

llvm/include/llvm/Support/raw_ostream.h

94 ↗

(On Diff #219344)

The struct will be an aggregate if you delete the constructor. It can still be brace initialized.

432 ↗

(On Diff #219344)

Delete enum

[jhenderson]

llvm/include/llvm/Support/raw_ostream.h

92 ↗

(On Diff #219344)

BG is not an obvious abbreviation for me. Could you use an unabbreviated name please?

[seiya]

Addressed review comments.

[ruiu]

This patch looks good, but I'd add a test for the new code. Is there any place in LLVM where the current color is saved and restored? If there's such code, you can rewrite the code to see if your code would work as expected.

[ruiu]

llvm/include/llvm/Support/WithColor.h

65–66

... then you can remove {}

llvm/include/llvm/Support/raw_ostream.h

90–93 ↗

(On Diff #219465)

C++11 allows to write default values for member variables, i.e.

Colors Color = Colors::RESET;
bool Bold = false;
bool Background = false;

309 ↗

(On Diff #219465)

and this {}

401 ↗

(On Diff #219465)

and this as well

[seiya]

Addressed review comments

[ruiu]

llvm/lib/Support/WithColor.cpp

74

Unrelated change?

[seiya]

In D67060#1664226, @ruiu wrote:

This patch looks good, but I'd add a test for the new code. Is there any place in LLVM where the current color is saved and restored? If there's such code, you can rewrite the code to see if your code would work as expected.

There're no such code AFAIK. How about committing this patch and D65191 at once? It will include a test case which depends on this feature.

[seiya]

• Removed an unrelated change from the diff.

[seiya]

llvm/lib/Support/WithColor.cpp

74

Yes. Fixed.

[ruiu]

LGTM

If you are going to submit another change to test this code, it is fine to submit this patch separately. Thank you for doign this.

llvm/include/llvm/Support/WithColor.h

 //===- WithColor.h ----------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_WITHCOLOR_H
 #define LLVM_SUPPORT_WITHCOLOR_H
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/CommandLine.h"
+#include <stack>
 
 namespace llvm {
 
 extern cl::OptionCategory ColorCategory;
 
 class raw_ostream;
 
 // Symbolic names for various syntax elements.
 enum class HighlightColor {
   Address,
   String,
   Tag,
   Attribute,
   Enumerator,
   Macro,
   Error,
   Warning,
   Note,
   Remark
 };
 
+/// A color settings saved in WithColorContext.
+struct ColorState {
+  raw_ostream::Colors Color;
+  bool Bold;
+  bool BG;
+
+  ColorState(raw_ostream::Colors Color, bool Bold, bool BG)
+      : Color(Color), Bold(Bold), BG(BG) {}
+};
+
+/// A context for NestableWithColor which holds saved colors.
+class WithColorContext {
+  std::stack<ColorState> Colors;
+
+public:
+  void pushColor(raw_ostream::Colors Color, bool Bold, bool BG = false) {
+    Colors.emplace(Color, Bold, BG);
+  }
+
+  void popColor() { Colors.pop(); }
+  bool empty() { return Colors.empty(); }
+  const ColorState &currentColor() { return Colors.top(); }
+};
+
 /// An RAII object that temporarily switches an output stream to a specific
 /// color.
 class WithColor {
   raw_ostream &OS;
   bool DisableColors;
+  WithColorContext *Context;
 
 public:
-  /// To be used like this: WithColor(OS, HighlightColor::String) << "text";
+  /// To change colors temporarily, use WithColor like this:
+  ///
+  ///   WithColor(OS, HighlightColor::String) << "text";
+  ///
+  /// If WithColorContext is given, it restores colors properly:
+  ///
+  ///   WithColorContext Ctx;
+  ///   {
+  ///     WithColor AddrColor(OS, HighlightColor::Address, false, &Ctx);
+  ///     AddrColor << "AAA";
+  ///     {
+  ///        WithColor StringColor(OS, HighlightColor::String, false, &Ctx);
+  ///        StringColor << "BBB";
+  ///        // The Blue's destructor is called here and it restores the
+  ///        // AddrColor's color state.
+  ///     }
+  ///     AddrColor << "CCC\n";
+  ///     // The AddrColor's destructor resets the colors to defaults.
+  ///   }
+  ///
+  /// Otherwise, WithColor resets colors in the destructor regardless of the
+  /// previous color. In the example above, "CCC" will be printed in the
+  /// terminal default colors, not in HighlightColor::String colors.
+  ///
   /// @param OS The output stream
   /// @param S Symbolic name for syntax element to color
   /// @param DisableColors Whether to ignore color changes regardless of -color
   /// and support in OS
-  WithColor(raw_ostream &OS, HighlightColor S, bool DisableColors = false);
-  /// To be used like this: WithColor(OS, raw_ostream::Black) << "text";
+  /// @param Context A WithColorContext object. Set nullptr if you don't need
+  ///                to restore colors.
+  WithColor(raw_ostream &OS, HighlightColor S, bool DisableColors = false,
+            WithColorContext *Context = nullptr);
+  /// To change colors temporarily, use WithColor like this:
+  ///
+  ///   WithColor(OS, raw_ostream::Black) << "text";
+  ///
+  /// If WithColorContext is given, it restores colors properly:
+  ///
+  ///   WithColorContext Ctx;
+  ///   {
+  ///     WithColor Red(OS, raw_ostream::RED, false, false, false, &Ctx);
+  ///     Red << "AAA";
+  ///     {
+  ///        WithColor Blue(OS, raw_ostream::BLUE, false, false, false, &Ctx);
+  ///        Blue << "BBB";
+  ///        // The Blue's destructor is called here and it restores the Red's
+  ///        // color state.
+  ///     }
+  ///     Red << "CCC\n";
+  ///     // The Red's destructor resets the colors to defaults.
+  ///   }
+  ///
+  /// Otherwise, WithColor resets colors in the destructor regardless of the
+  /// previous color. In the example above, "CCC" will be printed in the
+  /// terminal default color, not in red.
+  ///
   /// @param OS The output stream
   /// @param Color ANSI color to use, the special SAVEDCOLOR can be used to
   /// change only the bold attribute, and keep colors untouched
   /// @param Bold Bold/brighter text, default false
   /// @param BG If true, change the background, default: change foreground
   /// @param DisableColors Whether to ignore color changes regardless of -color
-  /// and support in OS
+  /// @param Context A WithColorContext object. Set nullptr if you don't need
+  ///                to restore colors.
   WithColor(raw_ostream &OS,
             raw_ostream::Colors Color = raw_ostream::SAVEDCOLOR,
-            bool Bold = false, bool BG = false, bool DisableColors = false)
-      : OS(OS), DisableColors(DisableColors) {
+            bool Bold = false, bool BG = false, bool DisableColors = false,
+            WithColorContext *Context = nullptr)
+      : OS(OS), DisableColors(DisableColors), Context(Context) {
     changeColor(Color, Bold, BG);
   }
   ~WithColor();
 
   raw_ostream &get() { return OS; }
   operator raw_ostream &() { return OS; }
   template <typename T> WithColor &operator<<(T &O) {
     OS << O;
public:
   static raw_ostream &warning();
   /// Convenience method for printing "note: " to stderr.
   static raw_ostream &note();
   /// Convenience method for printing "remark: " to stderr.
   static raw_ostream &remark();
 
   /// Convenience method for printing "error: " to the given stream.
   static raw_ostream &error(raw_ostream &OS, StringRef Prefix = "",
-                            bool DisableColors = false);
+                            bool DisableColors = false,
+                            WithColorContext *Context = nullptr);
   /// Convenience method for printing "warning: " to the given stream.
   static raw_ostream &warning(raw_ostream &OS, StringRef Prefix = "",
-                              bool DisableColors = false);
+                              bool DisableColors = false,
+                              WithColorContext *Context = nullptr);
   /// Convenience method for printing "note: " to the given stream.
   static raw_ostream &note(raw_ostream &OS, StringRef Prefix = "",
-                           bool DisableColors = false);
+                           bool DisableColors = false,
+                           WithColorContext *Context = nullptr);
   /// Convenience method for printing "remark: " to the given stream.
   static raw_ostream &remark(raw_ostream &OS, StringRef Prefix = "",
-                             bool DisableColors = false);
+                             bool DisableColors = false,
+                             WithColorContext *Context = nullptr);
 
   /// Determine whether colors are displayed.
   bool colorsEnabled();
 
   /// Change the color of text that will be output from this point forward.
   /// @param Color ANSI color to use, the special SAVEDCOLOR can be used to
   /// change only the bold attribute, and keep colors untouched
   /// @param Bold Bold/brighter text, default false

llvm/lib/Support/WithColor.cpp

 
 cl::OptionCategory llvm::ColorCategory("Color Options");
 
 static cl::opt<cl::boolOrDefault>
     UseColor("color", cl::cat(ColorCategory),
              cl::desc("Use colors in output (default=autodetect)"),
              cl::init(cl::BOU_UNSET));
 
-WithColor::WithColor(raw_ostream &OS, HighlightColor Color, bool DisableColors)
-    : OS(OS), DisableColors(DisableColors) {
+WithColor::WithColor(raw_ostream &OS, HighlightColor Color, bool DisableColors,
+                     WithColorContext *Context)
+    : OS(OS), DisableColors(DisableColors), Context(Context) {
   // Detect color from terminal type unless the user passed the --color option.
   if (colorsEnabled()) {
     switch (Color) {
     case HighlightColor::Address:
-      OS.changeColor(raw_ostream::YELLOW);
+      changeColor(raw_ostream::YELLOW);
       break;
     case HighlightColor::String:
-      OS.changeColor(raw_ostream::GREEN);
+      changeColor(raw_ostream::GREEN);
       break;
     case HighlightColor::Tag:
-      OS.changeColor(raw_ostream::BLUE);
+      changeColor(raw_ostream::BLUE);
       break;
     case HighlightColor::Attribute:
-      OS.changeColor(raw_ostream::CYAN);
+      changeColor(raw_ostream::CYAN);
       break;
     case HighlightColor::Enumerator:
-      OS.changeColor(raw_ostream::MAGENTA);
+      changeColor(raw_ostream::MAGENTA);
       break;
     case HighlightColor::Macro:
-      OS.changeColor(raw_ostream::RED);
+      changeColor(raw_ostream::RED);
       break;
     case HighlightColor::Error:
-      OS.changeColor(raw_ostream::RED, true);
+      changeColor(raw_ostream::RED, true);
       break;
     case HighlightColor::Warning:
-      OS.changeColor(raw_ostream::MAGENTA, true);
+      changeColor(raw_ostream::MAGENTA, true);
       break;
     case HighlightColor::Note:
-      OS.changeColor(raw_ostream::BLACK, true);
+      changeColor(raw_ostream::BLACK, true);
       break;
     case HighlightColor::Remark:
-      OS.changeColor(raw_ostream::BLUE, true);
+      changeColor(raw_ostream::BLUE, true);
       break;
     }
   }
 }
 
 raw_ostream &WithColor::error() { return error(errs()); }
 
 raw_ostream &WithColor::warning() { return warning(errs()); }
 
 raw_ostream &WithColor::note() { return note(errs()); }
 
 raw_ostream &WithColor::remark() { return remark(errs()); }
 
 raw_ostream &WithColor::error(raw_ostream &OS, StringRef Prefix,
-                              bool DisableColors) {
+                              bool DisableColors, WithColorContext *Context) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Error, DisableColors).get()
+  return WithColor(OS, HighlightColor::Error, DisableColors, Context).get()
          << "error: ";
 }
 
 raw_ostream &WithColor::warning(raw_ostream &OS, StringRef Prefix,
-                                bool DisableColors) {
+                                bool DisableColors, WithColorContext *Context) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Warning, DisableColors).get()
+  return WithColor(OS, HighlightColor::Warning, DisableColors, Context).get()
          << "warning: ";
 }
 
 raw_ostream &WithColor::note(raw_ostream &OS, StringRef Prefix,
-                             bool DisableColors) {
+                             bool DisableColors, WithColorContext *Context) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Note, DisableColors).get() << "note: ";
+  return WithColor(OS, HighlightColor::Note, DisableColors, Context).get()
+         << "note: ";
 }
 
 raw_ostream &WithColor::remark(raw_ostream &OS, StringRef Prefix,
-                               bool DisableColors) {
+                               bool DisableColors, WithColorContext *Context) {
   if (!Prefix.empty())
     OS << Prefix << ": ";
-  return WithColor(OS, HighlightColor::Remark, DisableColors).get()
+  return WithColor(OS, HighlightColor::Remark, DisableColors, Context).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())
+  if (colorsEnabled()) {
+    if (Context)
+      Context->pushColor(Color, Bold, BG);
+
     OS.changeColor(Color, Bold, BG);
+  }
+
   return *this;
 }
 
 WithColor &WithColor::resetColor() {
   if (colorsEnabled())
     OS.resetColor();
   return *this;
 }
 
-WithColor::~WithColor() { resetColor(); }
+WithColor::~WithColor() {
+  if (colorsEnabled()) {
+    OS.resetColor();
+    if (Context) {
+      Context->popColor();
+      if (!Context->empty()) {
+        const ColorState &State = Context->currentColor();
+        OS.changeColor(State.Color, State.Bold, State.BG);
+      }
+    }
+  }
+}