From d988b6eee2490311a88874bfe17b7611f793b0fb Mon Sep 17 00:00:00 2001
From: Segev Finer <segev@swimm.io>
Date: Tue, 26 Mar 2024 01:46:29 +0200
Subject: [PATCH] Fix for the NodeClass is not a constructor issue

---
 src/node.cc | 14 ++++++++------
 src/tree.cc | 12 +++++++++---
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/node.cc b/src/node.cc
index 7dba9d80..a169eead 100644
--- a/src/node.cc
+++ b/src/node.cc
@@ -61,7 +61,10 @@ Napi::Value GetMarshalNodes(const Napi::CallbackInfo &info,
   for (unsigned i = 0; i < node_count; i++) {
     TSNode node = nodes[i];
     const auto &cache_entry = tree->cached_nodes_.find(node.id);
-    if (cache_entry == tree->cached_nodes_.end()) {
+    Napi::Value value;
+    if (cache_entry != tree->cached_nodes_.end() && (value = cache_entry->second->node.Value(), !value.IsEmpty())) {
+      result[i] = value;
+    } else {
       MarshalNodeId(node.id, p);
       p += 2;
       *(p++) = node.context[0];
@@ -73,8 +76,6 @@ Napi::Value GetMarshalNodes(const Napi::CallbackInfo &info,
       } else {
         result[i] = env.Null();
       }
-    } else {
-      result[i] = cache_entry->second->node.Value();
     }
   }
   return result;
@@ -84,7 +85,10 @@ Napi::Value GetMarshalNode(const Napi::CallbackInfo &info, const Tree *tree, TSN
   Env env = info.Env();
   auto* data = env.GetInstanceData<AddonData>();
   const auto &cache_entry = tree->cached_nodes_.find(node.id);
-  if (cache_entry == tree->cached_nodes_.end()) {
+  Napi::Value value;
+  if (cache_entry != tree->cached_nodes_.end() && (value = cache_entry->second->node.Value(), !value.IsEmpty())) {
+    return value;
+  } else {
     setup_transfer_buffer(env, 1);
     uint32_t *p = data->transfer_buffer;
     MarshalNodeId(node.id, p);
@@ -96,8 +100,6 @@ Napi::Value GetMarshalNode(const Napi::CallbackInfo &info, const Tree *tree, TSN
     if (node.id != nullptr) {
       return Number::New(env, ts_node_symbol(node));
     }
-  } else {
-    return cache_entry->second->node.Value();
   }
   return env.Null();
 }
diff --git a/src/tree.cc b/src/tree.cc
index edf217c9..f0433f91 100644
--- a/src/tree.cc
+++ b/src/tree.cc
@@ -102,6 +102,9 @@ Napi::Value Tree::Edit(const Napi::CallbackInfo &info) {
 
   for (auto &entry : cached_nodes_) {
     Object js_node = entry.second->node.Value();
+    if (js_node.IsEmpty()) {
+      continue;
+    }
     TSNode node;
     node.id = entry.first;
     for (unsigned i = 0; i < 4; i++) {
@@ -256,12 +259,15 @@ void CacheNodeForTree(Tree *tree, Napi::Env _env, Object js_node) {
   };
   const void *key = UnmarshalNodeId(key_parts);
 
+  // A Node could have been garbage collected but the finalizer has not yet run to remove its cache entry
+  if (tree->cached_nodes_.count(key)) {
+    return;
+  }
+
   auto *cache_entry = new Tree::NodeCacheEntry{tree, key, {}};
-  cache_entry->node.Reset(js_node, 0);
+  cache_entry->node = Napi::Weak(js_node);
   js_node.AddFinalizer(&FinalizeNode, cache_entry);
 
-  assert(!tree->cached_nodes_.count(key));
-
   tree->cached_nodes_[key] = cache_entry;
 }