Represent 'follows' inputs explicitly in the lock file

This fixes an issue where lockfile generation was not idempotent:
after updating a lockfile, a "follows" node would end up pointing to a
new copy of the node, rather than to the original node.

5 files changed, 134 insertions, 86 deletions

diff --git a/src/libexpr/flake/call-flake.nix b/src/libexpr/flake/call-flake.nix
index 2084e3fb3..932ac5e90 100644
--- a/src/libexpr/flake/call-flake.nix
+++ b/src/libexpr/flake/call-flake.nix
@@ -8,14 +8,41 @@ let
     builtins.mapAttrs
       (key: node:
         let
+
           sourceInfo =
             if key == lockFile.root
             then rootSrc
             else fetchTree (node.info or {} // removeAttrs node.locked ["dir"]);
+
           subdir = if key == lockFile.root then rootSubdir else node.locked.dir or "";
+
           flake = import (sourceInfo + (if subdir != "" then "/" else "") + subdir + "/flake.nix");
-          inputs = builtins.mapAttrs (inputName: key: allNodes.${key}) (node.inputs or {});
+
+          inputs = builtins.mapAttrs
+            (inputName: inputSpec: allNodes.${resolveInput inputSpec})
+            (node.inputs or {});
+
+          # Resolve a input spec into a node name. An input spec is
+          # either a node name, or a 'follows' path from the root
+          # node.
+          resolveInput = inputSpec:
+              if builtins.isList inputSpec
+              then getInputByPath lockFile.root inputSpec
+              else inputSpec;
+
+          # Follow an input path (e.g. ["dwarffs" "nixpkgs"]) from the
+          # root node, returning the final node.
+          getInputByPath = nodeName: path:
+            if path == []
+            then nodeName
+            else
+              getInputByPath
+                # Since this could be a 'follows' input, call resolveInput.
+                (resolveInput lockFile.nodes.${nodeName}.inputs.${builtins.head path})
+                (builtins.tail path);
+
           outputs = flake.outputs (inputs // { self = result; });
+
           result = outputs // sourceInfo // { inherit inputs; inherit outputs; inherit sourceInfo; };
         in
           if node.flake or true then
diff --git a/src/libexpr/flake/flake.cc b/src/libexpr/flake/flake.cc
index 4538c03ff..b99e4794a 100644
--- a/src/libexpr/flake/flake.cc
+++ b/src/libexpr/flake/flake.cc
@@ -90,9 +90,7 @@ static FlakeInput parseFlakeInput(EvalState & state,
 {
     expectType(state, tAttrs, *value, pos);
 
-    FlakeInput input {
-        .ref = FlakeRef::fromAttrs({{"type", "indirect"}, {"id", inputName}})
-    };
+    FlakeInput input;
 
     auto sInputs = state.symbols.create("inputs");
     auto sUrl = state.symbols.create("url");
@@ -145,6 +143,9 @@ static FlakeInput parseFlakeInput(EvalState & state,
             input.ref = parseFlakeRef(*url, {}, true);
     }
 
+    if (!input.follows && !input.ref)
+        input.ref = FlakeRef::fromAttrs({{"type", "indirect"}, {"id", inputName}});
+
     return input;
 }
 
@@ -276,7 +277,6 @@ LockedFlake lockFlake(
     LockFile newLockFile;
 
     std::vector<FlakeRef> parents;
-    std::map<InputPath, InputPath> follows;
 
     std::function<void(
         const FlakeInputs & flakeInputs,
@@ -324,34 +324,36 @@ LockedFlake lockFlake(
             /* Resolve 'follows' later (since it may refer to an input
                path we haven't processed yet. */
             if (input.follows) {
-                if (hasOverride)
+                InputPath target;
+                if (hasOverride || input.absolute)
                     /* 'follows' from an override is relative to the
                        root of the graph. */
-                    follows.insert_or_assign(inputPath, *input.follows);
+                    target = *input.follows;
                 else {
                     /* Otherwise, it's relative to the current flake. */
-                    InputPath path(inputPathPrefix);
-                    for (auto & i : *input.follows) path.push_back(i);
-                    debug("input '%s' follows '%s'", inputPathS, printInputPath(path));
-                    follows.insert_or_assign(inputPath, path);
+                    target = inputPathPrefix;
+                    for (auto & i : *input.follows) target.push_back(i);
                 }
+                debug("input '%s' follows '%s'", inputPathS, printInputPath(target));
+                node->inputs.insert_or_assign(id, target);
                 continue;
             }
 
+            assert(input.ref);
+
             /* Do we have an entry in the existing lock file? And we
                don't have a --update-input flag for this input? */
-            std::shared_ptr<const LockedNode> oldLock;
+            std::shared_ptr<LockedNode> oldLock;
 
             updatesUsed.insert(inputPath);
 
-            if (oldNode && !lockFlags.inputUpdates.count(inputPath)) {
-                auto oldLockIt = oldNode->inputs.find(id);
-                if (oldLockIt != oldNode->inputs.end())
-                    oldLock = std::dynamic_pointer_cast<const LockedNode>(oldLockIt->second);
-            }
+            if (oldNode && !lockFlags.inputUpdates.count(inputPath))
+                if (auto oldLock2 = get(oldNode->inputs, id))
+                    if (auto oldLock3 = std::get_if<0>(&*oldLock2))
+                        oldLock = *oldLock3;
 
             if (oldLock
-                && oldLock->originalRef == input.ref
+                && oldLock->originalRef == *input.ref
                 && !hasOverride)
             {
                 debug("keeping existing input '%s'", inputPathS);
@@ -386,18 +388,16 @@ LockedFlake lockFlake(
                     FlakeInputs fakeInputs;
 
                     for (auto & i : oldLock->inputs) {
-                        auto lockedNode = std::dynamic_pointer_cast<LockedNode>(i.second);
-                        // Note: this node is not locked in case
-                        // of a circular reference back to the root.
-                        if (lockedNode)
+                        if (auto lockedNode = std::get_if<0>(&i.second)) {
                             fakeInputs.emplace(i.first, FlakeInput {
-                                .ref = lockedNode->originalRef,
-                                .isFlake = lockedNode->isFlake,
+                                .ref = (*lockedNode)->originalRef,
+                                .isFlake = (*lockedNode)->isFlake,
+                            });
+                        } else if (auto follows = std::get_if<1>(&i.second)) {
+                            fakeInputs.emplace(i.first, FlakeInput {
+                                .follows = *follows,
+                                .absolute = true
                             });
-                        else {
-                            InputPath path(inputPath);
-                            path.push_back(i.first);
-                            follows.insert_or_assign(path, InputPath());
                         }
                     }
 
@@ -409,11 +409,11 @@ LockedFlake lockFlake(
                    this input. */
                 debug("creating new input '%s'", inputPathS);
 
-                if (!lockFlags.allowMutable && !input.ref.input.isImmutable())
+                if (!lockFlags.allowMutable && !input.ref->input.isImmutable())
                     throw Error("cannot update flake input '%s' in pure mode", inputPathS);
 
                 if (input.isFlake) {
-                    auto inputFlake = getFlake(state, input.ref, lockFlags.useRegistries, flakeCache);
+                    auto inputFlake = getFlake(state, *input.ref, lockFlags.useRegistries, flakeCache);
 
                     /* Note: in case of an --override-input, we use
                        the *original* ref (input2.ref) for the
@@ -423,15 +423,15 @@ LockedFlake lockFlake(
                        file. That is, overrides are sticky unless you
                        use --no-write-lock-file. */
                     auto childNode = std::make_shared<LockedNode>(
-                        inputFlake.lockedRef, input2.ref);
+                        inputFlake.lockedRef, input2.ref ? *input2.ref : *input.ref);
 
                     node->inputs.insert_or_assign(id, childNode);
 
                     /* Guard against circular flake imports. */
                     for (auto & parent : parents)
-                        if (parent == input.ref)
+                        if (parent == *input.ref)
                             throw Error("found circular import of flake '%s'", parent);
-                    parents.push_back(input.ref);
+                    parents.push_back(*input.ref);
                     Finally cleanup([&]() { parents.pop_back(); });
 
                     /* Recursively process the inputs of this
@@ -448,9 +448,9 @@ LockedFlake lockFlake(
 
                 else {
                     auto [sourceInfo, resolvedRef, lockedRef] = fetchOrSubstituteTree(
-                        state, input.ref, lockFlags.useRegistries, flakeCache);
+                        state, *input.ref, lockFlags.useRegistries, flakeCache);
                     node->inputs.insert_or_assign(id,
-                        std::make_shared<LockedNode>(lockedRef, input.ref, false));
+                        std::make_shared<LockedNode>(lockedRef, *input.ref, false));
                 }
             }
         }
@@ -460,29 +460,6 @@ LockedFlake lockFlake(
         flake.inputs, newLockFile.root, {},
         lockFlags.recreateLockFile ? nullptr : oldLockFile.root);
 
-    /* Insert edges for 'follows' overrides. */
-    for (auto & [from, to] : follows) {
-        debug("adding 'follows' node from '%s' to '%s'",
-            printInputPath(from),
-            printInputPath(to));
-
-        assert(!from.empty());
-
-        InputPath fromParent(from);
-        fromParent.pop_back();
-
-        auto fromParentNode = newLockFile.root->findInput(fromParent);
-        assert(fromParentNode);
-
-        auto toNode = newLockFile.root->findInput(to);
-        if (!toNode)
-            throw Error("flake input '%s' follows non-existent flake input '%s'",
-                printInputPath(from),
-                printInputPath(to));
-
-        fromParentNode->inputs.insert_or_assign(from.back(), toNode);
-    }
-
     for (auto & i : lockFlags.inputOverrides)
         if (!overridesUsed.count(i.first))
             warn("the flag '--override-input %s %s' does not match any input",
@@ -514,9 +491,13 @@ LockedFlake lockFlake(
 
                     bool lockFileExists = pathExists(path);
 
-                    if (lockFileExists)
-                        warn("updating lock file '%s':\n%s", path, chomp(diff));
-                    else
+                    if (lockFileExists) {
+                        auto s = chomp(diff);
+                        if (s.empty())
+                            warn("updating lock file '%s'", path);
+                        else
+                            warn("updating lock file '%s':\n%s", path, s);
+                    } else
                         warn("creating lock file '%s'", path);
 
                     newLockFile.write(path);
diff --git a/src/libexpr/flake/flake.hh b/src/libexpr/flake/flake.hh
index ebf81362c..77f3abdeb 100644
--- a/src/libexpr/flake/flake.hh
+++ b/src/libexpr/flake/flake.hh
@@ -19,9 +19,10 @@ typedef std::map<FlakeId, FlakeInput> FlakeInputs;
 
 struct FlakeInput
 {
-    FlakeRef ref;
+    std::optional<FlakeRef> ref;
     bool isFlake = true;
     std::optional<InputPath> follows;
+    bool absolute = false; // whether 'follows' is relative to the flake root
     FlakeInputs overrides;
 };
 
diff --git a/src/libexpr/flake/lockfile.cc b/src/libexpr/flake/lockfile.cc
index 9ba12bff7..acde0ab7f 100644
--- a/src/libexpr/flake/lockfile.cc
+++ b/src/libexpr/flake/lockfile.cc
@@ -41,15 +41,22 @@ StorePath LockedNode::computeStorePath(Store & store) const
     return lockedRef.input.computeStorePath(store);
 }
 
-std::shared_ptr<Node> Node::findInput(const InputPath & path)
+std::shared_ptr<Node> LockFile::findInput(const InputPath & path)
 {
-    auto pos = shared_from_this();
+    auto pos = root;
+
+    if (!pos) return {};
 
     for (auto & elem : path) {
-        auto i = pos->inputs.find(elem);
-        if (i == pos->inputs.end())
+        if (auto i = get(pos->inputs, elem)) {
+            if (auto node = std::get_if<0>(&*i))
+                pos = *node;
+            else if (auto follows = std::get_if<1>(&*i)) {
+                pos = findInput(*follows);
+                if (!pos) return {};
+            }
+        } else
             return {};
-        pos = i->second;
     }
 
     return pos;
@@ -58,7 +65,7 @@ std::shared_ptr<Node> Node::findInput(const InputPath & path)
 LockFile::LockFile(const nlohmann::json & json, const Path & path)
 {
     auto version = json.value("version", 0);
-    if (version < 5 || version > 6)
+    if (version < 5 || version > 7)
         throw Error("lock file '%s' has unsupported version %d", path, version);
 
     std::unordered_map<std::string, std::shared_ptr<Node>> nodeMap;
@@ -69,21 +76,37 @@ LockFile::LockFile(const nlohmann::json & json, const Path & path)
     {
         if (jsonNode.find("inputs") == jsonNode.end()) return;
         for (auto & i : jsonNode["inputs"].items()) {
-            std::string inputKey = i.value();
-            auto k = nodeMap.find(inputKey);
-            if (k == nodeMap.end()) {
-                auto jsonNode2 = json["nodes"][inputKey];
-                auto input = std::make_shared<LockedNode>(jsonNode2);
-                k = nodeMap.insert_or_assign(inputKey, input).first;
-                getInputs(*input, jsonNode2);
+            if (i.value().is_array()) {
+                InputPath path;
+                for (auto & j : i.value())
+                    path.push_back(j);
+                node.inputs.insert_or_assign(i.key(), path);
+            } else {
+                std::string inputKey = i.value();
+                auto k = nodeMap.find(inputKey);
+                if (k == nodeMap.end()) {
+                    auto jsonNode2 = json["nodes"][inputKey];
+                    auto input = std::make_shared<LockedNode>(jsonNode2);
+                    k = nodeMap.insert_or_assign(inputKey, input).first;
+                    getInputs(*input, jsonNode2);
+                }
+                if (auto child = std::dynamic_pointer_cast<LockedNode>(k->second))
+                    node.inputs.insert_or_assign(i.key(), child);
+                else
+                    // FIXME: replace by follows node
+                    throw Error("lock file contains cycle to root node");
             }
-            node.inputs.insert_or_assign(i.key(), k->second);
         }
     };
 
     std::string rootKey = json["root"];
     nodeMap.insert_or_assign(rootKey, root);
     getInputs(*root, json["nodes"][rootKey]);
+
+    // FIXME: check that there are no cycles in version >= 7. Cycles
+    // between inputs are only possible using 'follows' indirections.
+    // Once we drop support for version <= 6, we can simplify the code
+    // a bit since we don't need to worry about cycles.
 }
 
 nlohmann::json LockFile::toJson() const
@@ -116,8 +139,16 @@ nlohmann::json LockFile::toJson() const
 
         if (!node->inputs.empty()) {
             auto inputs = nlohmann::json::object();
-            for (auto & i : node->inputs)
-                inputs[i.first] = dumpNode(i.first, i.second);
+            for (auto & i : node->inputs) {
+                if (auto child = std::get_if<0>(&i.second)) {
+                    inputs[i.first] = dumpNode(i.first, *child);
+                } else if (auto follows = std::get_if<1>(&i.second)) {
+                    auto arr = nlohmann::json::array();
+                    for (auto & x : *follows)
+                        arr.push_back(x);
+                    inputs[i.first] = std::move(arr);
+                }
+            }
             n["inputs"] = std::move(inputs);
         }
 
@@ -133,7 +164,7 @@ nlohmann::json LockFile::toJson() const
     };
 
     nlohmann::json json;
-    json["version"] = 6;
+    json["version"] = 7;
     json["root"] = dumpNode("root", root);
     json["nodes"] = std::move(nodes);
 
@@ -172,7 +203,9 @@ bool LockFile::isImmutable() const
     visit = [&](std::shared_ptr<const Node> node)
     {
         if (!nodes.insert(node).second) return;
-        for (auto & i : node->inputs) visit(i.second);
+        for (auto & i : node->inputs)
+            if (auto child = std::get_if<0>(&i.second))
+                visit(*child);
     };
 
     visit(root);
@@ -216,9 +249,11 @@ static void flattenLockFile(
     for (auto &[id, input] : node->inputs) {
         auto inputPath(prefix);
         inputPath.push_back(id);
-        if (auto lockedInput = std::dynamic_pointer_cast<const LockedNode>(input))
-            res.emplace(inputPath, lockedInput);
-        flattenLockFile(input, inputPath, done, res);
+        if (auto child = std::get_if<0>(&input)) {
+            if (auto lockedInput = std::dynamic_pointer_cast<const LockedNode>(*child))
+                res.emplace(inputPath, lockedInput);
+            flattenLockFile(*child, inputPath, done, res);
+        }
     }
 }
 
diff --git a/src/libexpr/flake/lockfile.hh b/src/libexpr/flake/lockfile.hh
index eb99ed997..04ac80f56 100644
--- a/src/libexpr/flake/lockfile.hh
+++ b/src/libexpr/flake/lockfile.hh
@@ -15,16 +15,18 @@ using namespace fetchers;
 
 typedef std::vector<FlakeId> InputPath;
 
+struct LockedNode;
+
 /* A node in the lock file. It has outgoing edges to other nodes (its
    inputs). Only the root node has this type; all other nodes have
    type LockedNode. */
 struct Node : std::enable_shared_from_this<Node>
 {
-    std::map<FlakeId, std::shared_ptr<Node>> inputs;
+    typedef std::variant<std::shared_ptr<LockedNode>, InputPath> Edge;
 
-    virtual ~Node() { }
+    std::map<FlakeId, Edge> inputs;
 
-    std::shared_ptr<Node> findInput(const InputPath & path);
+    virtual ~Node() { }
 };
 
 /* A non-root node in the lock file. */
@@ -63,6 +65,8 @@ struct LockFile
     bool isImmutable() const;
 
     bool operator ==(const LockFile & other) const;
+
+    std::shared_ptr<Node> findInput(const InputPath & path);
 };
 
 std::ostream & operator <<(std::ostream & stream, const LockFile & lockFile);