From e8f145ae691802498d30fbf7c4bcbaaefbe6946a Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Tue, 24 Jun 2025 15:14:40 +0200
Subject: [PATCH] nix flake prefetch-inputs: Add

This command fetches all inputs of a flake in parallel.

Example runtime for

  $ chmod -R u+w /tmp/nix2; rm -rf /tmp/nix2; rm ~/.cache/nix/fetcher-cache-v3.sqlite*; rm -rf ~/.cache/nix/tarball-cache/ ~/.cache/nix/gitv3/; time nix flake prefetch-inputs --store /tmp/nix2 https://api.flakehub.com/f/pinned/informalsystems/cosmos.nix/0.3.0/018ce9ed-d0be-7ce5-81b6-a3c6e3ae1187/source.tar.gz

with http-connections = 1:

real    4m11.859s
user    2m6.931s
sys     0m25.619s

and http-connections = 25 (the default):

real    0m57.146s
user    2m49.506s
sys     0m36.008s
---
 src/nix/flake-prefetch-inputs.md | 17 ++++++++++
 src/nix/flake.cc                 | 56 ++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)
 create mode 100644 src/nix/flake-prefetch-inputs.md

diff --git a/src/nix/flake-prefetch-inputs.md b/src/nix/flake-prefetch-inputs.md
new file mode 100644
index 000000000..a69f7d367
--- /dev/null
+++ b/src/nix/flake-prefetch-inputs.md
@@ -0,0 +1,17 @@
+R""(
+
+# Examples
+
+* Fetch the inputs of the `hydra` flake:
+
+  ```console
+  # nix flake prefetch-inputs github:NixOS/hydra
+  ```
+
+# Description
+
+Fetch the inputs of a flake. This ensures that they are already available for any subsequent evaluation of the flake.
+
+This operation is recursive: it will fetch not just the direct inputs of the top-level flake, but also transitive inputs.
+
+)""
diff --git a/src/nix/flake.cc b/src/nix/flake.cc
index 48e2ae392..c2aa442bc 100644
--- a/src/nix/flake.cc
+++ b/src/nix/flake.cc
@@ -19,6 +19,8 @@
 #include "nix/util/users.hh"
 #include "nix/fetchers/fetch-to-store.hh"
 #include "nix/store/local-fs-store.hh"
+#include "nix/util/thread-pool.hh"
+#include "nix/store/filetransfer.hh"
 
 #include <filesystem>
 #include <nlohmann/json.hpp>
@@ -1140,6 +1142,59 @@ struct CmdFlakeArchive : FlakeCommand, MixJSON, MixDryRun
     }
 };
 
+struct CmdFlakePrefetchInputs : FlakeCommand
+{
+    std::string description() override
+    {
+        return "fetch the inputs of a flake";
+    }
+
+    std::string doc() override
+    {
+        return
+          #include "flake-prefetch-inputs.md"
+          ;
+    }
+
+    void run(nix::ref<nix::Store> store) override
+    {
+        auto flake = lockFlake();
+
+        ThreadPool pool{fileTransferSettings.httpConnections};
+
+        struct State
+        {
+            std::set<const Node *> done;
+        };
+
+        Sync<State> state_;
+
+        std::function<void(const Node & node)> visit;
+        visit = [&](const Node & node)
+        {
+            if (!state_.lock()->done.insert(&node).second)
+                return;
+
+            if (auto lockedNode = dynamic_cast<const LockedNode *>(&node)) {
+                Activity act(*logger, lvlInfo, actUnknown,
+                    fmt("fetching '%s'", lockedNode->lockedRef));
+                auto accessor = lockedNode->lockedRef.input.getAccessor(store).first;
+                if (!evalSettings.lazyTrees)
+                    fetchToStore(*store, accessor, FetchMode::Copy, lockedNode->lockedRef.input.getName());
+            }
+
+            for (auto & [inputName, input] : node.inputs) {
+                if (auto inputNode = std::get_if<0>(&input))
+                    pool.enqueue(std::bind(visit, **inputNode));
+            }
+        };
+
+        pool.enqueue(std::bind(visit, *flake.lockFile.root));
+
+        pool.process();
+    }
+};
+
 struct CmdFlakeShow : FlakeCommand, MixJSON
 {
     bool showLegacy = false;
@@ -1543,6 +1598,7 @@ struct CmdFlake : NixMultiCommand
                 {"new", []() { return make_ref<CmdFlakeNew>(); }},
                 {"clone", []() { return make_ref<CmdFlakeClone>(); }},
                 {"archive", []() { return make_ref<CmdFlakeArchive>(); }},
+                {"prefetch-inputs", []() { return make_ref<CmdFlakePrefetchInputs>(); }},
                 {"show", []() { return make_ref<CmdFlakeShow>(); }},
                 {"prefetch", []() { return make_ref<CmdFlakePrefetch>(); }},
             })