1 files changed, 344 insertions, 0 deletions
diff --git a/js/public/UbiNodeShortestPaths.h b/js/public/UbiNodeShortestPaths.h
new file mode 100644
index 0000000000..35745ae9f9
--- /dev/null
+++ b/js/public/UbiNodeShortestPaths.h
@@ -0,0 +1,344 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*-
+ * vim: set ts=8 sts=2 et sw=2 tw=80:
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef js_UbiNodeShortestPaths_h
+#define js_UbiNodeShortestPaths_h
+
+#include "mozilla/Maybe.h"
+
+#include <utility>
+
+#include "js/AllocPolicy.h"
+#include "js/GCAPI.h"
+#include "js/UbiNode.h"
+#include "js/UbiNodeBreadthFirst.h"
+#include "js/UniquePtr.h"
+
+namespace JS {
+namespace ubi {
+
+/**
+ * A back edge along a path in the heap graph.
+ */
+struct JS_PUBLIC_API BackEdge {
+ private:
+  Node predecessor_;
+  EdgeName name_;
+
+ public:
+  using Ptr = js::UniquePtr<BackEdge>;
+
+  BackEdge() : predecessor_(), name_(nullptr) {}
+
+  [[nodiscard]] bool init(const Node& predecessor, Edge& edge) {
+    MOZ_ASSERT(!predecessor_);
+    MOZ_ASSERT(!name_);
+
+    predecessor_ = predecessor;
+    name_ = std::move(edge.name);
+    return true;
+  }
+
+  BackEdge(const BackEdge&) = delete;
+  BackEdge& operator=(const BackEdge&) = delete;
+
+  BackEdge(BackEdge&& rhs)
+      : predecessor_(rhs.predecessor_), name_(std::move(rhs.name_)) {
+    MOZ_ASSERT(&rhs != this);
+  }
+
+  BackEdge& operator=(BackEdge&& rhs) {
+    this->~BackEdge();
+    new (this) BackEdge(std::move(rhs));
+    return *this;
+  }
+
+  Ptr clone() const;
+
+  const EdgeName& name() const { return name_; }
+  EdgeName& name() { return name_; }
+
+  const JS::ubi::Node& predecessor() const { return predecessor_; }
+};
+
+/**
+ * A path is a series of back edges from which we discovered a target node.
+ */
+using Path = JS::ubi::Vector<BackEdge*>;
+
+/**
+ * The `JS::ubi::ShortestPaths` type represents a collection of up to N shortest
+ * retaining paths for each of a target set of nodes, starting from the same
+ * root node.
+ */
+struct JS_PUBLIC_API ShortestPaths {
+ private:
+  // Types, type aliases, and data members.
+
+  using BackEdgeVector = JS::ubi::Vector<BackEdge::Ptr>;
+  using NodeToBackEdgeVectorMap =
+      js::HashMap<Node, BackEdgeVector, js::DefaultHasher<Node>,
+                  js::SystemAllocPolicy>;
+
+  struct Handler;
+  using Traversal = BreadthFirst<Handler>;
+
+  /**
+   * A `JS::ubi::BreadthFirst` traversal handler that records back edges for
+   * how we reached each node, allowing us to reconstruct the shortest
+   * retaining paths after the traversal.
+   */
+  struct Handler {
+    using NodeData = BackEdge;
+
+    ShortestPaths& shortestPaths;
+    size_t totalMaxPathsToRecord;
+    size_t totalPathsRecorded;
+
+    explicit Handler(ShortestPaths& shortestPaths)
+        : shortestPaths(shortestPaths),
+          totalMaxPathsToRecord(shortestPaths.targets_.count() *
+                                shortestPaths.maxNumPaths_),
+          totalPathsRecorded(0) {}
+
+    bool operator()(Traversal& traversal, const JS::ubi::Node& origin,
+                    JS::ubi::Edge& edge, BackEdge* back, bool first) {
+      MOZ_ASSERT(back);
+      MOZ_ASSERT(origin == shortestPaths.root_ ||
+                 traversal.visited.has(origin));
+      MOZ_ASSERT(totalPathsRecorded < totalMaxPathsToRecord);
+
+      if (first && !back->init(origin, edge)) {
+        return false;
+      }
+
+      if (!shortestPaths.targets_.has(edge.referent)) {
+        return true;
+      }
+
+      // If `first` is true, then we moved the edge's name into `back` in
+      // the above call to `init`. So clone that back edge to get the
+      // correct edge name. If `first` is not true, then our edge name is
+      // still in `edge`. This accounts for the asymmetry between
+      // `back->clone()` in the first branch, and the `init` call in the
+      // second branch.
+
+      if (first) {
+        BackEdgeVector paths;
+        if (!paths.reserve(shortestPaths.maxNumPaths_)) {
+          return false;
+        }
+        auto cloned = back->clone();
+        if (!cloned) {
+          return false;
+        }
+        paths.infallibleAppend(std::move(cloned));
+        if (!shortestPaths.paths_.putNew(edge.referent, std::move(paths))) {
+          return false;
+        }
+        totalPathsRecorded++;
+      } else {
+        auto ptr = shortestPaths.paths_.lookup(edge.referent);
+        MOZ_ASSERT(ptr,
+                   "This isn't the first time we have seen the target node "
+                   "`edge.referent`. "
+                   "We should have inserted it into shortestPaths.paths_ the "
+                   "first time we "
+                   "saw it.");
+
+        if (ptr->value().length() < shortestPaths.maxNumPaths_) {
+          auto thisBackEdge = js::MakeUnique<BackEdge>();
+          if (!thisBackEdge || !thisBackEdge->init(origin, edge)) {
+            return false;
+          }
+          ptr->value().infallibleAppend(std::move(thisBackEdge));
+          totalPathsRecorded++;
+        }
+      }
+
+      MOZ_ASSERT(totalPathsRecorded <= totalMaxPathsToRecord);
+      if (totalPathsRecorded == totalMaxPathsToRecord) {
+        traversal.stop();
+      }
+
+      return true;
+    }
+  };
+
+  // The maximum number of paths to record for each node.
+  uint32_t maxNumPaths_;
+
+  // The root node we are starting the search from.
+  Node root_;
+
+  // The set of nodes we are searching for paths to.
+  NodeSet targets_;
+
+  // The resulting paths.
+  NodeToBackEdgeVectorMap paths_;
+
+  // Need to keep alive the traversal's back edges so we can walk them later
+  // when the traversal is over when recreating the shortest paths.
+  Traversal::NodeMap backEdges_;
+
+ private:
+  // Private methods.
+
+  ShortestPaths(uint32_t maxNumPaths, const Node& root, NodeSet&& targets)
+      : maxNumPaths_(maxNumPaths),
+        root_(root),
+        targets_(std::move(targets)),
+        paths_(targets_.count()),
+        backEdges_() {
+    MOZ_ASSERT(maxNumPaths_ > 0);
+    MOZ_ASSERT(root_);
+  }
+
+ public:
+  // Public methods.
+
+  ShortestPaths(ShortestPaths&& rhs)
+      : maxNumPaths_(rhs.maxNumPaths_),
+        root_(rhs.root_),
+        targets_(std::move(rhs.targets_)),
+        paths_(std::move(rhs.paths_)),
+        backEdges_(std::move(rhs.backEdges_)) {
+    MOZ_ASSERT(this != &rhs, "self-move is not allowed");
+  }
+
+  ShortestPaths& operator=(ShortestPaths&& rhs) {
+    this->~ShortestPaths();
+    new (this) ShortestPaths(std::move(rhs));
+    return *this;
+  }
+
+  ShortestPaths(const ShortestPaths&) = delete;
+  ShortestPaths& operator=(const ShortestPaths&) = delete;
+
+  /**
+   * Construct a new `JS::ubi::ShortestPaths`, finding up to `maxNumPaths`
+   * shortest retaining paths for each target node in `targets` starting from
+   * `root`.
+   *
+   * The resulting `ShortestPaths` instance must not outlive the
+   * `JS::ubi::Node` graph it was constructed from.
+   *
+   *   - For `JS::ubi::Node` graphs backed by the live heap graph, this means
+   *     that the `ShortestPaths`'s lifetime _must_ be contained within the
+   *     scope of the provided `AutoCheckCannotGC` reference because a GC will
+   *     invalidate the nodes.
+   *
+   *   - For `JS::ubi::Node` graphs backed by some other offline structure
+   *     provided by the embedder, the resulting `ShortestPaths`'s lifetime is
+   *     bounded by that offline structure's lifetime.
+   *
+   * Returns `mozilla::Nothing()` on OOM failure. It is the caller's
+   * responsibility to handle and report the OOM.
+   */
+  static mozilla::Maybe<ShortestPaths> Create(JSContext* cx,
+                                              AutoCheckCannotGC& noGC,
+                                              uint32_t maxNumPaths,
+                                              const Node& root,
+                                              NodeSet&& targets) {
+    MOZ_ASSERT(targets.count() > 0);
+    MOZ_ASSERT(maxNumPaths > 0);
+
+    ShortestPaths paths(maxNumPaths, root, std::move(targets));
+
+    Handler handler(paths);
+    Traversal traversal(cx, handler, noGC);
+    traversal.wantNames = true;
+    if (!traversal.addStart(root) || !traversal.traverse()) {
+      return mozilla::Nothing();
+    }
+
+    // Take ownership of the back edges we created while traversing the
+    // graph so that we can follow them from `paths_` and don't
+    // use-after-free.
+    paths.backEdges_ = std::move(traversal.visited);
+
+    return mozilla::Some(std::move(paths));
+  }
+
+  /**
+   * Get an iterator over each target node we searched for retaining paths
+   * for. The returned iterator must not outlive the `ShortestPaths`
+   * instance.
+   */
+  NodeSet::Iterator targetIter() const { return targets_.iter(); }
+
+  /**
+   * Invoke the provided functor/lambda/callable once for each retaining path
+   * discovered for `target`. The `func` is passed a single `JS::ubi::Path&`
+   * argument, which contains each edge along the path ordered starting from
+   * the root and ending at the target, and must not outlive the scope of the
+   * call.
+   *
+   * Note that it is possible that we did not find any paths from the root to
+   * the given target, in which case `func` will not be invoked.
+   */
+  template <class Func>
+  [[nodiscard]] bool forEachPath(const Node& target, Func func) {
+    MOZ_ASSERT(targets_.has(target));
+
+    auto ptr = paths_.lookup(target);
+
+    // We didn't find any paths to this target, so nothing to do here.
+    if (!ptr) {
+      return true;
+    }
+
+    MOZ_ASSERT(ptr->value().length() <= maxNumPaths_);
+
+    Path path;
+    for (const auto& backEdge : ptr->value()) {
+      path.clear();
+
+      if (!path.append(backEdge.get())) {
+        return false;
+      }
+
+      Node here = backEdge->predecessor();
+      MOZ_ASSERT(here);
+
+      while (here != root_) {
+        auto p = backEdges_.lookup(here);
+        MOZ_ASSERT(p);
+        if (!path.append(&p->value())) {
+          return false;
+        }
+        here = p->value().predecessor();
+        MOZ_ASSERT(here);
+      }
+
+      path.reverse();
+
+      if (!func(path)) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+};
+
+#ifdef DEBUG
+// A helper function to dump the first `maxNumPaths` shortest retaining paths to
+// `node` from the GC roots. Useful when GC things you expect to have been
+// reclaimed by the collector haven't been!
+//
+// Usage:
+//
+//     JSObject* foo = ...;
+//     JS::ubi::dumpPaths(rt, JS::ubi::Node(foo));
+JS_PUBLIC_API void dumpPaths(JSRuntime* rt, Node node,
+                             uint32_t maxNumPaths = 10);
+#endif
+
+}  // namespace ubi
+}  // namespace JS
+
+#endif  // js_UbiNodeShortestPaths_h