Merge pull request #14540 from lovesegfault/pre-compute-outputgraph

perf(libstore/derivation-builder): pre-compute outputGraph for linear complexity

Author: Ericson2314

src/libstore/local-store.cc

@@ -989,10 +989,10 @@ void LocalStore::registerValidPaths(const ValidPathInfos & infos)
            error if a cycle is detected and roll back the
            transaction.  Cycles can only occur when a derivation
            has multiple outputs. */
-        auto topoSortResult = topoSort(paths, {[&](const StorePath & path) {
-                                           auto i = infos.find(path);
-                                           return i == infos.end() ? StorePathSet() : i->second.references;
-                                       }});
+        auto topoSortResult = topoSort(paths, [&](const StorePath & path) {
+            auto i = infos.find(path);
+            return i == infos.end() ? StorePathSet() : i->second.references;
+        });
 
         std::visit(
             overloaded{

src/libstore/misc.cc

@@ -313,13 +313,13 @@ MissingPaths Store::queryMissing(const std::vector<DerivedPath> & targets)
 
 StorePaths Store::topoSortPaths(const StorePathSet & paths)
 {
-    auto result = topoSort(paths, {[&](const StorePath & path) {
-                               try {
-                                   return queryPathInfo(path)->references;
-                               } catch (InvalidPath &) {
-                                   return StorePathSet();
-                               }
-                           }});
+    auto result = topoSort(paths, [&](const StorePath & path) {
+        try {
+            return queryPathInfo(path)->references;
+        } catch (InvalidPath &) {
+            return StorePathSet();
+        }
+    });
 
     return std::visit(
         overloaded{

src/libstore/unix/build/derivation-builder.cc

@@ -1396,8 +1396,18 @@ SingleDrvOutputs DerivationBuilderImpl::registerOutputs()
     struct PerhapsNeedToRegister
     {
         StorePathSet refs;
+        /**
+         * References to other outputs. Built by looking up in
+         * `scratchOutputsInverse`.
+         */
+        StringSet otherOutputs;
     };
 
+    /* inverse map of scratchOutputs for efficient lookup */
+    std::map<StorePath, std::string> scratchOutputsInverse;
+    for (auto & [outputName, path] : scratchOutputs)
+        scratchOutputsInverse.insert_or_assign(path, outputName);
+
     std::map<std::string, std::variant<AlreadyRegistered, PerhapsNeedToRegister>> outputReferencesIfUnregistered;
     std::map<std::string, struct stat> outputStats;
     for (auto & [outputName, _] : drv.outputs) {
@@ -1466,36 +1476,40 @@ SingleDrvOutputs DerivationBuilderImpl::registerOutputs()
             references = scanForReferences(blank, actualPath, referenceablePaths);
         }
 
-        outputReferencesIfUnregistered.insert_or_assign(outputName, PerhapsNeedToRegister{.refs = references});
+        StringSet referencedOutputs;
+        for (auto & r : references)
+            if (auto * o = get(scratchOutputsInverse, r))
+                referencedOutputs.insert(*o);
+
+        outputReferencesIfUnregistered.insert_or_assign(
+            outputName,
+            PerhapsNeedToRegister{
+                .refs = references,
+                .otherOutputs = referencedOutputs,
+            });
         outputStats.insert_or_assign(outputName, std::move(st));
     }
 
-    auto topoSortResult = topoSort(outputsToSort, {[&](const std::string & name) {
-                                       auto orifu = get(outputReferencesIfUnregistered, name);
-                                       if (!orifu)
-                                           throw BuildError(
-                                               BuildResult::Failure::OutputRejected,
-                                               "no output reference for '%s' in build of '%s'",
-                                               name,
-                                               store.printStorePath(drvPath));
-                                       return std::visit(
-                                           overloaded{
-                                               /* Since we'll use the already installed versions of these, we
-                                                  can treat them as leaves and ignore any references they
-                                                  have. */
-                                               [&](const AlreadyRegistered &) { return StringSet{}; },
-                                               [&](const PerhapsNeedToRegister & refs) {
-                                                   StringSet referencedOutputs;
-                                                   /* FIXME build inverted map up front so no quadratic waste here */
-                                                   for (auto & r : refs.refs)
-                                                       for (auto & [o, p] : scratchOutputs)
-                                                           if (r == p)
-                                                               referencedOutputs.insert(o);
-                                                   return referencedOutputs;
-                                               },
-                                           },
-                                           *orifu);
-                                   }});
+    StringSet emptySet;
+
+    auto topoSortResult = topoSort(outputsToSort, [&](const std::string & name) -> const StringSet & {
+        auto * orifu = get(outputReferencesIfUnregistered, name);
+        if (!orifu)
+            throw BuildError(
+                BuildResult::Failure::OutputRejected,
+                "no output reference for '%s' in build of '%s'",
+                name,
+                store.printStorePath(drvPath));
+        return std::visit(
+            overloaded{
+                /* Since we'll use the already installed versions of these, we
+                   can treat them as leaves and ignore any references they
+                   have. */
+                [&](const AlreadyRegistered &) -> const StringSet & { return emptySet; },
+                [&](const PerhapsNeedToRegister & refs) -> const StringSet & { return refs.otherOutputs; },
+            },
+            *orifu);
+    });
 
     auto sortedOutputNames = std::visit(
         overloaded{

src/libutil/include/nix/util/topo-sort.hh

@@ -3,6 +3,7 @@
 
 #include "nix/util/error.hh"
 #include <variant>
+#include <concepts>
 
 namespace nix {
 
@@ -16,8 +17,9 @@ struct Cycle
 template<typename T>
 using TopoSortResult = std::variant<std::vector<T>, Cycle<T>>;
 
-template<typename T, typename Compare>
-TopoSortResult<T> topoSort(std::set<T, Compare> items, std::function<std::set<T, Compare>(const T &)> getChildren)
+template<typename T, typename Compare, std::invocable<const T &> F>
+    requires std::same_as<std::remove_cvref_t<std::invoke_result_t<F, const T &>>, std::set<T, Compare>>
+TopoSortResult<T> topoSort(std::set<T, Compare> items, F && getChildren)
 {
     std::vector<T> sorted;
     decltype(items) visited, parents;
@@ -34,7 +36,7 @@ TopoSortResult<T> topoSort(std::set<T, Compare> items, std::function<std::set<T,
         }
         parents.insert(path);
 
-        auto references = getChildren(path);
+        auto && references = std::invoke(getChildren, path);
 
         for (auto & i : references)
             /* Don't traverse into items that don't exist in our starting set. */