[lldb] Support format string in the prompt

[JDevlieghere]

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

[llvmbot]

@llvm/pr-subscribers-lldb

Author: Jonas Devlieghere (JDevlieghere)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/123430.diff

4 Files Affected:

• (modified) lldb/include/lldb/Utility/AnsiTerminal.h (+38)
• (modified) lldb/source/Host/common/Editline.cpp (+3-1)
• (modified) lldb/test/API/terminal/TestEditline.py (+16)
• (modified) lldb/unittests/Utility/AnsiTerminalTest.cpp (+15)

diff --git a/lldb/include/lldb/Utility/AnsiTerminal.h b/lldb/include/lldb/Utility/AnsiTerminal.h
index 67795971d2ca89..22601e873dfe9e 100644
--- a/lldb/include/lldb/Utility/AnsiTerminal.h
+++ b/lldb/include/lldb/Utility/AnsiTerminal.h
@@ -171,7 +171,45 @@ 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;
+
+    auto end = llvm::StringRef::npos;
+    for (size_t i = 0; i < right.size(); i++) {
+      char c = right[i];
+      if (c == 'm' || c == 'G') {
+        end = i;
+        break;
+      }
+      if (isdigit(c) || c == ';')
+        continue;
+
+      break;
+    }
+
+    // ANSI_ESC_END not found.
+    if (end != llvm::StringRef::npos) {
+      str = right.substr(end + 1);
+      continue;
+    }
+
+    stripped += ANSI_ESC_START;
+    str = right;
+  }
+  return stripped;
 }
+
+} // namespace ansi
 } // namespace lldb_private
 
 #endif
diff --git a/lldb/source/Host/common/Editline.cpp b/lldb/source/Host/common/Editline.cpp
index 6e35b15d69651d..d31fe3af946b37 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) {
diff --git a/lldb/test/API/terminal/TestEditline.py b/lldb/test/API/terminal/TestEditline.py
index aa7d827e599441..40f8a50e06bf6d 100644
--- a/lldb/test/API/terminal/TestEditline.py
+++ b/lldb/test/API/terminal/TestEditline.py
@@ -69,6 +69,22 @@ def test_prompt_color(self):
         # Column: 1....6.8
         self.child.expect(re.escape("\x1b[31m(lldb) \x1b[0m\x1b[8G"))
 
+    @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.child.send('settings set prompt-ansi-prefix ""\n')
+        self.child.send('settings set prompt-ansi-suffix ""\n')
+        self.child.send('settings set prompt "${ansi.fg.red}(lldb)${ansi.normal} "\n')
+        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 foo"))
+
     @skipIfAsan
     @skipIfEditlineSupportMissing
     def test_prompt_no_color(self):
diff --git a/lldb/unittests/Utility/AnsiTerminalTest.cpp b/lldb/unittests/Utility/AnsiTerminalTest.cpp
index a6dbfd61061420..1ba9565c3f6af3 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"));
 }

[JDevlieghere]

I need the same functionality for the Statusline (#121860) so I've factored this out into its own PR.