3 files changed, 52 insertions, 23 deletions

diff --git a/src/libexpr/flake/flake.cc b/src/libexpr/flake/flake.cc
index b99e4794a..fa3185d08 100644
--- a/src/libexpr/flake/flake.cc
+++ b/src/libexpr/flake/flake.cc
@@ -469,12 +469,15 @@ LockedFlake lockFlake(
         if (!updatesUsed.count(i))
             warn("the flag '--update-input %s' does not match any input", printInputPath(i));
 
+    /* Check 'follows' inputs. */
+    newLockFile.check();
+
     debug("new lock file: %s", newLockFile);
 
     /* Check whether we need to / can write the new lock file. */
     if (!(newLockFile == oldLockFile)) {
 
-        auto diff = diffLockFiles(oldLockFile, newLockFile);
+        auto diff = LockFile::diff(oldLockFile, newLockFile);
 
         if (lockFlags.writeLockFile) {
             if (auto sourcePath = topRef.input.getSourcePath()) {
diff --git a/src/libexpr/flake/lockfile.cc b/src/libexpr/flake/lockfile.cc
index fd5a68825..a74846944 100644
--- a/src/libexpr/flake/lockfile.cc
+++ b/src/libexpr/flake/lockfile.cc
@@ -238,21 +238,29 @@ InputPath parseInputPath(std::string_view s)
     return path;
 }
 
-static void flattenLockFile(
-    std::shared_ptr<const Node> node,
-    const InputPath & prefix,
-    std::unordered_set<std::shared_ptr<const Node>> & done,
-    std::map<InputPath, Node::Edge> & res)
+std::map<InputPath, Node::Edge> LockFile::getAllInputs() const
 {
-    if (!done.insert(node).second) return;
-
-    for (auto &[id, input] : node->inputs) {
-        auto inputPath(prefix);
-        inputPath.push_back(id);
-        res.emplace(inputPath, input);
-        if (auto child = std::get_if<0>(&input))
-            flattenLockFile(*child, inputPath, done, res);
-    }
+    std::unordered_set<std::shared_ptr<Node>> done;
+    std::map<InputPath, Node::Edge> res;
+
+    std::function<void(const InputPath & prefix, std::shared_ptr<Node> node)> recurse;
+
+    recurse = [&](const InputPath & prefix, std::shared_ptr<Node> node)
+    {
+        if (!done.insert(node).second) return;
+
+        for (auto &[id, input] : node->inputs) {
+            auto inputPath(prefix);
+            inputPath.push_back(id);
+            res.emplace(inputPath, input);
+            if (auto child = std::get_if<0>(&input))
+                recurse(inputPath, *child);
+        }
+    };
+
+    recurse({}, root);
+
+    return res;
 }
 
 std::ostream & operator <<(std::ostream & stream, const Node::Edge & edge)
@@ -275,13 +283,10 @@ static bool equals(const Node::Edge & e1, const Node::Edge & e2)
     return false;
 }
 
-std::string diffLockFiles(const LockFile & oldLocks, const LockFile & newLocks)
+std::string LockFile::diff(const LockFile & oldLocks, const LockFile & newLocks)
 {
-    std::unordered_set<std::shared_ptr<const Node>> done;
-    std::map<InputPath, Node::Edge> oldFlat, newFlat;
-    flattenLockFile(oldLocks.root, {}, done, oldFlat);
-    done.clear();
-    flattenLockFile(newLocks.root, {}, done, newFlat);
+    auto oldFlat = oldLocks.getAllInputs();
+    auto newFlat = newLocks.getAllInputs();
 
     auto i = oldFlat.begin();
     auto j = newFlat.begin();
@@ -309,6 +314,22 @@ std::string diffLockFiles(const LockFile & oldLocks, const LockFile & newLocks)
     return res;
 }
 
+void LockFile::check()
+{
+    auto inputs = getAllInputs();
+
+    for (auto & [inputPath, input] : inputs) {
+        if (auto follows = std::get_if<1>(&input)) {
+            if (!follows->empty() && !get(inputs, *follows))
+                throw Error("input '%s' follows a non-existent input '%s'",
+                    printInputPath(inputPath),
+                    printInputPath(*follows));
+        }
+    }
+}
+
+void check();
+
 std::string printInputPath(const InputPath & path)
 {
     return concatStringsSep("/", path);
diff --git a/src/libexpr/flake/lockfile.hh b/src/libexpr/flake/lockfile.hh
index 04ac80f56..5e7cfda3e 100644
--- a/src/libexpr/flake/lockfile.hh
+++ b/src/libexpr/flake/lockfile.hh
@@ -67,6 +67,13 @@ struct LockFile
     bool operator ==(const LockFile & other) const;
 
     std::shared_ptr<Node> findInput(const InputPath & path);
+
+    std::map<InputPath, Node::Edge> getAllInputs() const;
+
+    static std::string diff(const LockFile & oldLocks, const LockFile & newLocks);
+
+    /* Check that every 'follows' input target exists. */
+    void check();
 };
 
 std::ostream & operator <<(std::ostream & stream, const LockFile & lockFile);
@@ -75,6 +82,4 @@ InputPath parseInputPath(std::string_view s);
 
 std::string printInputPath(const InputPath & path);
 
-std::string diffLockFiles(const LockFile & oldLocks, const LockFile & newLocks);
-
 }