TypeGraph: Introduce NodeTracker for efficient cycle detection

Added to Flattener and TypeIdentifier passes for now as a
proof-of-concept. Other passes can come later.

oi/type_graph/Flattener.cpp

@@ -22,26 +22,19 @@ namespace type_graph {
 
 Pass Flattener::createPass() {
   auto fn = [](TypeGraph& typeGraph) {
-    Flattener flattener;
-    flattener.flatten(typeGraph.rootTypes());
-    // TODO should flatten just operate on a single type and we do the looping
-    // here?
+    Flattener flattener{typeGraph.resetTracker()};
+    for (auto& type : typeGraph.rootTypes()) {
+      flattener.accept(type);
+    }
   };
 
   return Pass("Flattener", fn);
 }
 
-void Flattener::flatten(std::vector<std::reference_wrapper<Type>>& types) {
-  for (auto& type : types) {
-    accept(type);
-  }
-}
-
 void Flattener::accept(Type& type) {
-  if (visited_.count(&type) != 0)
+  if (tracker_.visit(type))
     return;
 
-  visited_.insert(&type);
   type.accept(*this);
 }
 

oi/type_graph/Flattener.h

@@ -16,9 +16,9 @@
 #pragma once
 
 #include <string>
-#include <unordered_set>
 #include <vector>
 
+#include "NodeTracker.h"
 #include "PassManager.h"
 #include "Types.h"
 #include "Visitor.h"
@@ -28,14 +28,15 @@ namespace type_graph {
 /*
  * Flattener
  *
- * Flattens classes by removing parents and adding their members directly into
- * derived classes.
+ * Flattens classes by removing parents and adding their attributes directly
+ * into derived classes.
  */
 class Flattener : public RecursiveVisitor {
  public:
   static Pass createPass();
 
-  void flatten(std::vector<std::reference_wrapper<Type>>& types);
+  Flattener(NodeTracker& tracker) : tracker_(tracker) {
+  }
 
   using RecursiveVisitor::accept;
 
@@ -46,7 +47,7 @@ class Flattener : public RecursiveVisitor {
   static const inline std::string ParentPrefix = "__oi_parent";
 
  private:
-  std::unordered_set<Type*> visited_;
+  NodeTracker& tracker_;
   std::vector<Member> flattened_members_;
   std::vector<uint64_t> offset_stack_;
 };

oi/type_graph/NodeTracker.h

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <vector>
+
+#include "Types.h"
+
+namespace type_graph {
+
+/*
+ * NodeTracker
+ *
+ * Helper class for visitors. Efficiently tracks whether or not a graph node has
+ * been seen before, to avoid infinite looping on cycles.
+ */
+class NodeTracker {
+ public:
+  NodeTracker() = default;
+  NodeTracker(size_t size) : visited_(size) {
+  }
+
+  /*
+   * visit
+   *
+   * Marks a given node as visited.
+   * Returns true if this node has already been visited, false otherwise.
+   */
+  bool visit(const Type& type) {
+    auto id = type.id();
+    if (id < 0)
+      return false;
+    if (visited_.size() <= static_cast<size_t>(id))
+      visited_.resize(id + 1);
+    bool result = visited_[id];
+    visited_[id] = true;
+    return result;
+  }
+
+  /*
+   * reset
+   *
+   * Clears the contents of this NodeTracker and marks every node as unvisited.
+   */
+  void reset() {
+    std::fill(visited_.begin(), visited_.end(), false);
+  }
+
+  void resize(size_t size) {
+    visited_.resize(size);
+  }
+
+ private:
+  std::vector<bool> visited_;
+};
+
+}  // namespace type_graph

oi/type_graph/TypeGraph.cpp

@@ -17,6 +17,12 @@
 
 namespace type_graph {
 
+NodeTracker& TypeGraph::resetTracker() noexcept {
+  tracker_.reset();
+  tracker_.resize(size());
+  return tracker_;
+}
+
 template <>
 Primitive& TypeGraph::makeType<Primitive>(Primitive::Kind kind) {
   switch (kind) {

oi/type_graph/TypeGraph.h

@@ -19,6 +19,7 @@
 #include <memory>
 #include <vector>
 
+#include "NodeTracker.h"
 #include "Types.h"
 
 namespace type_graph {
@@ -47,6 +48,8 @@ class TypeGraph {
     rootTypes_.push_back(type);
   }
 
+  NodeTracker& resetTracker() noexcept;
+
   // Override of the generic makeType function that returns singleton Primitive
   // objects
   template <typename T>
@@ -83,6 +86,7 @@ class TypeGraph {
   std::vector<std::reference_wrapper<Type>> rootTypes_;
   // Store all type objects in vectors for ownership. Order is not significant.
   std::vector<std::unique_ptr<Type>> types_;
+  NodeTracker tracker_;
   NodeId next_id_ = 0;
 };
 

oi/type_graph/TypeIdentifier.cpp

@@ -23,7 +23,8 @@ namespace type_graph {
 Pass TypeIdentifier::createPass(
     const std::vector<ContainerInfo>& passThroughTypes) {
   auto fn = [&passThroughTypes](TypeGraph& typeGraph) {
-    TypeIdentifier typeId{typeGraph, passThroughTypes};
+    TypeIdentifier typeId{typeGraph.resetTracker(), typeGraph,
+                          passThroughTypes};
     for (auto& type : typeGraph.rootTypes()) {
       typeId.accept(type);
     }
@@ -48,10 +49,9 @@ bool TypeIdentifier::isAllocator(Type& t) {
 }
 
 void TypeIdentifier::accept(Type& type) {
-  if (visited_.count(&type) != 0)
+  if (tracker_.visit(type))
     return;
 
-  visited_.insert(&type);
   type.accept(*this);
 }
 

oi/type_graph/TypeIdentifier.h

@@ -15,10 +15,9 @@
  */
 #pragma once
 
-#include <array>
-#include <unordered_set>
 #include <vector>
 
+#include "NodeTracker.h"
 #include "PassManager.h"
 #include "Types.h"
 #include "Visitor.h"
@@ -38,9 +37,12 @@ class TypeIdentifier : public RecursiveVisitor {
   static Pass createPass(const std::vector<ContainerInfo>& passThroughTypes);
   static bool isAllocator(Type& t);
 
-  TypeIdentifier(TypeGraph& typeGraph,
+  TypeIdentifier(NodeTracker& tracker,
+                 TypeGraph& typeGraph,
                  const std::vector<ContainerInfo>& passThroughTypes)
-      : typeGraph_(typeGraph), passThroughTypes_(passThroughTypes) {
+      : tracker_(tracker),
+        typeGraph_(typeGraph),
+        passThroughTypes_(passThroughTypes) {
   }
 
   using RecursiveVisitor::accept;
@@ -49,7 +51,7 @@ class TypeIdentifier : public RecursiveVisitor {
   void visit(Container& c) override;
 
  private:
-  std::unordered_set<Type*> visited_;
+  NodeTracker& tracker_;
   TypeGraph& typeGraph_;
   const std::vector<ContainerInfo>& passThroughTypes_;
 };

oi/type_graph/Types.h

@@ -24,6 +24,9 @@
  * to remain valid (i.e. don't store nodes directly in vectors). It is
  * recommended to use the TypeGraph class when building a complete type graph as
  * this will the memory allocations safely.
+ *
+ * All non-leaf nodes have IDs for efficient cycle detection and to assist
+ * debugging.
  */
 
 #include <cstddef>
@@ -49,6 +52,7 @@ struct ContainerInfo;
 
 namespace type_graph {
 
+// Must be signed. "-1" represents "leaf node"
 using NodeId = int32_t;
 
 enum class Qualifier {
@@ -65,7 +69,11 @@ class ConstVisitor;
 
 // TODO delete copy and move ctors
 
-// TODO type qualifiers are needed for some stuff?
+/*
+ * Type
+ *
+ * Abstract type base class. Attributes common to all types belong here.
+ */
 class Type {
  public:
   virtual ~Type() = default;
@@ -76,8 +84,14 @@ class Type {
   virtual std::string name() const = 0;
   virtual size_t size() const = 0;
   virtual uint64_t align() const = 0;
+  virtual NodeId id() const = 0;
 };
 
+/*
+ * Member
+ *
+ * A class' member variable.
+ */
 class Member {
  public:
   Member(Type& type,
@@ -146,13 +160,18 @@ class TemplateParam {
   }
 
  private:
-  Type* type_ = nullptr;  // Note: type is not always set
+  Type* type_ = nullptr;  // Note: type is not set when this param holds a value
 
  public:
   QualifierSet qualifiers;
   std::optional<std::string> value;
 };
 
+/*
+ * Class
+ *
+ * A class, struct or union.
+ */
 class Class : public Type {
  public:
   enum class Kind {
@@ -205,6 +224,10 @@ class Class : public Type {
     return align_;
   }
 
+  virtual NodeId id() const override {
+    return id_;
+  }
+
   void setAlign(uint64_t alignment) {
     align_ = alignment;
   }
@@ -227,10 +250,6 @@ class Class : public Type {
 
   bool isDynamic() const;
 
-  NodeId id() const {
-    return id_;
-  }
-
   std::vector<TemplateParam> templateParams;
   std::vector<Parent> parents;  // Sorted by offset
   std::vector<Member> members;  // Sorted by offset
@@ -280,7 +299,7 @@ class Container : public Type {
     return 8;  // TODO not needed for containers?
   }
 
-  NodeId id() const {
+  virtual NodeId id() const override {
     return id_;
   }
 
@@ -313,6 +332,10 @@ class Enum : public Type {
     return size();
   }
 
+  virtual NodeId id() const override {
+    return -1;
+  }
+
  private:
   std::string name_;
   size_t size_;
@@ -338,6 +361,10 @@ class Array : public Type {
     return elementType_.size();
   }
 
+  virtual NodeId id() const override {
+    return id_;
+  }
+
   Type& elementType() const {
     return elementType_;
   }
@@ -346,10 +373,6 @@ class Array : public Type {
     return len_;
   }
 
-  NodeId id() const {
-    return id_;
-  }
-
  private:
   Type& elementType_;
   size_t len_;
@@ -388,6 +411,9 @@ class Primitive : public Type {
   virtual uint64_t align() const override {
     return size();
   }
+  virtual NodeId id() const override {
+    return -1;
+  }
 
  private:
   Kind kind_;
@@ -417,12 +443,12 @@ class Typedef : public Type {
     return underlyingType_.align();
   }
 
-  Type& underlyingType() const {
-    return underlyingType_;
+  virtual NodeId id() const override {
+    return id_;
   }
 
-  NodeId id() const {
-    return id_;
+  Type& underlyingType() const {
+    return underlyingType_;
   }
 
  private:
@@ -451,12 +477,12 @@ class Pointer : public Type {
     return size();
   }
 
-  Type& pointeeType() const {
-    return pointeeType_;
+  virtual NodeId id() const override {
+    return id_;
   }
 
-  NodeId id() const {
-    return id_;
+  Type& pointeeType() const {
+    return pointeeType_;
   }
 
  private:
@@ -464,6 +490,11 @@ class Pointer : public Type {
   NodeId id_ = -1;
 };
 
+/*
+ * Dummy
+ *
+ * A type that just has a given size and alignment.
+ */
 class Dummy : public Type {
  public:
   explicit Dummy(size_t size, uint64_t align) : size_(size), align_(align) {
@@ -484,11 +515,21 @@ class Dummy : public Type {
     return align_;
   }
 
+  virtual NodeId id() const override {
+    return -1;
+  }
+
  private:
   size_t size_;
   uint64_t align_;
 };
 
+/*
+ * DummyAllocator
+ *
+ * A type with a given size and alignment, which also satisfies the C++
+ * allocator completeness requirements for a given type.
+ */
 class DummyAllocator : public Type {
  public:
   explicit DummyAllocator(Type& type, size_t size, uint64_t align)
@@ -510,6 +551,10 @@ class DummyAllocator : public Type {
     return align_;
   }
 
+  virtual NodeId id() const override {
+    return -1;
+  }
+
   Type& allocType() const {
     return type_;
   }

oi/type_graph/Visitor.h

@@ -15,6 +15,15 @@
  */
 #pragma once
 
+/*
+ * Visitors are used to walk over graphs of Type nodes.
+ *
+ * To implement a new visitor:
+ * 1. Inherit from one of the base visitor classes in this file
+ * 2. Override the `visit()` functions for the types your visitor must deal with
+ * 3. Implement `accept(Type&)` to perform double-dispatch and cycle detection
+ */
+
 #include "Types.h"
 
 namespace type_graph {

test/CMakeLists.txt

@@ -56,6 +56,7 @@ add_executable(test_type_graph
   test_drgn_parser.cpp
   test_flattener.cpp
   test_name_gen.cpp
+  test_node_tracker.cpp
   test_remove_ignored.cpp
   test_remove_top_level_pointer.cpp
   test_topo_sorter.cpp

test/test_node_tracker.cpp

@@ -0,0 +1,84 @@
+#include <gtest/gtest.h>
+
+#include "oi/type_graph/NodeTracker.h"
+
+using namespace type_graph;
+
+TEST(NodeTrackerTest, LeafNodes) {
+  Primitive myint32{Primitive::Kind::Int32};
+  Primitive myint64{Primitive::Kind::Int64};
+
+  NodeTracker tracker;
+
+  // First visit
+  EXPECT_FALSE(tracker.visit(myint32));
+  EXPECT_FALSE(tracker.visit(myint64));
+
+  // Second visit
+  EXPECT_FALSE(tracker.visit(myint32));
+  EXPECT_FALSE(tracker.visit(myint64));
+}
+
+TEST(NodeTrackerTest, Basic) {
+  Class myclass{0, Class::Kind::Class, "myclass", 0};
+  Array myarray{1, myclass, 3};
+
+  NodeTracker tracker;
+
+  // First visit
+  EXPECT_FALSE(tracker.visit(myarray));
+  EXPECT_FALSE(tracker.visit(myclass));
+
+  // Second visit
+  EXPECT_TRUE(tracker.visit(myarray));
+  EXPECT_TRUE(tracker.visit(myclass));
+
+  // Third visit
+  EXPECT_TRUE(tracker.visit(myarray));
+  EXPECT_TRUE(tracker.visit(myclass));
+}
+
+TEST(NodeTrackerTest, Clear) {
+  Class myclass{0, Class::Kind::Class, "myclass", 0};
+  Array myarray{1, myclass, 3};
+
+  NodeTracker tracker;
+
+  // First visit
+  EXPECT_FALSE(tracker.visit(myarray));
+  EXPECT_FALSE(tracker.visit(myclass));
+
+  // Second visit
+  EXPECT_TRUE(tracker.visit(myarray));
+  EXPECT_TRUE(tracker.visit(myclass));
+
+  // Reset
+  tracker.reset();
+
+  // First visit
+  EXPECT_FALSE(tracker.visit(myarray));
+  EXPECT_FALSE(tracker.visit(myclass));
+
+  // Second visit
+  EXPECT_TRUE(tracker.visit(myarray));
+  EXPECT_TRUE(tracker.visit(myclass));
+}
+
+TEST(NodeTrackerTest, LargeIds) {
+  Class myclass1{100, Class::Kind::Class, "myclass1", 0};
+  Class myclass2{100000, Class::Kind::Class, "myclass2", 0};
+
+  NodeTracker tracker;
+
+  // First visit
+  EXPECT_FALSE(tracker.visit(myclass1));
+  EXPECT_FALSE(tracker.visit(myclass2));
+
+  // Second visit
+  EXPECT_TRUE(tracker.visit(myclass1));
+  EXPECT_TRUE(tracker.visit(myclass2));
+
+  // Third visit
+  EXPECT_TRUE(tracker.visit(myclass1));
+  EXPECT_TRUE(tracker.visit(myclass2));
+}