From eef0c3864bc8f81c95d85318bda779e806d78680 Mon Sep 17 00:00:00 2001
From: Jonas Devlieghere <jonas@devlieghere.com>
Date: Wed, 13 Sep 2023 18:50:39 -0700
Subject: [PATCH] [lldb] Support Unicode in the prompt

Account for Unicode when computing the prompt column width. Previously,
the string length (i.e. number of bytes) rather than the width of the
Glyph was used to compute the cursor position. The result was that the
cursor would be offset to the right when using a prompt containing
Unicode.
---
 lldb/include/lldb/Host/Editline.h             |  5 ++--
 .../Python/lldbsuite/test/lldbpexpect.py      |  2 ++
 lldb/source/Host/common/Editline.cpp          | 25 ++++++++++++-------
 lldb/test/API/terminal/TestEditline.py        | 11 ++++++++
 4 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/lldb/include/lldb/Host/Editline.h b/lldb/include/lldb/Host/Editline.h
index 35fd179abb509..64c384151facc 100644
--- a/lldb/include/lldb/Host/Editline.h
+++ b/lldb/include/lldb/Host/Editline.h
@@ -248,9 +248,8 @@ class Editline {
   void SetCurrentLine(int line_index);
 
   /// Determines the width of the prompt in characters.  The width is guaranteed
-  /// to be the same for
-  /// all lines of the current multi-line session.
-  int GetPromptWidth();
+  /// to be the same for all lines of the current multi-line session.
+  size_t GetPromptWidth();
 
   /// Returns true if the underlying EditLine session's keybindings are
   /// Emacs-based, or false if
diff --git a/lldb/packages/Python/lldbsuite/test/lldbpexpect.py b/lldb/packages/Python/lldbsuite/test/lldbpexpect.py
index 85102db7d0a05..9fedfa70c3f65 100644
--- a/lldb/packages/Python/lldbsuite/test/lldbpexpect.py
+++ b/lldb/packages/Python/lldbsuite/test/lldbpexpect.py
@@ -26,6 +26,7 @@ def launch(
         dimensions=None,
         run_under=None,
         post_spawn=None,
+        encoding=None,
         use_colors=False,
     ):
         logfile = getattr(sys.stdout, "buffer", sys.stdout) if self.TraceOn() else None
@@ -58,6 +59,7 @@ def launch(
             timeout=timeout,
             dimensions=dimensions,
             env=env,
+            encoding=encoding,
         )
         self.child.ptyproc.delayafterclose = timeout / 10
         self.child.ptyproc.delayafterterminate = timeout / 10
diff --git a/lldb/source/Host/common/Editline.cpp b/lldb/source/Host/common/Editline.cpp
index 544804db3879e..e84b451435a99 100644
--- a/lldb/source/Host/common/Editline.cpp
+++ b/lldb/source/Host/common/Editline.cpp
@@ -25,6 +25,7 @@
 #include "lldb/Utility/Timeout.h"
 
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Locale.h"
 #include "llvm/Support/Threading.h"
 
 using namespace lldb_private;
@@ -101,6 +102,10 @@ bool IsOnlySpaces(const EditLineStringType &content) {
   return true;
 }
 
+static size_t ColumnWidth(llvm::StringRef str) {
+  return llvm::sys::locale::columnWidth(str);
+}
+
 static int GetOperation(HistoryOperation op) {
   // The naming used by editline for the history operations is counter
   // intuitive to how it's used in LLDB's editline implementation.
@@ -328,14 +333,16 @@ std::string Editline::PromptForIndex(int line_index) {
   std::string continuation_prompt = prompt;
   if (m_set_continuation_prompt.length() > 0) {
     continuation_prompt = m_set_continuation_prompt;
-
     // Ensure that both prompts are the same length through space padding
-    while (continuation_prompt.length() < prompt.length()) {
-      continuation_prompt += ' ';
-    }
-    while (prompt.length() < continuation_prompt.length()) {
-      prompt += ' ';
-    }
+    const size_t prompt_width = ColumnWidth(prompt);
+    const size_t cont_prompt_width = ColumnWidth(continuation_prompt);
+    const size_t padded_prompt_width =
+        std::max(prompt_width, cont_prompt_width);
+    if (prompt_width < padded_prompt_width)
+      prompt += std::string(padded_prompt_width - prompt_width, ' ');
+    else if (cont_prompt_width < padded_prompt_width)
+      continuation_prompt +=
+          std::string(padded_prompt_width - cont_prompt_width, ' ');
   }
 
   if (use_line_numbers) {
@@ -353,7 +360,7 @@ void Editline::SetCurrentLine(int line_index) {
   m_current_prompt = PromptForIndex(line_index);
 }
 
-int Editline::GetPromptWidth() { return (int)PromptForIndex(0).length(); }
+size_t Editline::GetPromptWidth() { return ColumnWidth(PromptForIndex(0)); }
 
 bool Editline::IsEmacs() {
   const char *editor;
@@ -441,7 +448,7 @@ void Editline::DisplayInput(int firstIndex) {
 int Editline::CountRowsForLine(const EditLineStringType &content) {
   std::string prompt =
       PromptForIndex(0); // Prompt width is constant during an edit session
-  int line_length = (int)(content.length() + prompt.length());
+  int line_length = (int)(content.length() + ColumnWidth(prompt));
   return (line_length / m_terminal_width) + 1;
 }
 
diff --git a/lldb/test/API/terminal/TestEditline.py b/lldb/test/API/terminal/TestEditline.py
index 87fb3cba610c8..4e766cff286f6 100644
--- a/lldb/test/API/terminal/TestEditline.py
+++ b/lldb/test/API/terminal/TestEditline.py
@@ -44,3 +44,14 @@ def test_left_right_arrow(self):
             )
 
         self.quit()
+
+    @skipIfAsan
+    @skipIfEditlineSupportMissing
+    def test_prompt_unicode(self):
+        """Test that we can use Unicode in the LLDB prompt."""
+        self.launch(use_colors=True, encoding="utf-8")
+        self.child.send('settings set prompt "🐛 "\n')
+        # Check that the cursor is at position 4 ([4G)
+        # Prompt: 🐛 _
+        # Column: 1..4
+        self.child.expect(re.escape("🐛 \x1b[0m\x1b[4G"))