1 files changed, 206 insertions, 0 deletions
diff --git a/src/mongo/base/dependency_graph.cpp b/src/mongo/base/dependency_graph.cpp
new file mode 100644
index 00000000000..8d9107b5a10
--- /dev/null
+++ b/src/mongo/base/dependency_graph.cpp
@@ -0,0 +1,206 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/base/dependency_graph.h"
+
+#include <algorithm>
+#include <fmt/format.h>
+#include <fmt/ranges.h>
+#include <iostream>
+#include <iterator>
+#include <random>
+#include <sstream>
+
+#include "mongo/util/assert_util.h"
+#include "mongo/util/string_map.h"
+
+namespace mongo {
+
+void DependencyGraph::addNode(std::string name,
+                              std::vector<std::string> prerequisites,
+                              std::vector<std::string> dependents,
+                              std::unique_ptr<Payload> payload) {
+    if (!payload) {
+        struct DummyPayload : Payload {};
+        payload = std::make_unique<DummyPayload>();
+    }
+    auto& newNode = _nodes[name];
+    uassert(ErrorCodes::Error(50999), name, !newNode.payload);  // collision
+    for (auto& otherNode : prerequisites)
+        newNode.prerequisites.insert(otherNode);
+    for (auto& otherNode : dependents)
+        _nodes[otherNode].prerequisites.insert(name);
+    newNode.payload = std::move(payload);
+}
+
+namespace {
+
+using namespace fmt::literals;
+
+template <typename Seq>
+void strAppendJoin(std::string& out, StringData separator, const Seq& sequence) {
+    StringData currSep;
+    for (StringData str : sequence) {
+        out.append(currSep.rawData(), currSep.size());
+        out.append(str.rawData(), str.size());
+        currSep = separator;
+    }
+}
+
+// In the case of a cycle, copy the cycle node names into `*cycle`.
+template <typename Iter>
+void throwGraphContainsCycle(Iter first, Iter last, std::vector<std::string>* cycle) {
+    std::vector<std::string> names;
+    std::transform(first, last, std::back_inserter(names), [](auto& e) { return e->name(); });
+    if (cycle)
+        *cycle = names;
+    names.push_back((*first)->name());
+    uasserted(ErrorCodes::GraphContainsCycle,
+              format(FMT_STRING("Cycle in dependency graph: {}"), fmt::join(names, " -> ")));
+}
+
+}  // namespace
+
+std::vector<std::string> DependencyGraph::topSort(std::vector<std::string>* cycle) const {
+    // Topological sort via repeated depth-first traversal.
+    // All nodes must have an initFn before running topSort, or we return BadValue.
+    struct Element {
+        const std::string& name() const {
+            return nodeIter->first;
+        }
+        stdx::unordered_map<std::string, Node>::const_iterator nodeIter;
+        std::vector<Element*> children;
+        std::vector<Element*>::iterator membership;  // Position of this in `elements`.
+    };
+
+    std::vector<Element> elementsStore;
+    std::vector<Element*> elements;
+
+    // Swap the pointers in the `elements` vector that point to `a` and `b`.
+    // Update their 'membership' data members to reflect the change.
+    auto swapPositions = [](Element& a, Element& b) {
+        using std::swap;
+        swap(*a.membership, *b.membership);
+        swap(a.membership, b.membership);
+    };
+
+    elementsStore.reserve(_nodes.size());
+    for (auto iter = _nodes.begin(); iter != _nodes.end(); ++iter) {
+        uassert(ErrorCodes::BadValue,
+                "node {} was mentioned but never added"_format(iter->first),
+                iter->second.payload);
+        elementsStore.push_back(Element{iter});
+    }
+
+    // Wire up all the child relationships by pointer rather than by string names.
+    {
+        StringMap<Element*> byName;
+        for (Element& e : elementsStore)
+            byName[e.name()] = &e;
+        for (Element& element : elementsStore) {
+            const auto& prereqs = element.nodeIter->second.prerequisites;
+            std::transform(prereqs.begin(),
+                           prereqs.end(),
+                           std::back_inserter(element.children),
+                           [&](StringData childName) {
+                               auto iter = byName.find(childName);
+                               uassert(ErrorCodes::BadValue,
+                                       "node {} depends on missing node {}"_format(
+                                           element.nodeIter->first, childName),
+                                       iter != byName.end());
+                               return iter->second;
+                           });
+        }
+    }
+
+    elements.reserve(_nodes.size());
+    std::transform(elementsStore.begin(),
+                   elementsStore.end(),
+                   std::back_inserter(elements),
+                   [](auto& e) { return &e; });
+
+    // Shuffle the inputs to improve test coverage of undeclared dependencies.
+    {
+        std::random_device slowSeedGen;
+        std::mt19937 generator(slowSeedGen());
+        std::shuffle(elements.begin(), elements.end(), generator);
+        for (Element* e : elements)
+            std::shuffle(e->children.begin(), e->children.end(), generator);
+    }
+
+    // Initialize all the `membership` iterators. Must only happen after shuffle.
+    for (auto iter = elements.begin(); iter != elements.end(); ++iter)
+        (*iter)->membership = iter;
+
+    // The `elements` sequence is divided into 3 regions:
+    // elements:        [ sorted | unsorted | stack ]
+    //          unsortedBegin => [          )  <= unsortedEnd
+    // Each element of the stack region is a prerequisite of its neighbor to the right. Through
+    // 'swapPositions' calls and boundary increments, elements will transition from unsorted to
+    // stack to sorted. The unsorted region shinks to ultimately become an empty region on the
+    // right. No other moves are permitted.
+    auto unsortedBegin = elements.begin();
+    auto unsortedEnd = elements.end();
+
+    while (unsortedBegin != elements.end()) {
+        if (unsortedEnd == elements.end()) {
+            // The stack is empty but there's more work to do. Grow the stack region to enclose
+            // the rightmost unsorted element. Equivalent to pushing it.
+            --unsortedEnd;
+        }
+        auto top = unsortedEnd;
+        auto& children = (*top)->children;
+        if (!children.empty()) {
+            Element* picked = children.back();
+            children.pop_back();
+            if (picked->membership < unsortedBegin)
+                continue;
+            if (picked->membership >= unsortedEnd) {  // O(1) cycle detection
+                throwGraphContainsCycle(unsortedEnd, elements.end(), cycle);
+            }
+            swapPositions(**--unsortedEnd, *picked);  // unsorted push to stack
+            continue;
+        }
+        swapPositions(**unsortedEnd++, **unsortedBegin++);  // pop from stack to sorted
+    }
+    std::vector<std::string> sortedNames;
+    sortedNames.reserve(_nodes.size());
+    std::transform(elements.begin(),
+                   elements.end(),
+                   std::back_inserter(sortedNames),
+                   [](const Element* e) { return e->name(); });
+    return sortedNames;
+}
+
+DependencyGraph::Payload* DependencyGraph::find(const std::string& name) {
+    auto iter = _nodes.find(name);
+    return (iter == _nodes.end()) ? nullptr : iter->second.payload.get();
+}
+
+}  // namespace mongo