From 451678b19b5e481f04363c0d76e2ce39fafc8e15 Mon Sep 17 00:00:00 2001
From: Alastair Robertson <ajor@meta.com>
Date: Thu, 2 Nov 2023 06:09:07 -0700
Subject: [PATCH] TypeGraph: Create MutationTracker helper class

---
 oi/type_graph/NodeTracker.h | 43 +++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/oi/type_graph/NodeTracker.h b/oi/type_graph/NodeTracker.h
index baf7bdf..fdb79cd 100644
--- a/oi/type_graph/NodeTracker.h
+++ b/oi/type_graph/NodeTracker.h
@@ -108,4 +108,47 @@ class NodeTrackerHolder {
   NodeTracker& tracker_;
 };
 
+/*
+ * MutationTracker
+ *
+ * Helper class for mutators. Efficiently tracks visited and replaces nodes.
+ */
+class MutationTracker {
+ public:
+  MutationTracker(size_t size) : visited_(size) {
+  }
+
+  /*
+   * get
+   *
+   * Returns a type pointer if the given node has been visited or replaced.
+   * Returns nullptr if this node has not yet been seen.
+   */
+  Type* get(const Type& oldType) {
+    auto id = oldType.id();
+    if (id < 0)
+      return nullptr;
+    if (visited_.size() <= static_cast<size_t>(id))
+      visited_.resize(id + 1);
+    return visited_[id];
+  }
+
+  /*
+   * set
+   *
+   * Sets newType as the replacement node for oldType.
+   */
+  void set(const Type& oldType, Type& newType) {
+    auto id = oldType.id();
+    if (id < 0)
+      return;
+    if (visited_.size() <= static_cast<size_t>(id))
+      visited_.resize(id + 1);
+    visited_[id] = &newType;
+  }
+
+ private:
+  std::vector<Type*> visited_;
+};
+
 }  // namespace oi::detail::type_graph