From f9dfcdf9fc1a7890aa5ff85c73ef701eee3b3761 Mon Sep 17 00:00:00 2001
From: Danny Su <dsc@meta.com>
Date: Mon, 12 Aug 2024 18:44:10 -0700
Subject: [PATCH] Allow IDTracker to pause assigning unique IDs

Summary:
Adding a state to IDTracker to pause assigning unique IDs for
`getNumberID()`. This is to allow any Edge generating code to not have
to worry about whether numeric values should be captured in the
subsequent diff.

Reviewed By: neildhar

Differential Revision: D61048450

fbshipit-source-id: 1b987d2336f3f2f8e8372cec86ed6705ba867a91
---
 include/hermes/VM/GCBase.h | 20 ++++++++++++++++++--
 lib/VM/GCBase.cpp          | 32 ++++++++++++++++++++++++++------
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/include/hermes/VM/GCBase.h b/include/hermes/VM/GCBase.h
index 3c4c1983e3e..1f7b9b8185f 100644
--- a/include/hermes/VM/GCBase.h
+++ b/include/hermes/VM/GCBase.h
@@ -592,6 +592,7 @@ class GCBase {
       Null,
       True,
       False,
+      Number,
       FirstNonReservedID,
     };
 
@@ -646,6 +647,15 @@ class GCBase {
     /// Return true if Object IDs have been tracked.
     bool hasTrackedObjectIDs();
 
+    /// Return true if Number IDs are being tracked.
+    bool isTrackingNumberIDs();
+
+    /// Start assigning unique IDs to numbers passed to \p getNumberID
+    void startTrackingNumberIDs();
+
+    /// Stop assigning unique IDs to numbers passed to \p getNumberID
+    void stopTrackingNumberIDs();
+
     /// Get the unique object id of the given object.
     /// If one does not yet exist, start tracking it.
     HeapSnapshot::NodeID getObjectID(CompressedPointer cell);
@@ -674,8 +684,11 @@ class GCBase {
     llvh::SmallVector<HeapSnapshot::NodeID, 1> &getExtraNativeIDs(
         HeapSnapshot::NodeID node);
 
-    /// Assign a unique ID to a literal number value that occurs in the heap.
-    /// Can be used to make fake nodes that will display their numeric value.
+    /// If \p isTrackingNumberIDs_ is true, then assign a unique ID to a literal
+    /// number value that occurs in the heap. Can be used to make fake nodes
+    /// that will display their numeric value. Otherwise if \p
+    /// isTrackingNumberIDs_ is false, then simply returns the reserved ID
+    /// representing Number.
     HeapSnapshot::NodeID getNumberID(double num);
 
     /// Get the object pointer for the given ID. This is the inverse of \c
@@ -774,6 +787,9 @@ class GCBase {
     /// Map of numeric values to IDs. Used to give numbers in the heap a unique
     /// node.
     llvh::DenseMap<double, HeapSnapshot::NodeID, DoubleComparator> numberIDMap_;
+
+    /// Whether to assign an unique ID to the number given to \p getNumberID
+    bool isTrackingNumberIDs_ = true;
   };
 
   enum class HeapKind { HadesGC, MallocGC };
diff --git a/lib/VM/GCBase.cpp b/lib/VM/GCBase.cpp
index 7cd8757ef52..8b8759c7386 100644
--- a/lib/VM/GCBase.cpp
+++ b/lib/VM/GCBase.cpp
@@ -1156,13 +1156,18 @@ GCBase::IDTracker::getExtraNativeIDs(HeapSnapshot::NodeID node) {
 
 HeapSnapshot::NodeID GCBase::IDTracker::getNumberID(double num) {
   std::lock_guard<Mutex> lk{mtx_};
-  auto &numberRef = numberIDMap_[num];
-  // If the entry didn't exist, the value was initialized to 0.
-  if (numberRef != 0) {
-    return numberRef;
+  if (isTrackingNumberIDs_) {
+    auto &numberRef = numberIDMap_[num];
+    // If the entry didn't exist, the value was initialized to 0.
+    if (numberRef != 0) {
+      return numberRef;
+    }
+    // Else, it is a number that hasn't been seen before.
+    return numberRef = nextNumberID();
+  } else {
+    return GCBase::IDTracker::reserved(
+        GCBase::IDTracker::ReservedObjectID::Number);
   }
-  // Else, it is a number that hasn't been seen before.
-  return numberRef = nextNumberID();
 }
 
 llvh::Optional<CompressedPointer> GCBase::IDTracker::getObjectForID(
@@ -1198,6 +1203,21 @@ bool GCBase::IDTracker::hasTrackedObjectIDs() {
   return !objectIDMap_.empty();
 }
 
+bool GCBase::IDTracker::isTrackingNumberIDs() {
+  std::lock_guard<Mutex> lk{mtx_};
+  return isTrackingNumberIDs_;
+}
+
+void GCBase::IDTracker::startTrackingNumberIDs() {
+  std::lock_guard<Mutex> lk{mtx_};
+  isTrackingNumberIDs_ = true;
+}
+
+void GCBase::IDTracker::stopTrackingNumberIDs() {
+  std::lock_guard<Mutex> lk{mtx_};
+  isTrackingNumberIDs_ = false;
+}
+
 HeapSnapshot::NodeID GCBase::IDTracker::getObjectID(CompressedPointer cell) {
   std::lock_guard<Mutex> lk{mtx_};
   auto iter = objectIDMap_.find(cell.getRaw());