Fix flake input completion for InstallablesCommands

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>`.

src/libcmd/command.hh

@@ -77,12 +77,16 @@ struct MixFlakeOptions : virtual Args, EvalCommand
 {
     flake::LockFlags lockFlags;
 
+    std::optional<std::string> needsFlakeInputCompletion = {};
+
     MixFlakeOptions();
 
     virtual std::vector<std::string> getFlakesForCompletion()
     { return {}; }
 
     void completeFlakeInput(std::string_view prefix);
+
+    void completionHook() override;
 };
 
 struct SourceExprCommand : virtual Args, MixFlakeOptions

src/libcmd/installables.cc

@@ -23,18 +23,6 @@
 
 namespace nix {
 
-void MixFlakeOptions::completeFlakeInput(std::string_view prefix)
-{
-    auto evalState = getEvalState();
-    for (auto & flakeRefS : getFlakesForCompletion()) {
-        auto flakeRef = parseFlakeRefWithFragment(expandTilde(flakeRefS), absPath(".")).first;
-        auto flake = flake::getFlake(*evalState, flakeRef, true);
-        for (auto & input : flake.inputs)
-            if (hasPrefix(input.first, prefix))
-                completions->add(input.first);
-    }
-}
-
 MixFlakeOptions::MixFlakeOptions()
 {
     auto category = "Common flake-related options";
@@ -87,7 +75,7 @@ MixFlakeOptions::MixFlakeOptions()
             lockFlags.inputUpdates.insert(flake::parseInputPath(s));
         }},
         .completer = {[&](size_t, std::string_view prefix) {
-            completeFlakeInput(prefix);
+            needsFlakeInputCompletion = {std::string(prefix)};
         }}
     });
 
@@ -104,7 +92,7 @@ MixFlakeOptions::MixFlakeOptions()
         }},
         .completer = {[&](size_t n, std::string_view prefix) {
             if (n == 0)
-                completeFlakeInput(prefix);
+                needsFlakeInputCompletion = {std::string(prefix)};
             else if (n == 1)
                 completeFlakeRef(getEvalState()->store, prefix);
         }}
@@ -137,6 +125,24 @@ MixFlakeOptions::MixFlakeOptions()
     });
 }
 
+void MixFlakeOptions::completeFlakeInput(std::string_view prefix)
+{
+    auto evalState = getEvalState();
+    for (auto & flakeRefS : getFlakesForCompletion()) {
+        auto flakeRef = parseFlakeRefWithFragment(expandTilde(flakeRefS), absPath(".")).first;
+        auto flake = flake::getFlake(*evalState, flakeRef, true);
+        for (auto & input : flake.inputs)
+            if (hasPrefix(input.first, prefix))
+                completions->add(input.first);
+    }
+}
+
+void MixFlakeOptions::completionHook()
+{
+    if (auto & prefix = needsFlakeInputCompletion)
+        completeFlakeInput(*prefix);
+}
+
 SourceExprCommand::SourceExprCommand(bool supportReadOnlyMode)
 {
     addFlag({

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();

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;
 };
 

src/nix/main.cc

@@ -342,7 +342,10 @@ void mainWrapped(int argc, char * * argv)
         if (!completions) throw;
     }
 
-    if (completions) return;
+    if (completions) {
+        args.completionHook();
+        return;
+    }
 
     if (args.showVersion) {
         printVersion(programName);