LibCore: Make ArgsParser perform some completion if passed --complete

This makes it possible to autocomplete flags and options via the Shell.

Userland/Libraries/LibCore/ArgsParser.cpp

@@ -5,6 +5,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include "AK/JsonObject.h"
 #include <AK/Format.h>
 #include <AK/StringBuilder.h>
 #include <LibCore/ArgsParser.h>
@@ -30,6 +31,7 @@ ArgsParser::ArgsParser()
 {
     add_option(m_show_help, "Display help message and exit", "help", 0);
     add_option(m_show_version, "Print version", "version", 0);
+    add_option(m_perform_autocomplete, "Perform autocompletion", "complete", 0);
 }
 
 bool ArgsParser::parse(int argc, char* const* argv, FailureBehavior failure_behavior)
@@ -110,7 +112,7 @@ bool ArgsParser::parse(int argc, char* const* argv, FailureBehavior failure_beha
     }
 
     // We're done processing options.
-    // Now let's show version or help if requested.
+    // Now let's show version or help if requested, or perform autocompletion if needed.
 
     if (m_show_version) {
         print_version(stdout);
@@ -118,6 +120,7 @@ bool ArgsParser::parse(int argc, char* const* argv, FailureBehavior failure_beha
             exit(0);
         return false;
     }
+
     if (m_show_help) {
         print_usage(stdout, argv[0]);
         if (failure_behavior == FailureBehavior::Exit || failure_behavior == FailureBehavior::PrintUsageAndExit)
@@ -125,6 +128,13 @@ bool ArgsParser::parse(int argc, char* const* argv, FailureBehavior failure_beha
         return false;
     }
 
+    if (m_perform_autocomplete) {
+        autocomplete(stdout, argv[0], Span<char const* const> { argv + optind, static_cast<size_t>(argc - optind) });
+        if (failure_behavior == FailureBehavior::Exit || failure_behavior == FailureBehavior::PrintUsageAndExit)
+            exit(0);
+        return false;
+    }
+
     // Now let's parse positional arguments.
 
     int values_left = argc - optind;
@@ -628,4 +638,110 @@ void ArgsParser::add_positional_argument(Vector<StringView>& values, char const*
     add_positional_argument(move(arg));
 }
 
+void ArgsParser::autocomplete(FILE* file, StringView program_name, Span<char const* const> remaining_arguments)
+{
+    // We expect the full invocation of the program to be available as positional args,
+    // e.g. `foo --bar arg -b` (program invoked as `foo --complete -- foo --bar arg -b`)
+    auto first = true;
+    auto seen_all_options = false;
+    auto skip_next = false;
+
+    StringView argument_to_complete;
+    StringView option_to_complete;
+    auto completing_option = false;
+
+    for (auto& arg : remaining_arguments) {
+        StringView argument { arg };
+
+        completing_option = false;
+        if (skip_next) {
+            argument_to_complete = argument;
+            skip_next = false;
+            continue;
+        }
+
+        // Skip over the program name.
+        if (first && program_name == argument) {
+            first = false;
+            continue;
+        }
+
+        if (seen_all_options) {
+            argument_to_complete = argument;
+            continue;
+        }
+
+        if (argument.starts_with("--")) {
+            option_to_complete = argument;
+            completing_option = true;
+
+            if (argument == "--") {
+                seen_all_options = true;
+                continue;
+            }
+
+            // Look for a long option
+            auto option_pattern = argument.substring_view(2);
+            auto it = m_options.find_if([&](auto& option) { return StringView(option.long_name) == option_pattern; });
+            if (it.is_end())
+                continue;
+
+            if (it->requires_argument)
+                skip_next = true;
+            continue;
+        }
+
+        if (argument.starts_with("-")) {
+            option_to_complete = argument;
+            completing_option = true;
+
+            if (argument == "-") {
+                option_to_complete = argument;
+                continue;
+            }
+
+            // Look for a short option
+            auto option_pattern = argument[argument.length() - 1];
+            auto it = m_options.find_if([&](auto& option) { return option.short_name == option_pattern; });
+            if (it.is_end())
+                continue;
+
+            if (it->requires_argument)
+                skip_next = true;
+            continue;
+        }
+    }
+
+    // We don't know how to complete arguments quite yet.
+    if (!completing_option)
+        return;
+
+    auto write_completion = [&](auto format, auto... args) {
+        JsonObject object;
+        object.set("completion", String::formatted(format, args...));
+        object.set("static_offset", 0);
+        object.set("invariant_offset", option_to_complete.length());
+        outln(file, "{}", object.to_string());
+    };
+
+    if (option_to_complete.starts_with("--")) {
+        // Complete a long option.
+        auto option_pattern = option_to_complete.substring_view(2);
+        for (auto& option : m_options) {
+            StringView option_string = option.long_name;
+            if (option_string.starts_with(option_pattern)) {
+                write_completion("--{}", option_string);
+            }
+        }
+    } else {
+        // Complete a short option, note that we're not going to attempt to 'match' anything here.
+        for (auto& option : m_options) {
+            if (option.short_name == 0)
+                continue;
+
+            write_completion("-{}", option.short_name);
+        }
+    }
+}
+
 }

Userland/Libraries/LibCore/ArgsParser.h

@@ -32,11 +32,11 @@ public:
 
     struct Option {
         bool requires_argument { true };
-        const char* help_string { nullptr };
-        const char* long_name { nullptr };
+        char const* help_string { nullptr };
+        char const* long_name { nullptr };
         char short_name { 0 };
-        const char* value_name { nullptr };
-        Function<bool(const char*)> accept_value;
+        char const* value_name { nullptr };
+        Function<bool(char const*)> accept_value;
 
         String name_for_display() const
         {
@@ -47,11 +47,11 @@ public:
     };
 
     struct Arg {
-        const char* help_string { nullptr };
-        const char* name { nullptr };
+        char const* help_string { nullptr };
+        char const* name { nullptr };
         int min_values { 0 };
         int max_values { 1 };
-        Function<bool(const char*)> accept_value;
+        Function<bool(char const*)> accept_value;
     };
 
     bool parse(int argc, char* const* argv, FailureBehavior failure_behavior = FailureBehavior::PrintUsageAndExit);
@@ -61,42 +61,45 @@ public:
     }
 
     // *Without* trailing newline!
-    void set_general_help(const char* help_string) { m_general_help = help_string; };
+    void set_general_help(char const* help_string) { m_general_help = help_string; };
     void set_stop_on_first_non_option(bool stop_on_first_non_option) { m_stop_on_first_non_option = stop_on_first_non_option; }
-    void print_usage(FILE*, const char* argv0);
-    void print_usage_terminal(FILE*, const char* argv0);
-    void print_usage_markdown(FILE*, const char* argv0);
+    void print_usage(FILE*, char const* argv0);
+    void print_usage_terminal(FILE*, char const* argv0);
+    void print_usage_markdown(FILE*, char const* argv0);
     void print_version(FILE*);
 
     void add_option(Option&&);
-    void add_ignored(const char* long_name, char short_name);
-    void add_option(bool& value, const char* help_string, const char* long_name, char short_name);
-    void add_option(const char*& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
-    void add_option(String& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
+    void add_ignored(char const* long_name, char short_name);
+    void add_option(bool& value, char const* help_string, char const* long_name, char short_name);
+    void add_option(char const*& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
+    void add_option(String& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
     void add_option(StringView& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
-    void add_option(int& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
-    void add_option(unsigned& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
-    void add_option(double& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
-    void add_option(Optional<double>& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
-    void add_option(Optional<size_t>& value, const char* help_string, const char* long_name, char short_name, const char* value_name);
+    void add_option(int& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
+    void add_option(unsigned& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
+    void add_option(double& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
+    void add_option(Optional<double>& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
+    void add_option(Optional<size_t>& value, char const* help_string, char const* long_name, char short_name, char const* value_name);
 
     void add_positional_argument(Arg&&);
-    void add_positional_argument(const char*& value, const char* help_string, const char* name, Required required = Required::Yes);
-    void add_positional_argument(String& value, const char* help_string, const char* name, Required required = Required::Yes);
+    void add_positional_argument(char const*& value, char const* help_string, char const* name, Required required = Required::Yes);
+    void add_positional_argument(String& value, char const* help_string, char const* name, Required required = Required::Yes);
     void add_positional_argument(StringView& value, char const* help_string, char const* name, Required required = Required::Yes);
-    void add_positional_argument(int& value, const char* help_string, const char* name, Required required = Required::Yes);
-    void add_positional_argument(unsigned& value, const char* help_string, const char* name, Required required = Required::Yes);
-    void add_positional_argument(double& value, const char* help_string, const char* name, Required required = Required::Yes);
-    void add_positional_argument(Vector<const char*>& value, const char* help_string, const char* name, Required required = Required::Yes);
+    void add_positional_argument(int& value, char const* help_string, char const* name, Required required = Required::Yes);
+    void add_positional_argument(unsigned& value, char const* help_string, char const* name, Required required = Required::Yes);
+    void add_positional_argument(double& value, char const* help_string, char const* name, Required required = Required::Yes);
+    void add_positional_argument(Vector<char const*>& value, char const* help_string, char const* name, Required required = Required::Yes);
     void add_positional_argument(Vector<StringView>& value, char const* help_string, char const* name, Required required = Required::Yes);
 
 private:
+    void autocomplete(FILE*, StringView program_name, Span<char const* const> remaining_arguments);
+
     Vector<Option> m_options;
     Vector<Arg> m_positional_args;
 
     bool m_show_help { false };
     bool m_show_version { false };
-    const char* m_general_help { nullptr };
+    bool m_perform_autocomplete { false };
+    char const* m_general_help { nullptr };
     bool m_stop_on_first_non_option { false };
 };