From 259ff74bdeeaac5a987a4eb88654ba80d9553543 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Mon, 11 May 2020 21:49:02 +0200
Subject: [PATCH] Add completion for installables

This completes flakerefs using the registry (e.g. 'nix<TAB>' => 'nix
nixpkgs') and flake output attributes by evaluating the flake
(e.g. 'dwarffs#nix<TAB>' => 'dwarffs#nixosModules').
---
 src/nix/command.hh      |   7 ++-
 src/nix/installables.cc | 100 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/src/nix/command.hh b/src/nix/command.hh
index 0738c0f91..6a6c3fed9 100644
--- a/src/nix/command.hh
+++ b/src/nix/command.hh
@@ -63,6 +63,8 @@ struct SourceExprCommand : virtual Args, EvalCommand, MixFlakeOptions
     virtual Strings getDefaultFlakeAttrPaths();
 
     virtual Strings getDefaultFlakeAttrPathPrefixes();
+
+    void completeInstallable(std::string_view prefix);
 };
 
 enum RealiseMode { Build, NoBuild, DryRun };
@@ -89,10 +91,7 @@ struct InstallableCommand : virtual Args, SourceExprCommand
 {
     std::shared_ptr<Installable> installable;
 
-    InstallableCommand()
-    {
-        expectArg("installable", &_installable, true);
-    }
+    InstallableCommand();
 
     void prepare() override;
 
diff --git a/src/nix/installables.cc b/src/nix/installables.cc
index abc642e11..ade8f83f9 100644
--- a/src/nix/installables.cc
+++ b/src/nix/installables.cc
@@ -11,6 +11,7 @@
 #include "flake/flake.hh"
 #include "eval-cache.hh"
 #include "url.hh"
+#include "registry.hh"
 
 #include <regex>
 #include <queue>
@@ -106,6 +107,91 @@ Strings SourceExprCommand::getDefaultFlakeAttrPathPrefixes()
     };
 }
 
+void SourceExprCommand::completeInstallable(std::string_view prefix)
+{
+    completePath(0, prefix);
+
+    if (file) return; // FIXME
+
+    /* Look for flake output attributes that match the
+       prefix. */
+    try {
+        auto hash = prefix.find('#');
+        if (hash != std::string::npos) {
+            auto fragment = prefix.substr(hash + 1);
+            auto flakeRefS = std::string(prefix.substr(0, hash));
+            // FIXME: do tilde expansion.
+            auto flakeRef = parseFlakeRef(flakeRefS, absPath("."));
+
+            auto state = getEvalState();
+
+            auto evalCache = openEvalCache(*state,
+                std::make_shared<flake::LockedFlake>(lockFlake(*state, flakeRef, lockFlags)),
+                true);
+
+            auto root = evalCache->getRoot();
+
+            /* Complete 'fragment' relative to all the
+               attrpath prefixes as well as the root of the
+               flake. */
+            auto attrPathPrefixes = getDefaultFlakeAttrPathPrefixes();
+            attrPathPrefixes.push_back("");
+
+            for (auto & attrPathPrefixS : attrPathPrefixes) {
+                auto attrPathPrefix = parseAttrPath(*state, attrPathPrefixS);
+                auto attrPathS = attrPathPrefixS + std::string(fragment);
+                auto attrPath = parseAttrPath(*state, attrPathS);
+
+                std::string lastAttr;
+                if (!attrPath.empty() && !hasSuffix(attrPathS, ".")) {
+                    lastAttr = attrPath.back();
+                    attrPath.pop_back();
+                }
+
+                auto attr = root->findAlongAttrPath(attrPath);
+                if (!attr) continue;
+
+                auto attrs = attr->getAttrs();
+                for (auto & attr2 : attrs) {
+                    if (hasPrefix(attr2, lastAttr)) {
+                        auto attrPath2 = attr->getAttrPath(attr2);
+                        /* Strip the attrpath prefix. */
+                        attrPath2.erase(attrPath2.begin(), attrPath2.begin() + attrPathPrefix.size());
+                        completions->insert(flakeRefS + "#" + concatStringsSep(".", attrPath2));
+                    }
+                }
+            }
+
+            /* And add an empty completion for the default
+               attrpaths. */
+            if (fragment.empty()) {
+                for (auto & attrPath : getDefaultFlakeAttrPaths()) {
+                    auto attr = root->findAlongAttrPath(parseAttrPath(*state, attrPath));
+                    if (!attr) continue;
+                    completions->insert(flakeRefS + "#");
+                }
+            }
+        }
+    } catch (Error & e) {
+        warn(e.msg());
+    }
+
+    /* Look for registry entries that match the prefix. */
+    for (auto & registry : fetchers::getRegistries(getStore())) {
+        for (auto & entry : registry->entries) {
+            auto from = entry.from->to_string();
+            if (!hasPrefix(prefix, "flake:") && hasPrefix(from, "flake:")) {
+                std::string from2(from, 6);
+                if (hasPrefix(from2, prefix))
+                    completions->insert(from2);
+            } else {
+                if (hasPrefix(from, prefix))
+                    completions->insert(from);
+            }
+        }
+    }
+}
+
 ref<EvalState> EvalCommand::getEvalState()
 {
     if (!evalState)
@@ -576,6 +662,9 @@ InstallablesCommand::InstallablesCommand()
     expectArgs({
         .label = "installables",
         .handler = {&_installables},
+        .completer = {[&](size_t, std::string_view prefix) {
+            completeInstallable(prefix);
+        }}
     });
 }
 
@@ -588,6 +677,17 @@ void InstallablesCommand::prepare()
     installables = parseInstallables(getStore(), _installables);
 }
 
+InstallableCommand::InstallableCommand()
+{
+    expectArgs({
+        .label = "installable",
+        .handler = {&_installable},
+        .completer = {[&](size_t, std::string_view prefix) {
+            completeInstallable(prefix);
+        }}
+    });
+}
+
 void InstallableCommand::prepare()
 {
     installable = parseInstallable(getStore(), _installable);