1 files changed, 191 insertions, 0 deletions
diff --git a/js/public/UbiNodePostOrder.h b/js/public/UbiNodePostOrder.h
new file mode 100644
index 0000000000..10d8835be9
--- /dev/null
+++ b/js/public/UbiNodePostOrder.h
@@ -0,0 +1,191 @@
+/* -*- 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_UbiNodePostOrder_h
+#define js_UbiNodePostOrder_h
+
+#include "mozilla/Attributes.h"
+
+#include <utility>
+
+#include "js/UbiNode.h"
+#include "js/Vector.h"
+
+namespace js {
+class SystemAllocPolicy;
+}
+
+namespace JS {
+class AutoCheckCannotGC;
+
+namespace ubi {
+
+/**
+ * A post-order depth-first traversal of `ubi::Node` graphs.
+ *
+ * No GC may occur while an instance of `PostOrder` is live.
+ *
+ * The `NodeVisitor` type provided to `PostOrder::traverse` must have the
+ * following member:
+ *
+ *   bool operator()(Node& node)
+ *
+ *     The node visitor method. This method is called once for each `node`
+ *     reachable from the start set in post-order.
+ *
+ *     This visitor function should return true on success, or false if an error
+ *     occurs. A false return value terminates the traversal immediately, and
+ *     causes `PostOrder::traverse` to return false.
+ *
+ * The `EdgeVisitor` type provided to `PostOrder::traverse` must have the
+ * following member:
+ *
+ *   bool operator()(Node& origin, Edge& edge)
+ *
+ *     The edge visitor method. This method is called once for each outgoing
+ *     `edge` from `origin` that is reachable from the start set.
+ *
+ *     NB: UNLIKE NODES, THERE IS NO GUARANTEED ORDER IN WHICH EDGES AND THEIR
+ *     ORIGINS ARE VISITED!
+ *
+ *     This visitor function should return true on success, or false if an error
+ *     occurs. A false return value terminates the traversal immediately, and
+ *     causes `PostOrder::traverse` to return false.
+ */
+struct PostOrder {
+ private:
+  struct OriginAndEdges {
+    Node origin;
+    EdgeVector edges;
+
+    OriginAndEdges(const Node& node, EdgeVector&& edges)
+        : origin(node), edges(std::move(edges)) {}
+
+    OriginAndEdges(const OriginAndEdges& rhs) = delete;
+    OriginAndEdges& operator=(const OriginAndEdges& rhs) = delete;
+
+    OriginAndEdges(OriginAndEdges&& rhs)
+        : origin(rhs.origin), edges(std::move(rhs.edges)) {
+      MOZ_ASSERT(&rhs != this, "self-move disallowed");
+    }
+
+    OriginAndEdges& operator=(OriginAndEdges&& rhs) {
+      this->~OriginAndEdges();
+      new (this) OriginAndEdges(std::move(rhs));
+      return *this;
+    }
+  };
+
+  using Stack = js::Vector<OriginAndEdges, 256, js::SystemAllocPolicy>;
+  using Set = js::HashSet<Node, js::DefaultHasher<Node>, js::SystemAllocPolicy>;
+
+  JSContext* cx;
+  Set seen;
+  Stack stack;
+#ifdef DEBUG
+  bool traversed;
+#endif
+
+ private:
+  [[nodiscard]] bool fillEdgesFromRange(EdgeVector& edges,
+                                        js::UniquePtr<EdgeRange>& range) {
+    MOZ_ASSERT(range);
+    for (; !range->empty(); range->popFront()) {
+      if (!edges.append(std::move(range->front()))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  [[nodiscard]] bool pushForTraversing(const Node& node) {
+    EdgeVector edges;
+    auto range = node.edges(cx, /* wantNames */ false);
+    return range && fillEdgesFromRange(edges, range) &&
+           stack.append(OriginAndEdges(node, std::move(edges)));
+  }
+
+ public:
+  // Construct a post-order traversal object.
+  //
+  // The traversal asserts that no GC happens in its runtime during its
+  // lifetime via the `AutoCheckCannotGC&` parameter. We do nothing with it,
+  // other than require it to exist with a lifetime that encloses our own.
+  PostOrder(JSContext* cx, AutoCheckCannotGC&)
+      : cx(cx),
+        seen(),
+        stack()
+#ifdef DEBUG
+        ,
+        traversed(false)
+#endif
+  {
+  }
+
+  // Add `node` as a starting point for the traversal. You may add
+  // as many starting points as you like. Returns false on OOM.
+  [[nodiscard]] bool addStart(const Node& node) {
+    if (!seen.put(node)) {
+      return false;
+    }
+    return pushForTraversing(node);
+  }
+
+  // Traverse the graph in post-order, starting with the set of nodes passed
+  // to `addStart` and applying `onNode::operator()` for each node in the
+  // graph and `onEdge::operator()` for each edge in the graph, as described
+  // above.
+  //
+  // This should be called only once per instance of this class.
+  //
+  // Return false on OOM or error return from `onNode::operator()` or
+  // `onEdge::operator()`.
+  template <typename NodeVisitor, typename EdgeVisitor>
+  [[nodiscard]] bool traverse(NodeVisitor onNode, EdgeVisitor onEdge) {
+#ifdef DEBUG
+    MOZ_ASSERT(!traversed, "Can only traverse() once!");
+    traversed = true;
+#endif
+
+    while (!stack.empty()) {
+      auto& origin = stack.back().origin;
+      auto& edges = stack.back().edges;
+
+      if (edges.empty()) {
+        if (!onNode(origin)) {
+          return false;
+        }
+        stack.popBack();
+        continue;
+      }
+
+      Edge edge = std::move(edges.back());
+      edges.popBack();
+
+      if (!onEdge(origin, edge)) {
+        return false;
+      }
+
+      auto ptr = seen.lookupForAdd(edge.referent);
+      // We've already seen this node, don't follow its edges.
+      if (ptr) {
+        continue;
+      }
+
+      // Mark the referent as seen and follow its edges.
+      if (!seen.add(ptr, edge.referent) || !pushForTraversing(edge.referent)) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+};
+
+}  // namespace ubi
+}  // namespace JS
+
+#endif  // js_UbiNodePostOrder_h