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, 24 insertions, 14 deletions

diff --git a/src/libcmd/command.hh b/src/libcmd/command.hh
index cab379b84..ffa8e784f 100644
--- a/src/libcmd/command.hh
+++ b/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
diff --git a/src/libcmd/installables.cc b/src/libcmd/installables.cc
index 1bcef4172..81c8dd062 100644
--- a/src/libcmd/installables.cc
+++ b/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({