Fixed Flake data type and flake fetching

3 files changed, 101 insertions, 92 deletions

diff --git a/src/libexpr/primops/flake.cc b/src/libexpr/primops/flake.cc
index ab797fe01..89267e2b1 100644
--- a/src/libexpr/primops/flake.cc
+++ b/src/libexpr/primops/flake.cc
@@ -204,28 +204,32 @@ static FlakeRef lookupFlake(EvalState & state, const FlakeRef & flakeRef, const
     return flakeRef;
 }
 
-static FlakeSourceInfo fetchFlake(EvalState & state, const FlakeRef fRef, bool impureIsAllowed = false)
+// Lookups happen here too
+static SourceInfo fetchFlake(EvalState & state, const FlakeRef & flakeRef, bool impureIsAllowed = false)
 {
-    if (evalSettings.pureEval && !impureIsAllowed && !fRef.isImmutable())
-        throw Error("requested to fetch mutable flake '%s' in pure mode", fRef);
+    FlakeRef resolvedRef = lookupFlake(state, flakeRef,
+        impureIsAllowed ? state.getFlakeRegistries() : std::vector<std::shared_ptr<FlakeRegistry>>());
+
+    if (evalSettings.pureEval && !impureIsAllowed && !resolvedRef.isImmutable())
+        throw Error("requested to fetch mutable flake '%s' in pure mode", resolvedRef);
 
     // This only downloads only one revision of the repo, not the entire history.
-    if (auto refData = std::get_if<FlakeRef::IsGitHub>(&fRef.data)) {
+    if (auto refData = std::get_if<FlakeRef::IsGitHub>(&resolvedRef.data)) {
 
         // FIXME: use regular /archive URLs instead? api.github.com
         // might have stricter rate limits.
 
         auto url = fmt("https://api.github.com/repos/%s/%s/tarball/%s",
             refData->owner, refData->repo,
-            fRef.rev ? fRef.rev->to_string(Base16, false)
-                : fRef.ref ? *fRef.ref : "master");
+            resolvedRef.rev ? resolvedRef.rev->to_string(Base16, false)
+                : resolvedRef.ref ? *resolvedRef.ref : "master");
 
         std::string accessToken = settings.githubAccessToken.get();
         if (accessToken != "")
             url += "?access_token=" + accessToken;
 
         auto result = getDownloader()->downloadCached(state.store, url, true, "source",
-            Hash(), nullptr, fRef.rev ? 1000000000 : settings.tarballTtl);
+            Hash(), nullptr, resolvedRef.rev ? 1000000000 : settings.tarballTtl);
 
         if (!result.etag)
             throw Error("did not receive an ETag header from '%s'", url);
@@ -233,72 +237,60 @@ static FlakeSourceInfo fetchFlake(EvalState & state, const FlakeRef fRef, bool i
         if (result.etag->size() != 42 || (*result.etag)[0] != '"' || (*result.etag)[41] != '"')
             throw Error("ETag header '%s' from '%s' is not a Git revision", *result.etag, url);
 
-        FlakeSourceInfo info(fRef);
+        std::string rev = std::string(*result.etag, 1, result.etag->size() - 2);
+        const FlakeRef ref(resolvedRef.baseRef().to_string() + "/" + rev);
+        SourceInfo info(ref);
         info.storePath = result.path;
-        info.rev = Hash(std::string(*result.etag, 1, result.etag->size() - 2), htSHA1);
-        info.flakeRef.rev = info.rev;
-        info.flakeRef.ref = {};
 
         return info;
     }
 
     // This downloads the entire git history
-    else if (auto refData = std::get_if<FlakeRef::IsGit>(&fRef.data)) {
-        auto gitInfo = exportGit(state.store, refData->uri, fRef.ref, fRef.rev, "source");
-        FlakeSourceInfo info(fRef);
+    else if (auto refData = std::get_if<FlakeRef::IsGit>(&resolvedRef.data)) {
+        auto gitInfo = exportGit(state.store, refData->uri, resolvedRef.ref, resolvedRef.rev, "source");
+        const FlakeRef ref(resolvedRef.baseRef().to_string() + "/" + gitInfo.ref + "/" + gitInfo.rev.to_string(Base16, false));
+        SourceInfo info(ref);
         info.storePath = gitInfo.storePath;
-        info.rev = gitInfo.rev;
         info.revCount = gitInfo.revCount;
-        info.flakeRef.ref = gitInfo.ref;
-        info.flakeRef.rev = info.rev;
         return info;
     }
 
-    else if (auto refData = std::get_if<FlakeRef::IsPath>(&fRef.data)) {
+    else if (auto refData = std::get_if<FlakeRef::IsPath>(&resolvedRef.data)) {
         if (!pathExists(refData->path + "/.git"))
             throw Error("flake '%s' does not reference a Git repository", refData->path);
         auto gitInfo = exportGit(state.store, refData->path, {}, {}, "source");
-        FlakeSourceInfo info(fRef);
+        const FlakeRef ref(resolvedRef.baseRef().to_string() + "/" + gitInfo.ref + "/" + gitInfo.rev.to_string(Base16, false));
+        SourceInfo info(ref);
         info.storePath = gitInfo.storePath;
-        info.rev = gitInfo.rev;
         info.revCount = gitInfo.revCount;
-        info.flakeRef.ref = gitInfo.ref;
-        info.flakeRef.rev = info.rev;
         return info;
     }
 
     else abort();
 }
 
-// This will return the flake which corresponds to a given FlakeRef. The lookupFlake is done within this function.
+// This will return the flake which corresponds to a given FlakeRef. The lookupFlake is done within `fetchFlake`, which is used here.
 Flake getFlake(EvalState & state, const FlakeRef & flakeRef, bool impureIsAllowed = false)
 {
-    FlakeRef resolvedRef = lookupFlake(state, flakeRef,
-        impureIsAllowed ? state.getFlakeRegistries() : std::vector<std::shared_ptr<FlakeRegistry>>());
+    SourceInfo sourceInfo = fetchFlake(state, flakeRef, impureIsAllowed);
+    debug("got flake source '%s' with flakeref %s", sourceInfo.storePath, sourceInfo.resolvedRef.to_string());
 
-    FlakeSourceInfo sourceInfo = fetchFlake(state, resolvedRef, impureIsAllowed);
-    debug("got flake source '%s' with revision %s",
-        sourceInfo.storePath, sourceInfo.rev.value_or(Hash(htSHA1)).to_string(Base16, false));
+    FlakeRef resolvedRef = sourceInfo.resolvedRef;
 
-    resolvedRef = sourceInfo.flakeRef; // `resolvedRef` is now immutable
+    resolvedRef = sourceInfo.resolvedRef; // `resolvedRef` is now immutable
 
     state.store->assertStorePath(sourceInfo.storePath);
 
     if (state.allowedPaths)
         state.allowedPaths->insert(sourceInfo.storePath);
 
-    Flake flake(resolvedRef, std::move(sourceInfo));
-    if (std::get_if<FlakeRef::IsGitHub>(&resolvedRef.data)) {
-        // FIXME: ehm?
-        if (flake.sourceInfo.rev)
-            flake.ref = FlakeRef(resolvedRef.baseRef().to_string()
-                + "/" + flake.sourceInfo.rev->to_string(Base16, false));
-    }
-
     // Guard against symlink attacks.
-    auto flakeFile = canonPath(sourceInfo.storePath + "/" + resolvedRef.subdir + "/flake.nix");
+    Path flakeFile = canonPath(sourceInfo.storePath + "/" + resolvedRef.subdir + "/flake.nix");
     if (!isInDir(flakeFile, sourceInfo.storePath))
         throw Error("flake file '%s' escapes from '%s'", resolvedRef, sourceInfo.storePath);
+
+    Flake flake(flakeRef, sourceInfo);
+
     if (!pathExists(flakeFile))
         throw Error("source tree referenced by '%s' does not contain a '%s/flake.nix' file", resolvedRef, resolvedRef.subdir);
 
@@ -344,24 +336,18 @@ Flake getFlake(EvalState & state, const FlakeRef & flakeRef, bool impureIsAllowe
 // Get the `NonFlake` corresponding to a `FlakeRef`.
 NonFlake getNonFlake(EvalState & state, const FlakeRef & flakeRef, FlakeAlias alias)
 {
-    FlakeSourceInfo sourceInfo = fetchFlake(state, flakeRef);
-    debug("got non-flake source '%s' with revision %s",
-        sourceInfo.storePath, sourceInfo.rev.value_or(Hash(htSHA1)).to_string(Base16, false));
+    SourceInfo sourceInfo = fetchFlake(state, flakeRef);
+    debug("got non-flake source '%s' with flakeref %s", sourceInfo.storePath, sourceInfo.resolvedRef.to_string());
 
-    auto flakePath = sourceInfo.storePath;
-    state.store->assertStorePath(flakePath);
+    FlakeRef resolvedRef = sourceInfo.resolvedRef;
 
-    if (state.allowedPaths)
-        state.allowedPaths->insert(flakePath);
+    NonFlake nonFlake(flakeRef, sourceInfo);
 
-    NonFlake nonFlake(flakeRef);
-    if (std::get_if<FlakeRef::IsGitHub>(&flakeRef.data)) {
-        if (sourceInfo.rev)
-            nonFlake.ref = FlakeRef(flakeRef.baseRef().to_string()
-                + "/" + sourceInfo.rev->to_string(Base16, false));
-    }
+    nonFlake.storePath = sourceInfo.storePath;
+    state.store->assertStorePath(nonFlake.storePath);
 
-    nonFlake.path = flakePath;
+    if (state.allowedPaths)
+        state.allowedPaths->insert(nonFlake.storePath);
 
     nonFlake.alias = alias;
 
@@ -380,7 +366,7 @@ ResolvedFlake resolveFlake(EvalState & state, const FlakeRef & topRef,
     LockFile lockFile;
 
     if (isTopFlake)
-        lockFile = readLockFile(flake.sourceInfo.storePath + "/" + flake.ref.subdir + "/flake.lock"); // FIXME: symlink attack
+        lockFile = readLockFile(flake.storePath + flake.resolvedRef.subdir + "/flake.lock"); // FIXME: symlink attack
 
     ResolvedFlake deps(flake);
 
@@ -399,13 +385,13 @@ ResolvedFlake resolveFlake(EvalState & state, const FlakeRef & topRef,
 
 LockFile::FlakeEntry dependenciesToFlakeEntry(const ResolvedFlake & resolvedFlake)
 {
-    LockFile::FlakeEntry entry(resolvedFlake.flake.sourceInfo.flakeRef);
+    LockFile::FlakeEntry entry(resolvedFlake.flake.resolvedRef);
 
     for (auto & newResFlake : resolvedFlake.flakeDeps)
-        entry.flakeEntries.insert_or_assign(newResFlake.flake.id, dependenciesToFlakeEntry(newResFlake));
+        entry.flakeEntries.insert_or_assign(newResFlake.flake.originalRef, dependenciesToFlakeEntry(newResFlake));
 
     for (auto & nonFlake : resolvedFlake.nonFlakeDeps)
-        entry.nonFlakeEntries.insert_or_assign(nonFlake.alias, nonFlake.ref);
+        entry.nonFlakeEntries.insert_or_assign(nonFlake.alias, nonFlake.resolvedRef);
 
     return entry;
 }
@@ -453,18 +439,18 @@ void callFlake(EvalState & state, const ResolvedFlake & resFlake, Value & v)
         auto vNonFlake = state.allocAttr(v, nonFlake.alias);
         state.mkAttrs(*vNonFlake, 4);
 
-        state.store->isValidPath(nonFlake.path);
-        mkString(*state.allocAttr(*vNonFlake, state.sOutPath), nonFlake.path, {nonFlake.path});
+        state.store->isValidPath(nonFlake.storePath);
+        mkString(*state.allocAttr(*vNonFlake, state.sOutPath), nonFlake.storePath, {nonFlake.storePath});
     }
 
     mkString(*state.allocAttr(v, state.sDescription), resFlake.flake.description);
 
-    auto & path = resFlake.flake.sourceInfo.storePath;
+    auto & path = resFlake.flake.storePath;
     state.store->isValidPath(path);
     mkString(*state.allocAttr(v, state.sOutPath), path, {path});
 
-    if (resFlake.flake.sourceInfo.revCount)
-        mkInt(*state.allocAttr(v, state.symbols.create("revCount")), *resFlake.flake.sourceInfo.revCount);
+    if (resFlake.flake.revCount)
+        mkInt(*state.allocAttr(v, state.symbols.create("revCount")), *resFlake.flake.revCount);
 
     auto vProvides = state.allocAttr(v, state.symbols.create("provides"));
     mkApp(*vProvides, *resFlake.flake.vProvides, v);
diff --git a/src/libexpr/primops/flake.hh b/src/libexpr/primops/flake.hh
index 6329c36ec..f3bf76cc7 100644
--- a/src/libexpr/primops/flake.hh
+++ b/src/libexpr/primops/flake.hh
@@ -23,7 +23,7 @@ struct LockFile
     {
         FlakeRef ref;
         std::map<FlakeRef, FlakeEntry> flakeEntries;
-        std::map<FlakeId, FlakeRef> nonFlakeEntries;
+        std::map<FlakeAlias, FlakeRef> nonFlakeEntries;
         FlakeEntry(const FlakeRef & flakeRef) : ref(flakeRef) {};
     };
 
@@ -43,37 +43,43 @@ std::shared_ptr<FlakeRegistry> readRegistry(const Path &);
 
 void writeRegistry(const FlakeRegistry &, const Path &);
 
-struct FlakeSourceInfo
+struct SourceInfo
 {
-    FlakeRef flakeRef;
+    FlakeRef resolvedRef;
     Path storePath;
-    std::optional<Hash> rev;
     std::optional<uint64_t> revCount;
     // date
-    FlakeSourceInfo(const FlakeRef & flakeRef) : flakeRef(flakeRef) { }
+    SourceInfo(const FlakeRef & resolvRef) : resolvedRef(resolvRef) {};
 };
 
 struct Flake
 {
     FlakeId id;
-    FlakeRef ref;
+    FlakeRef originalRef;
+    FlakeRef resolvedRef;
     std::string description;
-    FlakeSourceInfo sourceInfo;
+    std::optional<uint64_t> revCount;
+    Path storePath;
     std::vector<FlakeRef> requires;
     std::map<FlakeAlias, FlakeRef> nonFlakeRequires;
     Value * vProvides; // FIXME: gc
-    Flake(const FlakeRef & flakeRef, FlakeSourceInfo && sourceInfo)
-        : ref(flakeRef), sourceInfo(sourceInfo) {};
+    // date
+    // content hash
+    Flake(const FlakeRef & origRef, const SourceInfo & sourceInfo) : originalRef(origRef),
+        resolvedRef(sourceInfo.resolvedRef), revCount(sourceInfo.revCount), storePath(sourceInfo.storePath) {};
 };
 
 struct NonFlake
 {
     FlakeAlias alias;
-    FlakeRef ref;
-    Path path;
+    FlakeRef originalRef;
+    FlakeRef resolvedRef;
+    std::optional<uint64_t> revCount;
+    Path storePath;
     // date
     // content hash
-    NonFlake(const FlakeRef flakeRef) : ref(flakeRef) {};
+    NonFlake(const FlakeRef & origRef, const SourceInfo & sourceInfo) : originalRef(origRef),
+        resolvedRef(sourceInfo.resolvedRef), revCount(sourceInfo.revCount), storePath(sourceInfo.storePath) {};
 };
 
 std::shared_ptr<FlakeRegistry> getGlobalRegistry();
diff --git a/src/nix/flake.cc b/src/nix/flake.cc
index e38c4db0b..0af368570 100644
--- a/src/nix/flake.cc
+++ b/src/nix/flake.cc
@@ -45,35 +45,53 @@ void printFlakeInfo(Flake & flake, bool json) {
     if (json) {
         nlohmann::json j;
         j["id"] = flake.id;
-        j["uri"] = flake.sourceInfo.flakeRef.to_string();
+        j["uri"] = flake.resolvedRef.to_string();
         j["description"] = flake.description;
-        if (flake.sourceInfo.rev)
-            j["revision"] = flake.sourceInfo.rev->to_string(Base16, false);
-        if (flake.sourceInfo.revCount)
-            j["revCount"] = *flake.sourceInfo.revCount;
-        j["path"] = flake.sourceInfo.storePath;
+        if (flake.resolvedRef.ref)
+            j["branch"] = *flake.resolvedRef.ref;
+        if (flake.resolvedRef.rev)
+            j["revision"] = flake.resolvedRef.rev->to_string(Base16, false);
+        if (flake.revCount)
+            j["revCount"] = *flake.revCount;
+        j["path"] = flake.storePath;
         std::cout << j.dump(4) << std::endl;
     } else {
         std::cout << "ID:          " << flake.id << "\n";
-        std::cout << "URI:         " << flake.sourceInfo.flakeRef << "\n";
+        std::cout << "URI:         " << flake.resolvedRef.to_string() << "\n";
         std::cout << "Description: " << flake.description << "\n";
-        if (flake.sourceInfo.rev)
-            std::cout << "Revision:    " << flake.sourceInfo.rev->to_string(Base16, false) << "\n";
-        if (flake.sourceInfo.revCount)
-            std::cout << "Revcount:    " << *flake.sourceInfo.revCount << "\n";
-        std::cout << "Path:        " << flake.sourceInfo.storePath << "\n";
+        if (flake.resolvedRef.ref)
+            std::cout << "Branch:      " << *flake.resolvedRef.ref;
+        if (flake.resolvedRef.rev)
+            std::cout << "Revision:    " << flake.resolvedRef.rev->to_string(Base16, false) << "\n";
+        if (flake.revCount)
+            std::cout << "Revcount:    " << *flake.revCount << "\n";
+        std::cout << "Path:        " << flake.storePath << "\n";
     }
 }
 
 void printNonFlakeInfo(NonFlake & nonFlake, bool json) {
     if (json) {
         nlohmann::json j;
-        j["name"] = nonFlake.alias;
-        j["location"] = nonFlake.path;
+        j["id"] = nonFlake.alias;
+        j["uri"] = nonFlake.resolvedRef.to_string();
+        if (nonFlake.resolvedRef.ref)
+            j["branch"] = *nonFlake.resolvedRef.ref;
+        if (nonFlake.resolvedRef.rev)
+            j["revision"] = nonFlake.resolvedRef.rev->to_string(Base16, false);
+        if (nonFlake.revCount)
+            j["revCount"] = *nonFlake.revCount;
+        j["path"] = nonFlake.storePath;
         std::cout << j.dump(4) << std::endl;
     } else {
-        std::cout << "name:        " << nonFlake.alias << "\n";
-        std::cout << "Location:    " << nonFlake.path << "\n";
+        std::cout << "ID:          " << nonFlake.alias << "\n";
+        std::cout << "URI:         " << nonFlake.resolvedRef.to_string() << "\n";
+        if (nonFlake.resolvedRef.ref)
+            std::cout << "Branch:      " << *nonFlake.resolvedRef.ref;
+        if (nonFlake.resolvedRef.rev)
+            std::cout << "Revision:    " << nonFlake.resolvedRef.rev->to_string(Base16, false) << "\n";
+        if (nonFlake.revCount)
+            std::cout << "Revcount:    " << *nonFlake.revCount << "\n";
+        std::cout << "Path:        " << nonFlake.storePath << "\n";
     }
 }
 
@@ -244,14 +262,13 @@ struct CmdFlakePin : virtual Args, StoreCommand, MixEvalArgs
         FlakeRegistry userRegistry = *readRegistry(userRegistryPath);
         auto it = userRegistry.entries.find(FlakeRef(alias));
         if (it != userRegistry.entries.end()) {
-            it->second = getFlake(*evalState, it->second, true).ref;
-            // The 'ref' in 'flake' is immutable.
+            it->second = getFlake(*evalState, it->second, true).resolvedRef;
             writeRegistry(userRegistry, userRegistryPath);
         } else {
             std::shared_ptr<FlakeRegistry> globalReg = getGlobalRegistry();
             it = globalReg->entries.find(FlakeRef(alias));
             if (it != globalReg->entries.end()) {
-                FlakeRef newRef = getFlake(*evalState, it->second, true).ref;
+                FlakeRef newRef = getFlake(*evalState, it->second, true).resolvedRef;
                 userRegistry.entries.insert_or_assign(alias, newRef);
                 writeRegistry(userRegistry, userRegistryPath);
             } else