From 11f0ace9b6bf7cf618a42de7e37d9608d2d86d99 Mon Sep 17 00:00:00 2001
From: Maxime Coste <mawww@kakoune.org>
Date: Tue, 7 Nov 2023 18:01:54 +1100
Subject: [PATCH] Make shell-script-candidates completer run in the background

Read output from the script as it comes and update the candidate
list progressively.

Disable updating of the list when a completion has been explicitely
selected.
---
 doc/pages/changelog.asciidoc                  |  3 +
 src/commands.cc                               | 72 +++++++++++++------
 src/input_handler.cc                          | 26 +++++--
 src/input_handler.hh                          |  2 +
 src/main.cc                                   |  1 +
 src/shell_manager.cc                          | 64 ++++++++---------
 src/shell_manager.hh                          | 23 ++++++
 src/unique_descriptor.hh                      | 23 ++++++
 .../cmd                                       |  2 +-
 .../rc                                        |  1 +
 .../script                                    |  6 ++
 test/shell/prompt-shell-script-candidates/cmd |  2 +-
 test/shell/prompt-shell-script-candidates/rc  |  1 +
 .../prompt-shell-script-candidates/script     |  7 ++
 14 files changed, 171 insertions(+), 62 deletions(-)
 create mode 100644 src/unique_descriptor.hh
 create mode 100644 test/regression/0-nothing-selected-on-prompt-initial-shift-tab/script
 create mode 100644 test/shell/prompt-shell-script-candidates/script

diff --git a/doc/pages/changelog.asciidoc b/doc/pages/changelog.asciidoc
index b5f937e3a..bc0a749e2 100644
--- a/doc/pages/changelog.asciidoc
+++ b/doc/pages/changelog.asciidoc
@@ -5,6 +5,9 @@ released versions.
 
 == Development version
 
+* `shell-script-candidates` completion now runs the script asynchronously
+  while displaying and updating results live.
+
 * `%val{window_range}` elements are now emitted as different strings
 
 * `+` only duplicates identical selections a single time to avoid surprising
diff --git a/src/commands.cc b/src/commands.cc
index 92e11e1b4..55becf4ad 100644
--- a/src/commands.cc
+++ b/src/commands.cc
@@ -280,53 +280,85 @@ struct ShellCandidatesCompleter
                              Completions::Flags flags = Completions::Flags::None)
       : m_shell_script{std::move(shell_script)}, m_flags(flags) {}
 
+    ShellCandidatesCompleter(const ShellCandidatesCompleter& other) : m_shell_script{other.m_shell_script}, m_flags(other.m_flags) {}
+    ShellCandidatesCompleter& operator=(const ShellCandidatesCompleter& other) {  m_shell_script = other.m_shell_script; m_flags = other.m_flags; return *this; }
+
     Completions operator()(const Context& context, CompletionFlags flags,
                            CommandParameters params, size_t token_to_complete,
                            ByteCount pos_in_token)
     {
-        if (m_token != token_to_complete)
+        if (m_last_token != token_to_complete)
         {
             ShellContext shell_context{
                 params,
                 { { "token_to_complete", to_string(token_to_complete) } }
             };
-            String output = ShellManager::instance().eval(m_shell_script, context, {},
-                                                          ShellManager::Flags::WaitForStdout,
-                                                          shell_context).first;
+            m_running_script.emplace(ShellManager::instance().spawn(m_shell_script, context, false, shell_context));
+            m_watcher.emplace((int)m_running_script->out, FdEvents::Read, EventMode::Urgent,
+                              [this, &input_handler=context.input_handler()](auto&&... args) { read_candidates(input_handler); });
             m_candidates.clear();
-            for (auto c : output | split<StringView>('\n')
-                                 | filter([](auto s) { return not s.empty(); }))
-                m_candidates.emplace_back(c.str(), used_letters(c));
-            m_token = token_to_complete;
+            m_last_token = token_to_complete;
+        }
+        return rank_candidates(params[token_to_complete].substr(0, pos_in_token));
+    }
+
+private:
+    void read_candidates(InputHandler& input_handler)
+    {
+        char buffer[2048];
+        bool closed = false;
+        int fd = (int)m_running_script->out;
+        while (fd_readable(fd))
+        {
+            int size = read(fd, buffer, sizeof(buffer));
+            if (size == 0)
+            {
+                closed = true;
+                break;
+            }
+            m_stdout_buffer.insert(m_stdout_buffer.end(), buffer, buffer + size);
         }
 
-        StringView query = params[token_to_complete].substr(0, pos_in_token);
+        auto end = closed ? m_stdout_buffer.end() : find(m_stdout_buffer | reverse(), '\n').base();
+        for (auto c : ArrayView(m_stdout_buffer.begin(), end) | split<StringView>('\n')
+                                                              | filter([](auto s) { return not s.empty(); }))
+            m_candidates.emplace_back(c.str(), used_letters(c));
+        m_stdout_buffer.erase(m_stdout_buffer.begin(), end);
+
+        input_handler.refresh_ifn();
+        if (not closed)
+            return;
+
+        m_running_script.reset();
+        m_watcher.reset();
+    }
+
+    Completions rank_candidates(StringView query)
+    {
         UsedLetters query_letters = used_letters(query);
-        Vector<RankedMatch> matches;
-        for (const auto& candidate : m_candidates)
-        {
-            if (RankedMatch match{candidate.first, candidate.second, query, query_letters})
-                matches.push_back(match);
-        }
+        auto matches = m_candidates | transform([&](const auto& c) { return RankedMatch{c.first, c.second, query, query_letters}; })
+                                    | filter([](const auto& m) { return (bool)m; })
+                                    | gather<Vector<RankedMatch>>();
 
         constexpr size_t max_count = 100;
         CandidateList res;
         // Gather best max_count matches
-        for_n_best(matches, max_count, [](auto& lhs, auto& rhs) { return rhs < lhs; },
-                   [&] (const RankedMatch& m) {
+        for_n_best(matches, max_count, [](auto& lhs, auto& rhs) { return rhs < lhs; }, [&] (const RankedMatch& m) {
             if (not res.empty() and res.back() == m.candidate())
                 return false;
             res.push_back(m.candidate().str());
             return true;
         });
 
-        return Completions{0_byte, pos_in_token, std::move(res), m_flags};
+        return Completions{0_byte, query.length(), std::move(res), m_flags};
     }
 
-private:
     String m_shell_script;
+    Vector<char, MemoryDomain::Completion> m_stdout_buffer;
+    Optional<Shell> m_running_script;
+    Optional<FDWatcher> m_watcher;
     Vector<std::pair<String, UsedLetters>, MemoryDomain::Completion> m_candidates;
-    int m_token = -1;
+    int m_last_token = -1;
     Completions::Flags m_flags;
 };
 
diff --git a/src/input_handler.cc b/src/input_handler.cc
index f1eeb58f5..9bc495a2d 100644
--- a/src/input_handler.cc
+++ b/src/input_handler.cc
@@ -37,6 +37,8 @@ public:
     virtual void on_enabled(bool from_pop) {}
     virtual void on_disabled(bool from_push) {}
 
+    virtual void refresh_ifn() {}
+
     bool enabled() const { return &m_input_handler.current_mode() == this; }
     Context& context() const { return m_input_handler.context(); }
 
@@ -1048,6 +1050,19 @@ public:
             m_idle_timer.set_next_date(Clock::now() + get_idle_timeout(context()));
     }
 
+    void refresh_ifn()
+    {
+        bool explicit_completion_selected = m_current_completion != -1 and
+            (not m_prefix_in_completions or m_current_completion != m_completions.candidates.size() - 1);
+        if (not enabled() or (context().flags() & Context::Flags::Draft) or explicit_completion_selected)
+            return;
+
+        if (auto next_date = Clock::now() + get_idle_timeout(context());
+            next_date < m_idle_timer.next_date())
+            m_idle_timer.set_next_date(next_date);
+        m_refresh_completion_pending = true;
+    }
+
     void paste(StringView content) override
     {
         m_line_editor.insert(content);
@@ -1122,14 +1137,12 @@ private:
             if (menu)
                 context().client().menu_select(0);
             auto prefix = line.substr(m_completions.start, m_completions.end - m_completions.start);
-            if (not menu and not contains(m_completions.candidates, prefix))
+            m_prefix_in_completions = not menu and not contains(m_completions.candidates, prefix);
+            if (m_prefix_in_completions)
             {
                 m_current_completion = m_completions.candidates.size();
                 m_completions.candidates.push_back(prefix.str());
-                m_prefix_in_completions = true;
             }
-            else
-                m_prefix_in_completions = false;
         } catch (runtime_error&) {}
     }
 
@@ -1814,6 +1827,11 @@ void InputHandler::handle_key(Key key)
     }
 }
 
+void InputHandler::refresh_ifn()
+{
+    current_mode().refresh_ifn();
+}
+
 void InputHandler::start_recording(char reg)
 {
     kak_assert(m_recording_reg == 0);
diff --git a/src/input_handler.hh b/src/input_handler.hh
index afcc24d49..a26d392c5 100644
--- a/src/input_handler.hh
+++ b/src/input_handler.hh
@@ -99,6 +99,8 @@ public:
     // process the given key
     void handle_key(Key key);
 
+    void refresh_ifn();
+
     void start_recording(char reg);
     bool is_recording() const;
     void stop_recording();
diff --git a/src/main.cc b/src/main.cc
index bc12f335e..09df53612 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -46,6 +46,7 @@ struct {
     StringView notes;
 } constexpr version_notes[] = { {
         0,
+        "» asynchronous {+u}shell-script-candidates{} completion\n"
         "» {+b}%val{window_range}{} is now emitted as separate strings\n"
         "» {+b}+{} only duplicates identical selections a single time\n"
         "» {+u}daemonize-session{} command\n"
diff --git a/src/shell_manager.cc b/src/shell_manager.cc
index 0d012538e..e505fdef1 100644
--- a/src/shell_manager.cc
+++ b/src/shell_manager.cc
@@ -86,26 +86,6 @@ ShellManager::ShellManager(ConstArrayView<EnvVarDesc> builtin_env_vars)
 namespace
 {
 
-struct UniqueFd
-{
-    UniqueFd(int fd = -1) : fd{fd} {}
-    UniqueFd(UniqueFd&& other) : fd{other.fd} { other.fd = -1; }
-    UniqueFd& operator=(UniqueFd&& other) { std::swap(fd, other.fd); other.close(); return *this; }
-    ~UniqueFd() { close(); }
-
-    explicit operator bool() const { return fd != -1; }
-    void close() { if (fd != -1) { ::close(fd); fd = -1; } }
-    int fd;
-};
-
-struct Shell
-{
-    pid_t pid;
-    UniqueFd in;
-    UniqueFd out;
-    UniqueFd err;
-};
-
 Shell spawn_shell(const char* shell, StringView cmdline,
                   ConstArrayView<String> params,
                   ConstArrayView<String> kak_env,
@@ -145,13 +125,13 @@ Shell spawn_shell(const char* shell, StringView cmdline,
         close(oldfd);
     };
 
-    renamefd(stdin_pipe[0].fd, 0);
-    renamefd(stdout_pipe[1].fd, 1);
-    renamefd(stderr_pipe[1].fd, 2);
+    renamefd((int)stdin_pipe[0], 0);
+    renamefd((int)stdout_pipe[1], 1);
+    renamefd((int)stderr_pipe[1], 2);
 
-    close(stdin_pipe[1].fd);
-    close(stdout_pipe[0].fd);
-    close(stderr_pipe[0].fd);
+    close((int)stdin_pipe[1]);
+    close((int)stdout_pipe[0]);
+    close((int)stderr_pipe[0]);
 
     execve(shell, (char* const*)execparams.data(), (char* const*)envptrs.data());
     char buffer[1024];
@@ -215,13 +195,13 @@ FDWatcher make_reader(int fd, String& contents, OnClose&& on_close)
 
 FDWatcher make_pipe_writer(UniqueFd& fd, StringView contents)
 {
-    int flags = fcntl(fd.fd, F_GETFL, 0);
-    fcntl(fd.fd, F_SETFL, flags | O_NONBLOCK);
-    return {fd.fd, FdEvents::Write, EventMode::Urgent,
+    int flags = fcntl((int)fd, F_GETFL, 0);
+    fcntl((int)fd, F_SETFL, flags | O_NONBLOCK);
+    return {(int)fd, FdEvents::Write, EventMode::Urgent,
             [contents, &fd](FDWatcher& watcher, FdEvents, EventMode) mutable {
-        while (fd_writable(fd.fd))
+        while (fd_writable((int)fd))
         {
-            ssize_t size = ::write(fd.fd, contents.begin(),
+            ssize_t size = ::write((int)fd, contents.begin(),
                                    (size_t)contents.length());
             if (size > 0)
                 contents = contents.substr(ByteCount{(int)size});
@@ -310,8 +290,8 @@ std::pair<String, int> ShellManager::eval(
     auto wait_time = Clock::now();
 
     String stdout_contents, stderr_contents;
-    auto stdout_reader = make_reader(shell.out.fd, stdout_contents, [&](bool){ shell.out.close(); });
-    auto stderr_reader = make_reader(shell.err.fd, stderr_contents, [&](bool){ shell.err.close(); });
+    auto stdout_reader = make_reader((int)shell.out, stdout_contents, [&](bool){ shell.out.close(); });
+    auto stderr_reader = make_reader((int)shell.err, stderr_contents, [&](bool){ shell.err.close(); });
     auto stdin_writer = make_pipe_writer(shell.in, input);
 
     // block SIGCHLD to make sure we wont receive it before
@@ -324,7 +304,7 @@ std::pair<String, int> ShellManager::eval(
 
     int status = 0;
     // check for termination now that SIGCHLD is blocked
-    bool terminated = waitpid(shell.pid, &status, WNOHANG) != 0;
+    bool terminated = waitpid((int)shell.pid, &status, WNOHANG) != 0;
     bool failed = false;
 
     using namespace std::chrono;
@@ -357,7 +337,7 @@ std::pair<String, int> ShellManager::eval(
         }
         catch (cancel&)
         {
-            kill(shell.pid, SIGINT);
+            kill((int)shell.pid, SIGINT);
             cancelling = true;
         }
         catch (runtime_error& error)
@@ -366,7 +346,7 @@ std::pair<String, int> ShellManager::eval(
             failed = true;
         }
         if (not terminated)
-            terminated = waitpid(shell.pid, &status, WNOHANG) == shell.pid;
+            terminated = waitpid((int)shell.pid, &status, WNOHANG) == (int)shell.pid;
     }
 
     if (not stderr_contents.empty())
@@ -394,6 +374,18 @@ std::pair<String, int> ShellManager::eval(
     return { std::move(stdout_contents), WIFEXITED(status) ? WEXITSTATUS(status) : -1 };
 }
 
+Shell ShellManager::spawn(StringView cmdline, const Context& context,
+                          bool open_stdin, const ShellContext& shell_context)
+{
+    auto kak_env = generate_env(cmdline, context, [&](StringView name, Quoting quoting) {
+        if (auto it = shell_context.env_vars.find(name); it != shell_context.env_vars.end())
+            return it->value;
+        return join(get_val(name, context) | transform(quoter(quoting)), ' ', false);
+    });
+
+    return spawn_shell(m_shell.c_str(), cmdline, shell_context.params, kak_env, open_stdin);
+}
+
 Vector<String> ShellManager::get_val(StringView name, const Context& context) const
 {
     auto env_var = find_if(m_env_vars, [name](const EnvVarDesc& desc) {
diff --git a/src/shell_manager.hh b/src/shell_manager.hh
index 7dc00f839..6b45961a8 100644
--- a/src/shell_manager.hh
+++ b/src/shell_manager.hh
@@ -5,8 +5,13 @@
 #include "env_vars.hh"
 #include "string.hh"
 #include "utils.hh"
+#include "unique_descriptor.hh"
 #include "completion.hh"
 
+#include <signal.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
 namespace Kakoune
 {
 
@@ -27,6 +32,19 @@ struct EnvVarDesc
     Retriever func;
 };
 
+inline void closepid(int pid){ kill(pid, SIGTERM); int status = 0; waitpid(pid, &status, 0); }
+
+using UniqueFd = UniqueDescriptor<::close>;
+using UniquePid = UniqueDescriptor<closepid>;
+
+struct Shell
+{
+    UniquePid pid;
+    UniqueFd in;
+    UniqueFd out;
+    UniqueFd err;
+};
+
 class ShellManager : public Singleton<ShellManager>
 {
 public:
@@ -44,6 +62,11 @@ public:
                                 Flags flags = Flags::WaitForStdout,
                                 const ShellContext& shell_context = {});
 
+    Shell spawn(StringView cmdline,
+                const Context& context,
+                bool open_stdin,
+                const ShellContext& shell_complete = {});
+
     Vector<String> get_val(StringView name, const Context& context) const;
 
     CandidateList complete_env_var(StringView prefix, ByteCount cursor_pos) const;
diff --git a/src/unique_descriptor.hh b/src/unique_descriptor.hh
new file mode 100644
index 000000000..c099c17e7
--- /dev/null
+++ b/src/unique_descriptor.hh
@@ -0,0 +1,23 @@
+#ifndef fd_hh_INCLUDED
+#define fd_hh_INCLUDED
+
+namespace Kakoune
+{
+
+template<auto close_fn>
+struct UniqueDescriptor
+{
+    UniqueDescriptor(int descriptor = -1) : descriptor{descriptor} {}
+    UniqueDescriptor(UniqueDescriptor&& other) : descriptor{other.descriptor} { other.descriptor = -1; }
+    UniqueDescriptor& operator=(UniqueDescriptor&& other) { std::swap(descriptor, other.descriptor); other.close(); return *this; }
+    ~UniqueDescriptor() { close(); }
+
+    explicit operator int() const { return descriptor; }
+    explicit operator bool() const { return descriptor != -1; }
+    void close() { if (descriptor != -1) { close_fn(descriptor); descriptor = -1; } }
+    int descriptor;
+};
+
+}
+
+#endif // fd_hh_INCLUDED
diff --git a/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/cmd b/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/cmd
index 3a6c3b743..8b1378917 100644
--- a/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/cmd
+++ b/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/cmd
@@ -1 +1 @@
-:my-command <s-tab><s-tab><ret>
+
diff --git a/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/rc b/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/rc
index 53ad0d140..16ae8a38b 100644
--- a/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/rc
+++ b/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/rc
@@ -1 +1,2 @@
+set global autocomplete prompt
 def my-command -params 0..1 -shell-script-candidates %{ printf "aaa\nbbb\nccc" } %{ exec i %arg{1} <esc> }
diff --git a/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/script b/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/script
new file mode 100644
index 000000000..f73d1c08d
--- /dev/null
+++ b/test/regression/0-nothing-selected-on-prompt-initial-shift-tab/script
@@ -0,0 +1,6 @@
+ui_out -ignore 7
+ui_in '{ "jsonrpc": "2.0", "method": "keys", "params": [ ":my-command " ] }'
+ui_out -ignore 7
+ui_out '{ "jsonrpc": "2.0", "method": "refresh", "params": [false] }'
+ui_out '{ "jsonrpc": "2.0", "method": "menu_show", "params": [[[{ "face": { "fg": "default", "bg": "default", "underline": "default", "attributes": [] }, "contents": "aaa" }], [{ "face": { "fg": "default", "bg": "default", "underline": "default", "attributes": [] }, "contents": "bbb" }], [{ "face": { "fg": "default", "bg": "default", "underline": "default", "attributes": [] }, "contents": "ccc" }]], { "line": 0, "column": 0 }, { "fg": "white", "bg": "blue", "underline": "default", "attributes": [] }, { "fg": "blue", "bg": "white", "underline": "default", "attributes": [] }, "prompt"] }'
+ui_in '{ "jsonrpc": "2.0", "method": "keys", "params": [ "<s-tab><ret>" ] }'
diff --git a/test/shell/prompt-shell-script-candidates/cmd b/test/shell/prompt-shell-script-candidates/cmd
index 7cc383905..8b1378917 100644
--- a/test/shell/prompt-shell-script-candidates/cmd
+++ b/test/shell/prompt-shell-script-candidates/cmd
@@ -1 +1 @@
-:foo<ret>b<tab><tab><ret>
+
diff --git a/test/shell/prompt-shell-script-candidates/rc b/test/shell/prompt-shell-script-candidates/rc
index 6aaaf259c..82e4e62ec 100644
--- a/test/shell/prompt-shell-script-candidates/rc
+++ b/test/shell/prompt-shell-script-candidates/rc
@@ -1,3 +1,4 @@
+set-option global autocomplete prompt
 define-command foo %{
     prompt -shell-script-candidates %{ printf 'foo\nbar\nhaz\n' } ': ' %{exec i %val{text} <esc>}
 }
diff --git a/test/shell/prompt-shell-script-candidates/script b/test/shell/prompt-shell-script-candidates/script
new file mode 100644
index 000000000..231c87255
--- /dev/null
+++ b/test/shell/prompt-shell-script-candidates/script
@@ -0,0 +1,7 @@
+ui_out -ignore 7
+ui_in '{ "jsonrpc": "2.0", "method": "keys", "params": [ ":foo<ret>b" ] }'
+ui_out -ignore 4
+ui_out '{ "jsonrpc": "2.0", "method": "refresh", "params": [false] }'
+ui_out -ignore 3
+ui_out '{ "jsonrpc": "2.0", "method": "menu_show", "params": [[[{ "face": { "fg": "default", "bg": "default", "underline": "default", "attributes": [] }, "contents": "bar" }]], { "line": 0, "column": 0 }, { "fg": "white", "bg": "blue", "underline": "default", "attributes": [] }, { "fg": "blue", "bg": "white", "underline": "default", "attributes": [] }, "prompt"] }'
+ui_in '{ "jsonrpc": "2.0", "method": "keys", "params": [ "<tab><ret>" ] }'