From d6d0e781bbade76f6ea3f310cb36973f4013826d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Na=C3=AFm=20Favier?= <n@monade.li>
Date: Sun, 19 Jun 2022 17:54:27 +0200
Subject: [PATCH] Complete flake inputs for all given flakes

Allow `nix build flake1 flake2 --update-input <Tab>` to complete the
inputs of both flakes.

Also do tilde expansion so that `nix build ~/flake --update-input <Tab>`
works.
---
 src/libcmd/command.hh      | 10 ++++++----
 src/libcmd/installables.cc | 34 ++++++++++++++++------------------
 src/nix/flake.cc           |  4 ++--
 3 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/src/libcmd/command.hh b/src/libcmd/command.hh
index 8982f21d0..cab379b84 100644
--- a/src/libcmd/command.hh
+++ b/src/libcmd/command.hh
@@ -79,8 +79,10 @@ struct MixFlakeOptions : virtual Args, EvalCommand
 
     MixFlakeOptions();
 
-    virtual std::optional<FlakeRef> getFlakeRefForCompletion()
+    virtual std::vector<std::string> getFlakesForCompletion()
     { return {}; }
+
+    void completeFlakeInput(std::string_view prefix);
 };
 
 struct SourceExprCommand : virtual Args, MixFlakeOptions
@@ -119,7 +121,7 @@ struct InstallablesCommand : virtual Args, SourceExprCommand
 
     virtual bool useDefaultInstallables() { return true; }
 
-    std::optional<FlakeRef> getFlakeRefForCompletion() override;
+    std::vector<std::string> getFlakesForCompletion() override;
 
 private:
 
@@ -135,9 +137,9 @@ struct InstallableCommand : virtual Args, SourceExprCommand
 
     void prepare() override;
 
-    std::optional<FlakeRef> getFlakeRefForCompletion() override
+    std::vector<std::string> getFlakesForCompletion() override
     {
-        return parseFlakeRefWithFragment(_installable, absPath(".")).first;
+        return {_installable};
     }
 
 private:
diff --git a/src/libcmd/installables.cc b/src/libcmd/installables.cc
index ffc25135e..1bcef4172 100644
--- a/src/libcmd/installables.cc
+++ b/src/libcmd/installables.cc
@@ -23,15 +23,16 @@
 
 namespace nix {
 
-void completeFlakeInputPath(
-    ref<EvalState> evalState,
-    const FlakeRef & flakeRef,
-    std::string_view prefix)
+void MixFlakeOptions::completeFlakeInput(std::string_view prefix)
 {
-    auto flake = flake::getFlake(*evalState, flakeRef, true);
-    for (auto & input : flake.inputs)
-        if (hasPrefix(input.first, prefix))
-            completions->add(input.first);
+    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()
@@ -86,8 +87,7 @@ MixFlakeOptions::MixFlakeOptions()
             lockFlags.inputUpdates.insert(flake::parseInputPath(s));
         }},
         .completer = {[&](size_t, std::string_view prefix) {
-            if (auto flakeRef = getFlakeRefForCompletion())
-                completeFlakeInputPath(getEvalState(), *flakeRef, prefix);
+            completeFlakeInput(prefix);
         }}
     });
 
@@ -103,12 +103,10 @@ MixFlakeOptions::MixFlakeOptions()
                 parseFlakeRef(flakeRef, absPath("."), true));
         }},
         .completer = {[&](size_t n, std::string_view prefix) {
-            if (n == 0) {
-                if (auto flakeRef = getFlakeRefForCompletion())
-                    completeFlakeInputPath(getEvalState(), *flakeRef, prefix);
-            } else if (n == 1) {
+            if (n == 0)
+                completeFlakeInput(prefix);
+            else if (n == 1)
                 completeFlakeRef(getEvalState()->store, prefix);
-            }
         }}
     });
 
@@ -1043,14 +1041,14 @@ void InstallablesCommand::prepare()
     installables = parseInstallables(getStore(), _installables);
 }
 
-std::optional<FlakeRef> InstallablesCommand::getFlakeRefForCompletion()
+std::vector<std::string> InstallablesCommand::getFlakesForCompletion()
 {
     if (_installables.empty()) {
         if (useDefaultInstallables())
-            return parseFlakeRefWithFragment(".", absPath(".")).first;
+            return {"."};
         return {};
     }
-    return parseFlakeRefWithFragment(_installables.front(), absPath(".")).first;
+    return _installables;
 }
 
 InstallableCommand::InstallableCommand(bool supportReadOnlyMode)
diff --git a/src/nix/flake.cc b/src/nix/flake.cc
index 8370b8dcf..439eb53ba 100644
--- a/src/nix/flake.cc
+++ b/src/nix/flake.cc
@@ -50,9 +50,9 @@ public:
         return flake::lockFlake(*getEvalState(), getFlakeRef(), lockFlags);
     }
 
-    std::optional<FlakeRef> getFlakeRefForCompletion() override
+    std::vector<std::string> getFlakesForCompletion() override
     {
-        return getFlakeRef();
+        return {flakeUrl};
     }
 };