Fix flake input completion for `InstallablesCommand`s

Defers completion of flake inputs until the whole command line is parsed
so that we know what flakes we need to complete the inputs of.

Previously, `nix build flake --update-input <Tab>` always behaved like
`nix build . --update-input <Tab>`.

2 files changed, 16 insertions, 1 deletions

diff --git a/src/libutil/args.cc b/src/libutil/args.cc
index 4b8c55686..44b63f0f6 100644
--- a/src/libutil/args.cc
+++ b/src/libutil/args.cc
@@ -124,7 +124,7 @@ bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
         bool anyCompleted = false;
         for (size_t n = 0 ; n < flag.handler.arity; ++n) {
             if (pos == end) {
-                if (flag.handler.arity == ArityAny) break;
+                if (flag.handler.arity == ArityAny || anyCompleted) break;
                 throw UsageError("flag '%s' requires %d argument(s)", name, flag.handler.arity);
             }
             if (auto prefix = needsCompletion(*pos)) {
@@ -362,6 +362,14 @@ bool MultiCommand::processArgs(const Strings & args, bool finish)
         return Args::processArgs(args, finish);
 }
 
+void MultiCommand::completionHook()
+{
+    if (command)
+        return command->second->completionHook();
+    else
+        return Args::completionHook();
+}
+
 nlohmann::json MultiCommand::toJSON()
 {
     auto cmds = nlohmann::json::object();
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index 07c017719..84866f12b 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -148,6 +148,11 @@ protected:
        argument (if any) have been processed. */
     virtual void initialFlagsProcessed() {}
 
+    /* Called after the command line has been processed if we need to generate
+       completions. Useful for commands that need to know the whole command line
+       in order to know what completions to generate. */
+    virtual void completionHook() { }
+
 public:
 
     void addFlag(Flag && flag);
@@ -223,6 +228,8 @@ public:
 
     bool processArgs(const Strings & args, bool finish) override;
 
+    void completionHook() override;
+
     nlohmann::json toJSON() override;
 };