From 84777fbe62d7cb8f4cfd57f38cd7f44a7fe7f331 Mon Sep 17 00:00:00 2001
From: Ali Mohammad Pur <ali.mpfard@gmail.com>
Date: Thu, 31 Aug 2023 13:34:21 +0330
Subject: [PATCH] Shell: Allow the user to set the prompt using PROMPT()

This allows the prompt to be dynamically configurable, making it
possible to display various bits of information in the prompt.
---
 Userland/Shell/Builtin.cpp | 22 ++++++++++++++++++++++
 Userland/Shell/Shell.cpp   | 16 ++++++++++++++++
 Userland/Shell/Shell.h     |  8 +++++++-
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/Userland/Shell/Builtin.cpp b/Userland/Shell/Builtin.cpp
index edd772ebe83..9e60e0d5aa7 100644
--- a/Userland/Shell/Builtin.cpp
+++ b/Userland/Shell/Builtin.cpp
@@ -1967,6 +1967,28 @@ ErrorOr<int> Shell::builtin_run_with_env(Main::Arguments arguments)
     return exit_code;
 }
 
+ErrorOr<int> Shell::builtin_shell_set_active_prompt(Main::Arguments arguments)
+{
+    StringView new_prompt;
+
+    Core::ArgsParser parser;
+    parser.add_positional_argument(new_prompt, "New prompt text", "prompt", Core::ArgsParser::Required::Yes);
+
+    if (!parser.parse(arguments, Core::ArgsParser::FailureBehavior::Ignore))
+        return 1;
+
+    if (!m_editor) {
+        warnln("shell_set_active_prompt: No active prompt");
+        return 1;
+    }
+
+    if (m_editor->is_editing())
+        m_editor->set_prompt(new_prompt);
+    else
+        m_next_scheduled_prompt_text = new_prompt;
+    return 0;
+}
+
 bool Shell::has_builtin(StringView name) const
 {
     if (name == ":"sv || (m_in_posix_mode && name == "."sv))
diff --git a/Userland/Shell/Shell.cpp b/Userland/Shell/Shell.cpp
index 7b082df2e68..4d582322e1b 100644
--- a/Userland/Shell/Shell.cpp
+++ b/Userland/Shell/Shell.cpp
@@ -79,6 +79,9 @@ void Shell::print_path(StringView path)
 
 DeprecatedString Shell::prompt() const
 {
+    if (m_next_scheduled_prompt_text.has_value())
+        return m_next_scheduled_prompt_text.release_value();
+
     auto build_prompt = [&]() -> DeprecatedString {
         auto* ps1 = getenv("PROMPT");
         if (!ps1) {
@@ -2086,9 +2089,22 @@ void Shell::setup_keybinds()
     });
 }
 
+void Shell::set_user_prompt()
+{
+    if (!has_function("PROMPT"sv))
+        return;
+
+    if (!m_prompt_command_node)
+        m_prompt_command_node = Parser { "shell_set_active_prompt -- ${join \"\\n\" $(PROMPT)}"sv }.parse();
+
+    (void)m_prompt_command_node->run(this);
+}
+
 bool Shell::read_single_line()
 {
     while (true) {
+        set_user_prompt();
+
         restore_ios();
         bring_cursor_to_beginning_of_a_line();
         m_editor->initialize();
diff --git a/Userland/Shell/Shell.h b/Userland/Shell/Shell.h
index 58f712dd952..e47b5d8a7a1 100644
--- a/Userland/Shell/Shell.h
+++ b/Userland/Shell/Shell.h
@@ -60,7 +60,8 @@
     __ENUMERATE_SHELL_BUILTIN(continue, OnlyInPOSIXMode)     \
     __ENUMERATE_SHELL_BUILTIN(read, OnlyInPOSIXMode)         \
     __ENUMERATE_SHELL_BUILTIN(run_with_env, OnlyInPOSIXMode) \
-    __ENUMERATE_SHELL_BUILTIN(argsparser_parse, InAllModes)
+    __ENUMERATE_SHELL_BUILTIN(argsparser_parse, InAllModes)  \
+    __ENUMERATE_SHELL_BUILTIN(shell_set_active_prompt, InAllModes)
 
 #define ENUMERATE_SHELL_OPTIONS()                                                                                    \
     __ENUMERATE_SHELL_OPTION(inline_exec_keep_empty_segments, false, "Keep empty segments in inline execute $(...)") \
@@ -422,6 +423,8 @@ private:
 
     void timer_event(Core::TimerEvent&) override;
 
+    void set_user_prompt();
+
     bool is_allowed_to_modify_termios(const AST::Command&) const;
 
     void bring_cursor_to_beginning_of_a_line() const;
@@ -510,6 +513,9 @@ private:
     Optional<size_t> m_history_autosave_time;
 
     StackInfo m_completion_stack_info;
+
+    RefPtr<AST::Node> m_prompt_command_node;
+    mutable Optional<DeprecatedString> m_next_scheduled_prompt_text;
 };
 
 [[maybe_unused]] static constexpr bool is_word_character(char c)