nix: Implement basic bash completion

9 files changed, 95 insertions, 4 deletions

diff --git a/Makefile b/Makefile
index e3057c36c..545397b23 100644
--- a/Makefile
+++ b/Makefile
@@ -11,6 +11,7 @@ makefiles = \
   src/resolve-system-dependencies/local.mk \
   scripts/local.mk \
   corepkgs/local.mk \
+  misc/bash/local.mk \
   misc/systemd/local.mk \
   misc/launchd/local.mk \
   misc/upstart/local.mk \
diff --git a/misc/bash/completion.sh b/misc/bash/completion.sh
new file mode 100644
index 000000000..097353b50
--- /dev/null
+++ b/misc/bash/completion.sh
@@ -0,0 +1,7 @@
+function _complete_nix {
+    while IFS= read -r line; do
+        COMPREPLY+=("$line")
+    done < <(NIX_GET_COMPLETIONS=$COMP_CWORD "${COMP_WORDS[@]}")
+}
+
+complete -F _complete_nix nix
diff --git a/misc/bash/local.mk b/misc/bash/local.mk
new file mode 100644
index 000000000..99ada5108
--- /dev/null
+++ b/misc/bash/local.mk
@@ -0,0 +1 @@
+$(eval $(call install-file-as, $(d)/completion.sh, $(datarootdir)/bash-completion/completions/_nix3, 0644))
diff --git a/src/libmain/common-args.cc b/src/libmain/common-args.cc
index 51e199ea5..f3f508ff4 100644
--- a/src/libmain/common-args.cc
+++ b/src/libmain/common-args.cc
@@ -31,9 +31,17 @@ MixCommonArgs::MixCommonArgs(const string & programName)
         .labels = {"name", "value"},
         .handler = {[](std::string name, std::string value) {
             try {
+                if (auto prefix = needsCompletion(name)) {
+                    std::map<std::string, Config::SettingInfo> settings;
+                    globalConfig.getSettings(settings);
+                    for (auto & s : settings)
+                        if (hasPrefix(s.first, *prefix))
+                            completions->insert(s.first);
+                }
                 globalConfig.set(name, value);
             } catch (UsageError & e) {
-                warn(e.what());
+                if (!completions)
+                    warn(e.what());
             }
         }},
     });
diff --git a/src/libutil/args.cc b/src/libutil/args.cc
index f829415d1..ceaabcfca 100644
--- a/src/libutil/args.cc
+++ b/src/libutil/args.cc
@@ -13,6 +13,19 @@ void Args::addFlag(Flag && flag_)
     if (flag->shortName) shortFlags[flag->shortName] = flag;
 }
 
+std::shared_ptr<std::set<std::string>> completions;
+
+std::string completionMarker = "___COMPLETE___";
+
+std::optional<std::string> needsCompletion(std::string_view s)
+{
+    if (!completions) return {};
+    auto i = s.find(completionMarker);
+    if (i != std::string::npos)
+        return std::string(s.begin(), i);
+    return {};
+}
+
 void Args::parseCmdline(const Strings & _cmdline)
 {
     Strings pendingArgs;
@@ -20,6 +33,13 @@ void Args::parseCmdline(const Strings & _cmdline)
 
     Strings cmdline(_cmdline);
 
+    if (auto s = getEnv("NIX_GET_COMPLETIONS")) {
+        size_t n = std::stoi(*s);
+        assert(n > 0 && n <= cmdline.size());
+        *std::next(cmdline.begin(), n - 1) += completionMarker;
+        completions = std::make_shared<decltype(completions)::element_type>();
+    }
+
     for (auto pos = cmdline.begin(); pos != cmdline.end(); ) {
 
         auto arg = *pos;
@@ -99,18 +119,32 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
     auto process = [&](const std::string & name, const Flag & flag) -> bool {
         ++pos;
         std::vector<std::string> args;
+        bool anyNeedsCompletion = false;
         for (size_t n = 0 ; n < flag.handler.arity; ++n) {
             if (pos == end) {
                 if (flag.handler.arity == ArityAny) break;
-                throw UsageError("flag '%s' requires %d argument(s)", name, flag.handler.arity);
+                if (anyNeedsCompletion)
+                    args.push_back("");
+                else
+                    throw UsageError("flag '%s' requires %d argument(s)", name, flag.handler.arity);
+            } else {
+                if (needsCompletion(*pos))
+                    anyNeedsCompletion = true;
+                args.push_back(*pos++);
             }
-            args.push_back(*pos++);
         }
         flag.handler.fun(std::move(args));
         return true;
     };
 
     if (string(*pos, 0, 2) == "--") {
+        if (auto prefix = needsCompletion(*pos)) {
+            for (auto & [name, flag] : longFlags) {
+                if (!hiddenCategories.count(flag->category)
+                    && hasPrefix(name, std::string(*prefix, 2)))
+                    completions->insert("--" + name);
+            }
+        }
         auto i = longFlags.find(string(*pos, 2));
         if (i == longFlags.end()) return false;
         return process("--" + i->first, *i->second);
@@ -123,6 +157,14 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
         return process(std::string("-") + c, *i->second);
     }
 
+    if (auto prefix = needsCompletion(*pos)) {
+        if (prefix == "-") {
+            completions->insert("--");
+            for (auto & [flag, _] : shortFlags)
+                completions->insert(std::string("-") + flag);
+        }
+    }
+
     return false;
 }
 
@@ -161,6 +203,10 @@ Args::Flag Args::Flag::mkHashTypeFlag(std::string && longName, HashType * ht)
         .description = "hash algorithm ('md5', 'sha1', 'sha256', or 'sha512')",
         .labels = {"hash-algo"},
         .handler = {[ht](std::string s) {
+            if (auto prefix = needsCompletion(s))
+                for (auto & type : hashTypes)
+                    if (hasPrefix(type, *prefix))
+                        completions->insert(type);
             *ht = parseHashType(s);
             if (*ht == htUnknown)
                 throw UsageError("unknown hash type '%1%'", s);
@@ -217,6 +263,11 @@ MultiCommand::MultiCommand(const Commands & commands)
 {
     expectedArgs.push_back(ExpectedArg{"command", 1, true, [=](std::vector<std::string> ss) {
         assert(!command);
+        if (auto prefix = needsCompletion(ss[0])) {
+            for (auto & [name, command] : commands)
+                if (hasPrefix(name, *prefix))
+                    completions->insert(name);
+        }
         auto i = commands.find(ss[0]);
         if (i == commands.end())
             throw UsageError("'%s' is not a recognised command", ss[0]);
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index 1932e6a8a..ae2875e72 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -256,4 +256,8 @@ typedef std::vector<std::pair<std::string, std::string>> Table2;
 
 void printTable(std::ostream & out, const Table2 & table);
 
+extern std::shared_ptr<std::set<std::string>> completions;
+
+std::optional<std::string> needsCompletion(std::string_view s);
+
 }
diff --git a/src/libutil/hash.cc b/src/libutil/hash.cc
index 7caee1da7..606c78ed7 100644
--- a/src/libutil/hash.cc
+++ b/src/libutil/hash.cc
@@ -16,6 +16,9 @@
 namespace nix {
 
 
+std::set<std::string> hashTypes = { "md5", "sha1", "sha256", "sha512" };
+
+
 void Hash::init()
 {
     if (type == htMD5) hashSize = md5HashSize;
diff --git a/src/libutil/hash.hh b/src/libutil/hash.hh
index ea9fca3e7..e1a16ba22 100644
--- a/src/libutil/hash.hh
+++ b/src/libutil/hash.hh
@@ -18,6 +18,8 @@ const int sha1HashSize = 20;
 const int sha256HashSize = 32;
 const int sha512HashSize = 64;
 
+extern std::set<std::string> hashTypes;
+
 extern const string base32Chars;
 
 enum Base : int { Base64, Base32, Base16, SRI };
diff --git a/src/nix/main.cc b/src/nix/main.cc
index 3915a4896..c491bc264 100644
--- a/src/nix/main.cc
+++ b/src/nix/main.cc
@@ -166,7 +166,21 @@ void mainWrapped(int argc, char * * argv)
 
     NixArgs args;
 
-    args.parseCmdline(argvToStrings(argc, argv));
+    Finally printCompletions([&]()
+    {
+        if (completions) {
+            for (auto & s : *completions)
+                std::cout << s << "\n";
+        }
+    });
+
+    try {
+        args.parseCmdline(argvToStrings(argc, argv));
+    } catch (UsageError &) {
+        if (!completions) throw;
+    }
+
+    if (completions) return;
 
     settings.requireExperimentalFeature("nix-command");