[lldb] Support format string in the prompt (#123430)

Implement ansi::StripAnsiTerminalCodes and fix a long standing bug where
using format strings in lldb's prompt resulted in an incorrect prompt
column width.

7 files changed, 82 insertions, 24 deletions

diff --git a/lldb/include/lldb/Host/Editline.h b/lldb/include/lldb/Host/Editline.h
index 26deba3..27b8638 100644
--- a/lldb/include/lldb/Host/Editline.h
+++ b/lldb/include/lldb/Host/Editline.h
@@ -152,7 +152,7 @@ using namespace line_editor;
 class Editline {
 public:
   Editline(const char *editor_name, FILE *input_file, FILE *output_file,
-           FILE *error_file, std::recursive_mutex &output_mutex);
+           FILE *error_file, bool color, std::recursive_mutex &output_mutex);
 
   ~Editline();
 
@@ -212,19 +212,23 @@ public:
   }
 
   void SetPromptAnsiPrefix(std::string prefix) {
-    m_prompt_ansi_prefix = std::move(prefix);
+    if (m_color)
+      m_prompt_ansi_prefix = std::move(prefix);
   }
 
   void SetPromptAnsiSuffix(std::string suffix) {
-    m_prompt_ansi_suffix = std::move(suffix);
+    if (m_color)
+      m_prompt_ansi_suffix = std::move(suffix);
   }
 
   void SetSuggestionAnsiPrefix(std::string prefix) {
-    m_suggestion_ansi_prefix = std::move(prefix);
+    if (m_color)
+      m_suggestion_ansi_prefix = std::move(prefix);
   }
 
   void SetSuggestionAnsiSuffix(std::string suffix) {
-    m_suggestion_ansi_suffix = std::move(suffix);
+    if (m_color)
+      m_suggestion_ansi_suffix = std::move(suffix);
   }
 
   /// Prompts for and reads a single line of user input.
@@ -400,6 +404,7 @@ private:
   CompleteCallbackType m_completion_callback;
   SuggestionCallbackType m_suggestion_callback;
 
+  bool m_color;
   std::string m_prompt_ansi_prefix;
   std::string m_prompt_ansi_suffix;
   std::string m_suggestion_ansi_prefix;
diff --git a/lldb/include/lldb/Utility/AnsiTerminal.h b/lldb/include/lldb/Utility/AnsiTerminal.h
index 6779597..1939c49 100644
--- a/lldb/include/lldb/Utility/AnsiTerminal.h
+++ b/lldb/include/lldb/Utility/AnsiTerminal.h
@@ -171,7 +171,32 @@ inline std::string FormatAnsiTerminalCodes(llvm::StringRef format,
   }
   return fmt;
 }
+
+inline std::string StripAnsiTerminalCodes(llvm::StringRef str) {
+  std::string stripped;
+  while (!str.empty()) {
+    llvm::StringRef left, right;
+
+    std::tie(left, right) = str.split(ANSI_ESC_START);
+    stripped += left;
+
+    // ANSI_ESC_START not found.
+    if (left == str && right.empty())
+      break;
+
+    size_t end = right.find_first_not_of("0123456789;");
+    if (end < right.size() && (right[end] == 'm' || right[end] == 'G')) {
+      str = right.substr(end + 1);
+    } else {
+      // ANSI_ESC_END not found.
+      stripped += ANSI_ESC_START;
+      str = right;
+    }
+  }
+  return stripped;
 }
+
+} // namespace ansi
 } // namespace lldb_private
 
 #endif
diff --git a/lldb/source/Core/IOHandler.cpp b/lldb/source/Core/IOHandler.cpp
index 695c248..ca06b52 100644
--- a/lldb/source/Core/IOHandler.cpp
+++ b/lldb/source/Core/IOHandler.cpp
@@ -264,7 +264,7 @@ IOHandlerEditline::IOHandlerEditline(
   if (use_editline) {
     m_editline_up = std::make_unique<Editline>(editline_name, GetInputFILE(),
                                                GetOutputFILE(), GetErrorFILE(),
-                                               GetOutputMutex());
+                                               m_color, GetOutputMutex());
     m_editline_up->SetIsInputCompleteCallback(
         [this](Editline *editline, StringList &lines) {
           return this->IsInputCompleteCallback(editline, lines);
@@ -278,12 +278,10 @@ IOHandlerEditline::IOHandlerEditline(
       m_editline_up->SetSuggestionCallback([this](llvm::StringRef line) {
         return this->SuggestionCallback(line);
       });
-      if (m_color) {
-        m_editline_up->SetSuggestionAnsiPrefix(ansi::FormatAnsiTerminalCodes(
-            debugger.GetAutosuggestionAnsiPrefix()));
-        m_editline_up->SetSuggestionAnsiSuffix(ansi::FormatAnsiTerminalCodes(
-            debugger.GetAutosuggestionAnsiSuffix()));
-      }
+      m_editline_up->SetSuggestionAnsiPrefix(ansi::FormatAnsiTerminalCodes(
+          debugger.GetAutosuggestionAnsiPrefix()));
+      m_editline_up->SetSuggestionAnsiSuffix(ansi::FormatAnsiTerminalCodes(
+          debugger.GetAutosuggestionAnsiSuffix()));
     }
     // See if the delegate supports fixing indentation
     const char *indent_chars = delegate.IOHandlerGetFixIndentationCharacters();
@@ -478,12 +476,10 @@ bool IOHandlerEditline::SetPrompt(llvm::StringRef prompt) {
 #if LLDB_ENABLE_LIBEDIT
   if (m_editline_up) {
     m_editline_up->SetPrompt(m_prompt.empty() ? nullptr : m_prompt.c_str());
-    if (m_color) {
-      m_editline_up->SetPromptAnsiPrefix(
-          ansi::FormatAnsiTerminalCodes(m_debugger.GetPromptAnsiPrefix()));
-      m_editline_up->SetPromptAnsiSuffix(
-          ansi::FormatAnsiTerminalCodes(m_debugger.GetPromptAnsiSuffix()));
-    }
+    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 6e35b15..73da1d84 100644
--- a/lldb/source/Host/common/Editline.cpp
+++ b/lldb/source/Host/common/Editline.cpp
@@ -14,6 +14,7 @@
 #include "lldb/Host/Editline.h"
 #include "lldb/Host/FileSystem.h"
 #include "lldb/Host/Host.h"
+#include "lldb/Utility/AnsiTerminal.h"
 #include "lldb/Utility/CompletionRequest.h"
 #include "lldb/Utility/FileSpec.h"
 #include "lldb/Utility/LLDBAssert.h"
@@ -85,7 +86,8 @@ bool IsOnlySpaces(const EditLineStringType &content) {
 }
 
 static size_t ColumnWidth(llvm::StringRef str) {
-  return llvm::sys::locale::columnWidth(str);
+  std::string stripped = ansi::StripAnsiTerminalCodes(str);
+  return llvm::sys::locale::columnWidth(stripped);
 }
 
 static int GetOperation(HistoryOperation op) {
@@ -610,7 +612,7 @@ int Editline::GetCharacter(EditLineGetCharType *c) {
 }
 
 const char *Editline::Prompt() {
-  if (!m_prompt_ansi_prefix.empty() || !m_prompt_ansi_suffix.empty())
+  if (m_color)
     m_needs_prompt_repaint = true;
   return m_current_prompt.c_str();
 }
@@ -1471,11 +1473,11 @@ Editline *Editline::InstanceFor(EditLine *editline) {
 }
 
 Editline::Editline(const char *editline_name, FILE *input_file,
-                   FILE *output_file, FILE *error_file,
+                   FILE *output_file, FILE *error_file, bool color,
                    std::recursive_mutex &output_mutex)
     : m_editor_status(EditorStatus::Complete), m_input_file(input_file),
       m_output_file(output_file), m_error_file(error_file),
-      m_input_connection(fileno(input_file), false),
+      m_input_connection(fileno(input_file), false), m_color(color),
       m_output_mutex(output_mutex) {
   // Get a shared history instance
   m_editor_name = (editline_name == nullptr) ? "lldb-tmp" : editline_name;
diff --git a/lldb/test/API/terminal/TestEditline.py b/lldb/test/API/terminal/TestEditline.py
index aa7d827..ddaa441 100644
--- a/lldb/test/API/terminal/TestEditline.py
+++ b/lldb/test/API/terminal/TestEditline.py
@@ -2,7 +2,6 @@
 Test that the lldb editline handling is configured correctly.
 """
 
-
 import lldb
 from lldbsuite.test.decorators import *
 from lldbsuite.test.lldbtest import *
@@ -71,6 +70,22 @@ class EditlineTest(PExpectTest):
 
     @skipIfAsan
     @skipIfEditlineSupportMissing
+    def test_prompt_format_color(self):
+        """Test that we can change the prompt color with a format string."""
+        self.launch(use_colors=True)
+        # Clear the prefix and suffix setting to simplify the output.
+        self.expect('settings set prompt-ansi-prefix ""')
+        self.expect('settings set prompt-ansi-suffix ""')
+        self.expect('settings set prompt "${ansi.fg.red}(lldb) ${ansi.normal}"')
+        self.child.send("foo")
+        # 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[8Gfoo"))
+
+    @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)
diff --git a/lldb/unittests/Editline/EditlineTest.cpp b/lldb/unittests/Editline/EditlineTest.cpp
index 333ad77..1327b58 100644
--- a/lldb/unittests/Editline/EditlineTest.cpp
+++ b/lldb/unittests/Editline/EditlineTest.cpp
@@ -118,7 +118,7 @@ EditlineAdapter::EditlineAdapter()
   // Create an Editline instance.
   _editline_sp.reset(new lldb_private::Editline(
       "gtest editor", *_el_secondary_file, *_el_secondary_file,
-      *_el_secondary_file, output_mutex));
+      *_el_secondary_file, /*color=*/false, output_mutex));
   _editline_sp->SetPrompt("> ");
 
   // Hookup our input complete callback.
diff --git a/lldb/unittests/Utility/AnsiTerminalTest.cpp b/lldb/unittests/Utility/AnsiTerminalTest.cpp
index a6dbfd6..1ba9565 100644
--- a/lldb/unittests/Utility/AnsiTerminalTest.cpp
+++ b/lldb/unittests/Utility/AnsiTerminalTest.cpp
@@ -16,16 +16,21 @@ TEST(AnsiTerminal, Empty) { EXPECT_EQ("", ansi::FormatAnsiTerminalCodes("")); }
 
 TEST(AnsiTerminal, WhiteSpace) {
   EXPECT_EQ(" ", ansi::FormatAnsiTerminalCodes(" "));
+  EXPECT_EQ(" ", ansi::StripAnsiTerminalCodes(" "));
 }
 
 TEST(AnsiTerminal, AtEnd) {
   EXPECT_EQ("abc\x1B[30m",
             ansi::FormatAnsiTerminalCodes("abc${ansi.fg.black}"));
+
+  EXPECT_EQ("abc", ansi::StripAnsiTerminalCodes("abc\x1B[30m"));
 }
 
 TEST(AnsiTerminal, AtStart) {
   EXPECT_EQ("\x1B[30mabc",
             ansi::FormatAnsiTerminalCodes("${ansi.fg.black}abc"));
+
+  EXPECT_EQ("abc", ansi::StripAnsiTerminalCodes("\x1B[30mabc"));
 }
 
 TEST(AnsiTerminal, KnownPrefix) {
@@ -45,10 +50,20 @@ TEST(AnsiTerminal, Incomplete) {
 TEST(AnsiTerminal, Twice) {
   EXPECT_EQ("\x1B[30m\x1B[31mabc",
             ansi::FormatAnsiTerminalCodes("${ansi.fg.black}${ansi.fg.red}abc"));
+
+  EXPECT_EQ("abc", ansi::StripAnsiTerminalCodes("\x1B[30m\x1B[31mabc"));
 }
 
 TEST(AnsiTerminal, Basic) {
   EXPECT_EQ(
       "abc\x1B[31mabc\x1B[0mabc",
       ansi::FormatAnsiTerminalCodes("abc${ansi.fg.red}abc${ansi.normal}abc"));
+
+  EXPECT_EQ("abcabcabc",
+            ansi::StripAnsiTerminalCodes("abc\x1B[31mabc\x1B[0mabc"));
+}
+
+TEST(AnsiTerminal, InvalidEscapeCode) {
+  EXPECT_EQ("abc\x1B[31kabcabc",
+            ansi::StripAnsiTerminalCodes("abc\x1B[31kabc\x1B[0mabc"));
 }