From 6470450ab48196668099eeca2f171fec110d647b Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Mon, 8 Jun 2020 16:20:00 +0200
Subject: [PATCH] Add completion for --update-input

---
 src/libexpr/flake/lockfile.cc |  2 +-
 src/libexpr/flake/lockfile.hh |  1 -
 src/libutil/args.hh           |  2 --
 src/libutil/types.hh          |  1 +
 src/nix/command.hh            | 14 ++++++++++++--
 src/nix/flake.cc              |  7 ++++++-
 src/nix/installables.cc       | 25 +++++++++++++++++++++++++
 7 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/src/libexpr/flake/lockfile.cc b/src/libexpr/flake/lockfile.cc
index f26baa697..5c2f6f04c 100644
--- a/src/libexpr/flake/lockfile.cc
+++ b/src/libexpr/flake/lockfile.cc
@@ -198,7 +198,7 @@ InputPath parseInputPath(std::string_view s)
 
     for (auto & elem : tokenizeString<std::vector<std::string>>(s, "/")) {
         if (!std::regex_match(elem, flakeIdRegex))
-            throw Error("invalid flake input path element '%s'", elem);
+            throw UsageError("invalid flake input path element '%s'", elem);
         path.push_back(elem);
     }
 
diff --git a/src/libexpr/flake/lockfile.hh b/src/libexpr/flake/lockfile.hh
index ba47f9b89..6b05e8c24 100644
--- a/src/libexpr/flake/lockfile.hh
+++ b/src/libexpr/flake/lockfile.hh
@@ -72,4 +72,3 @@ InputPath parseInputPath(std::string_view s);
 std::string diffLockFiles(const LockFile & oldLocks, const LockFile & newLocks);
 
 }
-
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index 59541df99..2d49dbf66 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -8,8 +8,6 @@
 
 namespace nix {
 
-MakeError(UsageError, Error);
-
 enum HashType : char;
 
 class Args
diff --git a/src/libutil/types.hh b/src/libutil/types.hh
index 250c9581d..f70590b44 100644
--- a/src/libutil/types.hh
+++ b/src/libutil/types.hh
@@ -131,6 +131,7 @@ public:
     }
 
 MakeError(Error, BaseError);
+MakeError(UsageError, Error);
 
 class SysError : public Error
 {
diff --git a/src/nix/command.hh b/src/nix/command.hh
index fe8e40835..a8779b0e6 100644
--- a/src/nix/command.hh
+++ b/src/nix/command.hh
@@ -40,14 +40,17 @@ struct EvalCommand : virtual StoreCommand, MixEvalArgs
     std::shared_ptr<EvalState> evalState;
 };
 
-struct MixFlakeOptions : virtual Args
+struct MixFlakeOptions : virtual Args, EvalCommand
 {
     flake::LockFlags lockFlags;
 
     MixFlakeOptions();
+
+    virtual std::optional<FlakeRef> getFlakeRefForCompletion()
+    { return {}; }
 };
 
-struct SourceExprCommand : virtual Args, EvalCommand, MixFlakeOptions
+struct SourceExprCommand : virtual Args, MixFlakeOptions
 {
     std::optional<Path> file;
     std::optional<std::string> expr;
@@ -81,6 +84,8 @@ struct InstallablesCommand : virtual Args, SourceExprCommand
 
     virtual bool useDefaultInstallables() { return true; }
 
+    std::optional<FlakeRef> getFlakeRefForCompletion() override;
+
 private:
 
     std::vector<std::string> _installables;
@@ -95,6 +100,11 @@ struct InstallableCommand : virtual Args, SourceExprCommand
 
     void prepare() override;
 
+    std::optional<FlakeRef> getFlakeRefForCompletion() override
+    {
+        return parseFlakeRef(_installable, absPath("."));
+    }
+
 private:
 
     std::string _installable{"."};
diff --git a/src/nix/flake.cc b/src/nix/flake.cc
index 33bcc9604..57ce996cd 100644
--- a/src/nix/flake.cc
+++ b/src/nix/flake.cc
@@ -20,7 +20,7 @@
 using namespace nix;
 using namespace nix::flake;
 
-class FlakeCommand : virtual Args, public EvalCommand, public MixFlakeOptions
+class FlakeCommand : virtual Args, public MixFlakeOptions
 {
     std::string flakeUrl = ".";
 
@@ -53,6 +53,11 @@ public:
     {
         return flake::lockFlake(*getEvalState(), getFlakeRef(), lockFlags);
     }
+
+    std::optional<FlakeRef> getFlakeRefForCompletion() override
+    {
+        return getFlakeRef();
+    }
 };
 
 static void printFlakeInfo(const Store & store, const Flake & flake)
diff --git a/src/nix/installables.cc b/src/nix/installables.cc
index 4b1e31000..583b9e021 100644
--- a/src/nix/installables.cc
+++ b/src/nix/installables.cc
@@ -18,6 +18,17 @@
 
 namespace nix {
 
+void completeFlakeInputPath(
+    ref<EvalState> evalState,
+    const FlakeRef & flakeRef,
+    std::string_view prefix)
+{
+    auto flake = flake::getFlake(*evalState, flakeRef, true);
+    for (auto & input : flake.inputs)
+        if (hasPrefix(input.first, prefix))
+            completions->insert(input.first);
+}
+
 MixFlakeOptions::MixFlakeOptions()
 {
     addFlag({
@@ -56,6 +67,10 @@ MixFlakeOptions::MixFlakeOptions()
         .labels = {"input-path"},
         .handler = {[&](std::string s) {
             lockFlags.inputUpdates.insert(flake::parseInputPath(s));
+        }},
+        .completer = {[&](size_t, std::string_view prefix) {
+            if (auto flakeRef = getFlakeRefForCompletion())
+                completeFlakeInputPath(getEvalState(), *flakeRef, prefix);
         }}
     });
 
@@ -707,6 +722,16 @@ void InstallablesCommand::prepare()
     installables = parseInstallables(getStore(), _installables);
 }
 
+std::optional<FlakeRef> InstallablesCommand::getFlakeRefForCompletion()
+{
+    if (_installables.empty()) {
+        if (useDefaultInstallables())
+            return parseFlakeRef(".", absPath("."));
+        return {};
+    }
+    return parseFlakeRef(_installables.front(), absPath("."));
+}
+
 InstallableCommand::InstallableCommand()
 {
     expectArgs({