From 5cc2ad518685570db98ef31ec75af335f2818321 Mon Sep 17 00:00:00 2001
From: Nick Lefever <lefever@meta.com>
Date: Fri, 31 May 2024 05:52:38 -0700
Subject: [PATCH] Make runtime shadow node references updatable from native

Summary:
Changelog: [Internal]

React native clones shadow nodes internally without providing the new instances to the React renderer (current fiber tree). To support updating the shadow node references held by the JS side, this diff wraps the returned shadow nodes and adds a link to the runtime reference on the shadow node instance.

This will allow for updating the shadow node references held within the JS runtime from the native side.

Differential Revision: D57860869
---
 .../react/renderer/core/ShadowNode.cpp        | 15 +++++++++++
 .../react/renderer/core/ShadowNode.h          | 27 +++++++++++++++++++
 .../react/renderer/uimanager/primitives.h     | 13 ++++++---
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
index d84c512643a5..846e3f2c3e48 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
+++ b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
@@ -309,6 +309,21 @@ bool ShadowNode::progressStateIfNecessary() {
   return false;
 }
 
+void ShadowNode::setRuntimeShadowNodeReference(
+    ShadowNodeWrapper* runtimeShadowNodeReference) const {
+  runtimeShadowNodeReference_ = runtimeShadowNodeReference;
+}
+
+void ShadowNode::transferRuntimeShadowNodeReference(
+    const Shared& destinationShadowNode) const {
+  destinationShadowNode->runtimeShadowNodeReference_ =
+      runtimeShadowNodeReference_;
+
+  if (runtimeShadowNodeReference_) {
+    runtimeShadowNodeReference_->shadowNode = destinationShadowNode;
+  }
+}
+
 const ShadowNodeFamily& ShadowNode::getFamily() const {
   return *family_;
 }
diff --git a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
index 644159952bb4..1333c9c7a662 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
+++ b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
@@ -26,6 +26,7 @@ namespace facebook::react {
 
 class ComponentDescriptor;
 struct ShadowNodeFragment;
+struct ShadowNodeWrapper;
 
 class ShadowNode : public Sealable,
                    public DebugStringConvertible,
@@ -188,6 +189,20 @@ class ShadowNode : public Sealable,
    */
   bool progressStateIfNecessary();
 
+  /*
+   * Bind the runtime reference to this `ShadowNode` with a raw pointer,
+   * allowing to update the reference to this `ShadowNode` when cloned.
+   */
+  void setRuntimeShadowNodeReference(
+      ShadowNodeWrapper* runtimeShadowNodeReference) const;
+
+  /*
+   * Transfer the runtime reference to this `ShadowNode` to a new instance,
+   * updating the reference to point to the new `ShadowNode` referencing it.
+   */
+  void transferRuntimeShadowNodeReference(
+      const Shared& destinationShadowNode) const;
+
 #pragma mark - DebugStringConvertible
 
 #if RN_DEBUG_STRING_CONVERTIBLE
@@ -245,10 +260,22 @@ class ShadowNode : public Sealable,
    * that class.
    */
   ShadowNodeTraits traits_;
+
+  /*
+   * Pointer to the runtime reference to this `ShadowNode`.
+   */
+  mutable ShadowNodeWrapper* runtimeShadowNodeReference_{};
 };
 
 static_assert(
     std::has_virtual_destructor<ShadowNode>::value,
     "ShadowNode must have a virtual destructor");
 
+struct ShadowNodeWrapper : public jsi::NativeState {
+  explicit ShadowNodeWrapper(ShadowNode::Shared shadowNode)
+      : shadowNode(std::move(shadowNode)) {}
+
+  ShadowNode::Shared shadowNode;
+};
+
 } // namespace facebook::react
diff --git a/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h b/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h
index 03552d2b140f..6c7e8b3dbafe 100644
--- a/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h
+++ b/packages/react-native/ReactCommon/react/renderer/uimanager/primitives.h
@@ -37,16 +37,21 @@ inline static ShadowNode::Shared shadowNodeFromValue(
     return nullptr;
   }
 
-  return value.getObject(runtime).getNativeState<ShadowNode>(runtime);
+  return value.getObject(runtime)
+      .getNativeState<ShadowNodeWrapper>(runtime)
+      ->shadowNode;
 }
 
 inline static jsi::Value valueFromShadowNode(
     jsi::Runtime& runtime,
     ShadowNode::Shared shadowNode) {
+  // Wrap the shadow node so that we can update JS references from native
+  auto wrappedShadowNode =
+      std::make_shared<ShadowNodeWrapper>(std::move(shadowNode));
+  wrappedShadowNode->shadowNode->setRuntimeShadowNodeReference(
+      &*wrappedShadowNode);
   jsi::Object obj(runtime);
-  // Need to const_cast since JSI only allows non-const pointees
-  obj.setNativeState(
-      runtime, std::const_pointer_cast<ShadowNode>(std::move(shadowNode)));
+  obj.setNativeState(runtime, std::move(wrappedShadowNode));
   return obj;
 }