[lldb] Add a setting to customize the prompt color (#66218)

Users often want to change the look of their prompt and currently the
only way to do that is by using ANSI escape codes in the prompt itself.
This is not only tedious, it also results in extra whitespace because
our Editline wrapper, when computing the cursor column, doesn't ignore
the invisible escape codes.

We already have various *-ansi-{prefix,suffix} settings that allow the
users to customize the color of auto-suggestions and progress events,
using mnemonics like ${ansi.fg.yellow}. This patch brings the same
mechanism to the prompt.

rdar://115390406

7 files changed, 87 insertions, 19 deletions

diff --git a/lldb/include/lldb/Core/Debugger.h b/lldb/include/lldb/Core/Debugger.h
index f0908ee..e4ee948 100644
--- a/lldb/include/lldb/Core/Debugger.h
+++ b/lldb/include/lldb/Core/Debugger.h
@@ -289,6 +289,10 @@ public:
 
   llvm::StringRef GetPrompt() const;
 
+  llvm::StringRef GetPromptAnsiPrefix() const;
+
+  llvm::StringRef GetPromptAnsiSuffix() const;
+
   void SetPrompt(llvm::StringRef p);
   void SetPrompt(const char *) = delete;
 
diff --git a/lldb/include/lldb/Host/Editline.h b/lldb/include/lldb/Host/Editline.h
index 64c38415..ba28968 100644
--- a/lldb/include/lldb/Host/Editline.h
+++ b/lldb/include/lldb/Host/Editline.h
@@ -211,6 +211,14 @@ public:
     m_fix_indentation_callback_chars = indent_chars;
   }
 
+  void SetPromptAnsiPrefix(std::string prefix) {
+    m_prompt_ansi_prefix = std::move(prefix);
+  }
+
+  void SetPromptAnsiSuffix(std::string suffix) {
+    m_prompt_ansi_suffix = std::move(suffix);
+  }
+
   void SetSuggestionAnsiPrefix(std::string prefix) {
     m_suggestion_ansi_prefix = std::move(prefix);
   }
@@ -398,6 +406,8 @@ private:
   CompleteCallbackType m_completion_callback;
   SuggestionCallbackType m_suggestion_callback;
 
+  std::string m_prompt_ansi_prefix;
+  std::string m_prompt_ansi_suffix;
   std::string m_suggestion_ansi_prefix;
   std::string m_suggestion_ansi_suffix;
 
diff --git a/lldb/source/Core/CoreProperties.td b/lldb/source/Core/CoreProperties.td
index c7d6967..9288425 100644
--- a/lldb/source/Core/CoreProperties.td
+++ b/lldb/source/Core/CoreProperties.td
@@ -65,6 +65,14 @@ let Definition = "debugger" in {
     DefaultEnumValue<"OptionValueString::eOptionEncodeCharacterEscapeSequences">,
     DefaultStringValue<"(lldb) ">,
     Desc<"The debugger command line prompt displayed for the user.">;
+  def PromptAnsiPrefix: Property<"prompt-ansi-prefix", "String">,
+    Global,
+    DefaultStringValue<"${ansi.faint}">,
+    Desc<"When in a color-enabled terminal, use the ANSI terminal code specified in this format immediately before the prompt.">;
+  def PromptAnsiSuffix: Property<"prompt-ansi-suffix", "String">,
+    Global,
+    DefaultStringValue<"${ansi.normal}">,
+    Desc<"When in a color-enabled terminal, use the ANSI terminal code specified in this format immediately after the prompt.">;
   def ScriptLanguage: Property<"script-lang", "Enum">,
     Global,
     DefaultEnumValue<"eScriptLanguagePython">,
diff --git a/lldb/source/Core/Debugger.cpp b/lldb/source/Core/Debugger.cpp
index 1e989e1..21f71e4 100644
--- a/lldb/source/Core/Debugger.cpp
+++ b/lldb/source/Core/Debugger.cpp
@@ -236,6 +236,13 @@ Status Debugger::SetPropertyValue(const ExecutionContext *exe_ctx,
       // codes.
       SetPrompt(GetPrompt());
     } else if (property_path ==
+                   g_debugger_properties[ePropertyPromptAnsiPrefix].name ||
+               property_path ==
+                   g_debugger_properties[ePropertyPromptAnsiSuffix].name) {
+      // Prompt colors changed. Ping the prompt so it can reset the ansi
+      // terminal codes.
+      SetPrompt(GetPrompt());
+    } else if (property_path ==
                g_debugger_properties[ePropertyUseSourceCache].name) {
       // use-source-cache changed. Wipe out the cache contents if it was
       // disabled.
@@ -301,6 +308,18 @@ llvm::StringRef Debugger::GetPrompt() const {
       idx, g_debugger_properties[idx].default_cstr_value);
 }
 
+llvm::StringRef Debugger::GetPromptAnsiPrefix() const {
+  const uint32_t idx = ePropertyPromptAnsiPrefix;
+  return GetPropertyAtIndexAs<llvm::StringRef>(
+      idx, g_debugger_properties[idx].default_cstr_value);
+}
+
+llvm::StringRef Debugger::GetPromptAnsiSuffix() const {
+  const uint32_t idx = ePropertyPromptAnsiSuffix;
+  return GetPropertyAtIndexAs<llvm::StringRef>(
+      idx, g_debugger_properties[idx].default_cstr_value);
+}
+
 void Debugger::SetPrompt(llvm::StringRef p) {
   constexpr uint32_t idx = ePropertyPrompt;
   SetPropertyAtIndex(idx, p);
diff --git a/lldb/source/Core/IOHandler.cpp b/lldb/source/Core/IOHandler.cpp
index ac9a9bb..7272e5d 100644
--- a/lldb/source/Core/IOHandler.cpp
+++ b/lldb/source/Core/IOHandler.cpp
@@ -474,8 +474,15 @@ bool IOHandlerEditline::SetPrompt(llvm::StringRef prompt) {
   m_prompt = std::string(prompt);
 
 #if LLDB_ENABLE_LIBEDIT
-  if (m_editline_up)
+  if (m_editline_up) {
     m_editline_up->SetPrompt(m_prompt.empty() ? nullptr : m_prompt.c_str());
+    if (m_debugger.GetUseColor()) {
+      m_editline_up->SetPromptAnsiPrefix(
+          ansi::FormatAnsiTerminalCodes(m_debugger.GetPromptAnsiPrefix()));
+      m_editline_up->SetPromptAnsiSuffix(
+          ansi::FormatAnsiTerminalCodes(m_debugger.GetPromptAnsiSuffix()));
+    }
+  }
 #endif
   return true;
 }
diff --git a/lldb/source/Host/common/Editline.cpp b/lldb/source/Host/common/Editline.cpp
index e84b451..32dd471 100644
--- a/lldb/source/Host/common/Editline.cpp
+++ b/lldb/source/Host/common/Editline.cpp
@@ -53,10 +53,6 @@ int setupterm(char *term, int fildes, int *errret);
 
 /// https://www.ecma-international.org/publications/files/ECMA-ST/Ecma-048.pdf
 #define ESCAPE "\x1b"
-/// Faint, decreased intensity or second colour.
-#define ANSI_FAINT ESCAPE "[2m"
-/// Normal colour or normal intensity (neither bold nor faint).
-#define ANSI_UNFAINT ESCAPE "[0m"
 #define ANSI_CLEAR_BELOW ESCAPE "[J"
 #define ANSI_CLEAR_RIGHT ESCAPE "[K"
 #define ANSI_SET_COLUMN_N ESCAPE "[%dG"
@@ -431,15 +427,13 @@ void Editline::MoveCursor(CursorLocation from, CursorLocation to) {
 void Editline::DisplayInput(int firstIndex) {
   fprintf(m_output_file, ANSI_SET_COLUMN_N ANSI_CLEAR_BELOW, 1);
   int line_count = (int)m_input_lines.size();
-  const char *faint = m_color_prompts ? ANSI_FAINT : "";
-  const char *unfaint = m_color_prompts ? ANSI_UNFAINT : "";
-
   for (int index = firstIndex; index < line_count; index++) {
-    fprintf(m_output_file, "%s"
-                           "%s"
-                           "%s" EditLineStringFormatSpec " ",
-            faint, PromptForIndex(index).c_str(), unfaint,
-            m_input_lines[index].c_str());
+    fprintf(m_output_file,
+            "%s"
+            "%s"
+            "%s" EditLineStringFormatSpec " ",
+            m_prompt_ansi_prefix.c_str(), PromptForIndex(index).c_str(),
+            m_prompt_ansi_suffix.c_str(), m_input_lines[index].c_str());
     if (index < line_count - 1)
       fprintf(m_output_file, "\n");
   }
@@ -548,14 +542,16 @@ unsigned char Editline::RecallHistory(HistoryOperation op) {
 int Editline::GetCharacter(EditLineGetCharType *c) {
   const LineInfoW *info = el_wline(m_editline);
 
-  // Paint a faint version of the desired prompt over the version libedit draws
-  // (will only be requested if colors are supported)
+  // Paint a ANSI formatted version of the desired prompt over the version
+  // libedit draws. (will only be requested if colors are supported)
   if (m_needs_prompt_repaint) {
     MoveCursor(CursorLocation::EditingCursor, CursorLocation::EditingPrompt);
-    fprintf(m_output_file, "%s"
-                           "%s"
-                           "%s",
-            ANSI_FAINT, Prompt(), ANSI_UNFAINT);
+    fprintf(m_output_file,
+            "%s"
+            "%s"
+            "%s",
+            m_prompt_ansi_prefix.c_str(), Prompt(),
+            m_prompt_ansi_suffix.c_str());
     MoveCursor(CursorLocation::EditingPrompt, CursorLocation::EditingCursor);
     m_needs_prompt_repaint = false;
   }
diff --git a/lldb/test/API/terminal/TestEditline.py b/lldb/test/API/terminal/TestEditline.py
index 4e766cf..0c6d16f 100644
--- a/lldb/test/API/terminal/TestEditline.py
+++ b/lldb/test/API/terminal/TestEditline.py
@@ -55,3 +55,27 @@ class EditlineTest(PExpectTest):
         # Prompt: 🐛 _
         # Column: 1..4
         self.child.expect(re.escape("🐛 \x1b[0m\x1b[4G"))
+
+    @skipIfAsan
+    @skipIfEditlineSupportMissing
+    def test_prompt_color(self):
+        """Test that we can change the prompt color with prompt-ansi-prefix."""
+        self.launch(use_colors=True)
+        self.child.send('settings set prompt-ansi-prefix "${ansi.fg.red}"\n')
+        # Make sure this change is reflected immediately. Check that the color
+        # is set (31) and the cursor position (8) is correct.
+        # Prompt: (lldb) _
+        # Column: 1....6.8
+        self.child.expect(re.escape("\x1b[31m(lldb) \x1b[0m\x1b[8G"))
+
+    @skipIfAsan
+    @skipIfEditlineSupportMissing
+    def test_prompt_no_color(self):
+        """Test that prompt-ansi-prefix doesn't color the prompt when colors are off."""
+        self.launch(use_colors=False)
+        self.child.send('settings set prompt-ansi-prefix "${ansi.fg.red}"\n')
+        # Send foo so we can match the newline before the prompt and the foo
+        # after the prompt.
+        self.child.send("foo")
+        # Check that there are no escape codes.
+        self.child.expect(re.escape("\n(lldb) foo"))