Improve flake references

1 files changed, 61 insertions, 46 deletions

diff --git a/src/libexpr/primops/flake.cc b/src/libexpr/primops/flake.cc
index 1367fa420..5e92b1da3 100644
--- a/src/libexpr/primops/flake.cc
+++ b/src/libexpr/primops/flake.cc
@@ -1,3 +1,4 @@
+#include "flake.hh"
 #include "primops.hh"
 #include "eval-inline.hh"
 #include "fetchGit.hh"
@@ -9,7 +10,7 @@
 
 namespace nix {
 
-const EvalState::FlakeRegistry & EvalState::getFlakeRegistry()
+const FlakeRegistry & EvalState::getFlakeRegistry()
 {
     std::call_once(_flakeRegistryInit, [&]()
     {
@@ -33,10 +34,7 @@ const EvalState::FlakeRegistry & EvalState::getFlakeRegistry()
 
             auto flakes = json["flakes"];
             for (auto i = flakes.begin(); i != flakes.end(); ++i) {
-                FlakeRegistry::Entry entry;
-                entry.uri = i->value("uri", "");
-                if (entry.uri.empty())
-                    throw Error("invalid flake registry entry");
+                FlakeRegistry::Entry entry{FlakeRef(i->value("uri", ""))};
                 _flakeRegistry->entries.emplace(i.key(), entry);
             }
         }
@@ -54,7 +52,7 @@ static void prim_flakeRegistry(EvalState & state, const Pos & pos, Value * * arg
     for (auto & entry : registry.entries) {
         auto vEntry = state.allocAttr(v, entry.first);
         state.mkAttrs(*vEntry, 2);
-        mkString(*state.allocAttr(*vEntry, state.symbols.create("uri")), entry.second.uri);
+        mkString(*state.allocAttr(*vEntry, state.symbols.create("uri")), entry.second.ref.to_string());
         vEntry->attrs->sort();
     }
 
@@ -63,44 +61,53 @@ static void prim_flakeRegistry(EvalState & state, const Pos & pos, Value * * arg
 
 static RegisterPrimOp r1("__flakeRegistry", 0, prim_flakeRegistry);
 
+static FlakeRef lookupFlake(EvalState & state, const FlakeRef & flakeRef)
+{
+    if (auto refData = std::get_if<FlakeRef::IsFlakeId>(&flakeRef.data)) {
+        auto registry = state.getFlakeRegistry();
+        auto i = registry.entries.find(refData->id);
+        if (i == registry.entries.end())
+            throw Error("cannot find flake '%s' in the flake registry", refData->id);
+        auto newRef = FlakeRef(i->second.ref);
+        if (!newRef.isDirect())
+            throw Error("found indirect flake URI '%s' in the flake registry", i->second.ref.to_string());
+        return newRef;
+    } else
+        return flakeRef;
+}
+
 struct Flake
 {
-    std::string name;
+    FlakeId id;
     std::string description;
     Path path;
     std::set<std::string> requires;
     Value * vProvides; // FIXME: gc
+    // commit hash
+    // date
+    // content hash
 };
 
-std::regex flakeRegex("^flake:([a-zA-Z][a-zA-Z0-9_-]*)(/[a-zA-Z][a-zA-Z0-9_.-]*)?$");
-std::regex githubRegex("^github:([a-zA-Z][a-zA-Z0-9_-]*)/([a-zA-Z][a-zA-Z0-9_-]*)(/([a-zA-Z][a-zA-Z0-9_-]*))?$");
-
-static Path fetchFlake(EvalState & state, const std::string & flakeUri)
+static Path fetchFlake(EvalState & state, const FlakeRef & flakeRef)
 {
-    std::smatch match;
-
-    if (std::regex_match(flakeUri, match, flakeRegex)) {
-        auto flakeName = match[1];
-        auto revOrRef = match[2];
-        auto registry = state.getFlakeRegistry();
-        auto i = registry.entries.find(flakeName);
-        if (i == registry.entries.end())
-            throw Error("unknown flake '%s'", flakeName);
-        return fetchFlake(state, i->second.uri);
-    }
-
-    else if (std::regex_match(flakeUri, match, githubRegex)) {
-        auto owner = match[1];
-        auto repo = match[2];
-        auto revOrRef = match[4].str();
-        if (revOrRef.empty()) revOrRef = "master";
+    assert(flakeRef.isDirect());
 
+    if (auto refData = std::get_if<FlakeRef::IsGitHub>(&flakeRef.data)) {
         // FIXME: require hash in pure mode.
 
         // FIXME: use regular /archive URLs instead? api.github.com
         // might have stricter rate limits.
+
+        // FIXME: support passing auth tokens for private repos.
+
         auto storePath = getDownloader()->downloadCached(state.store,
-            fmt("https://api.github.com/repos/%s/%s/tarball/%s", owner, repo, revOrRef),
+            fmt("https://api.github.com/repos/%s/%s/tarball/%s",
+                refData->owner, refData->repo,
+                refData->rev
+                  ? refData->rev->to_string(Base16, false)
+                  : refData->ref
+                    ? *refData->ref
+                    : "master"),
             true, "source");
 
         // FIXME: extract revision hash from ETag.
@@ -108,18 +115,18 @@ static Path fetchFlake(EvalState & state, const std::string & flakeUri)
         return storePath;
     }
 
-    else if (hasPrefix(flakeUri, "/") || hasPrefix(flakeUri, "git://")) {
-        auto gitInfo = exportGit(state.store, flakeUri, {}, "", "source");
+    else if (auto refData = std::get_if<FlakeRef::IsGit>(&flakeRef.data)) {
+        auto gitInfo = exportGit(state.store, refData->uri, refData->ref,
+            refData->rev ? refData->rev->to_string(Base16, false) : "", "source");
         return gitInfo.storePath;
     }
 
-    else
-        throw Error("unsupported flake URI '%s'", flakeUri);
+    else abort();
 }
 
-static Flake getFlake(EvalState & state, const std::string & flakeUri)
+static Flake getFlake(EvalState & state, const FlakeRef & flakeRef)
 {
-    auto flakePath = fetchFlake(state, flakeUri);
+    auto flakePath = fetchFlake(state, flakeRef);
     state.store->assertStorePath(flakePath);
 
     Flake flake;
@@ -130,7 +137,7 @@ static Flake getFlake(EvalState & state, const std::string & flakeUri)
     state.forceAttrs(vInfo);
 
     if (auto name = vInfo.attrs->get(state.sName))
-        flake.name = state.forceStringNoCtx(*(**name).value, *(**name).pos);
+        flake.id = state.forceStringNoCtx(*(**name).value, *(**name).pos);
     else
         throw Error("flake lacks attribute 'name'");
 
@@ -153,23 +160,31 @@ static Flake getFlake(EvalState & state, const std::string & flakeUri)
     return flake;
 }
 
-static std::map<std::string, Flake> resolveFlakes(EvalState & state, const StringSet & flakeUris)
+/* Given a set of flake references, recursively fetch them and their
+   dependencies. */
+static std::map<FlakeId, Flake> resolveFlakes(EvalState & state, const std::vector<FlakeRef> & flakeRefs)
 {
-    std::map<std::string, Flake> done;
-    std::queue<std::string> todo;
-    for (auto & i : flakeUris) todo.push(i);
+    std::map<FlakeId, Flake> done;
+    std::queue<FlakeRef> todo;
+    for (auto & i : flakeRefs) todo.push(i);
 
     while (!todo.empty()) {
-        auto flakeUri = todo.front();
+        auto flakeRef = todo.front();
         todo.pop();
-        if (done.count(flakeUri)) continue;
 
-        auto flake = getFlake(state, flakeUri);
+        if (auto refData = std::get_if<FlakeRef::IsFlakeId>(&flakeRef.data)) {
+            if (done.count(refData->id)) continue; // optimization
+            flakeRef = lookupFlake(state, flakeRef);
+        }
+
+        auto flake = getFlake(state, flakeRef);
+
+        if (done.count(flake.id)) continue;
 
         for (auto & require : flake.requires)
             todo.push(require);
 
-        done.emplace(flake.name, flake);
+        done.emplace(flake.id, flake);
     }
 
     return done;
@@ -177,7 +192,7 @@ static std::map<std::string, Flake> resolveFlakes(EvalState & state, const Strin
 
 static void prim_getFlake(EvalState & state, const Pos & pos, Value * * args, Value & v)
 {
-    std::string flakeUri = state.forceStringNoCtx(*args[0], pos);
+    auto flakeUri = FlakeRef(state.forceStringNoCtx(*args[0], pos));
 
     auto flakes = resolveFlakes(state, {flakeUri});
 
@@ -186,7 +201,7 @@ static void prim_getFlake(EvalState & state, const Pos & pos, Value * * args, Va
     state.mkAttrs(*vResult, flakes.size());
 
     for (auto & flake : flakes) {
-        auto vFlake = state.allocAttr(*vResult, flake.second.name);
+        auto vFlake = state.allocAttr(*vResult, flake.second.id);
         state.mkAttrs(*vFlake, 2);
         mkString(*state.allocAttr(*vFlake, state.sDescription), flake.second.description);
         auto vProvides = state.allocAttr(*vFlake, state.symbols.create("provides"));