Make runtime shadow node references updatable from native

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_;
 }

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

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;
 }