From 649c2db308f16ce4b2cbefe4a8760577541cfb47 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Mon, 11 May 2020 22:10:33 +0200
Subject: [PATCH] nix flake: Add completion support

---
 src/nix/command.hh      |  2 ++
 src/nix/flake.cc        |  9 ++++++++-
 src/nix/installables.cc | 26 +++++++++++++++++---------
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/nix/command.hh b/src/nix/command.hh
index 6a6c3fed9..faa19c8ea 100644
--- a/src/nix/command.hh
+++ b/src/nix/command.hh
@@ -38,6 +38,8 @@ struct EvalCommand : virtual StoreCommand, MixEvalArgs
     ref<EvalState> getEvalState();
 
     std::shared_ptr<EvalState> evalState;
+
+    void completeFlakeRef(std::string_view prefix);
 };
 
 struct MixFlakeOptions : virtual Args
diff --git a/src/nix/flake.cc b/src/nix/flake.cc
index 6eee781aa..b6cc7eb54 100644
--- a/src/nix/flake.cc
+++ b/src/nix/flake.cc
@@ -28,7 +28,14 @@ public:
 
     FlakeCommand()
     {
-        expectArg("flake-url", &flakeUrl, true);
+        expectArgs({
+            .label = "flake-url",
+            .optional = true,
+            .handler = {&flakeUrl},
+            .completer = {[&](size_t, std::string_view prefix) {
+                completeFlakeRef(prefix);
+            }}
+        });
     }
 
     FlakeRef getFlakeRef()
diff --git a/src/nix/installables.cc b/src/nix/installables.cc
index 5ba48d69a..2d23b33b7 100644
--- a/src/nix/installables.cc
+++ b/src/nix/installables.cc
@@ -109,8 +109,6 @@ Strings SourceExprCommand::getDefaultFlakeAttrPathPrefixes()
 
 void SourceExprCommand::completeInstallable(std::string_view prefix)
 {
-    completeDir(0, prefix);
-
     if (file) return; // FIXME
 
     /* Look for flake output attributes that match the
@@ -176,6 +174,23 @@ void SourceExprCommand::completeInstallable(std::string_view prefix)
         warn(e.msg());
     }
 
+    completeFlakeRef(prefix);
+}
+
+ref<EvalState> EvalCommand::getEvalState()
+{
+    if (!evalState)
+        evalState = std::make_shared<EvalState>(searchPath, getStore());
+    return ref<EvalState>(evalState);
+}
+
+void EvalCommand::completeFlakeRef(std::string_view prefix)
+{
+    if (prefix == "")
+        completions->insert(".");
+
+    completeDir(0, prefix);
+
     /* Look for registry entries that match the prefix. */
     for (auto & registry : fetchers::getRegistries(getStore())) {
         for (auto & entry : registry->entries) {
@@ -192,13 +207,6 @@ void SourceExprCommand::completeInstallable(std::string_view prefix)
     }
 }
 
-ref<EvalState> EvalCommand::getEvalState()
-{
-    if (!evalState)
-        evalState = std::make_shared<EvalState>(searchPath, getStore());
-    return ref<EvalState>(evalState);
-}
-
 Buildable Installable::toBuildable()
 {
     auto buildables = toBuildables();