Make runtime shadow node references updatable from native (#44772)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/44772

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.

Reviewed By: sammy-SC

Differential Revision: D57860869

fbshipit-source-id: 1703f0cd0183e2760436920a122857e17fda8dbb

packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp

@@ -309,6 +309,30 @@ 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_ != nullptr) {
+    runtimeShadowNodeReference_->shadowNode = destinationShadowNode;
+  }
+}
+
+void ShadowNode::transferRuntimeShadowNodeReference(
+    const Shared& destinationShadowNode,
+    const ShadowNodeFragment& fragment) const {
+  if (fragment.runtimeShadowNodeReference &&
+      ReactNativeFeatureFlags::useRuntimeShadowNodeReferenceUpdate()) {
+    transferRuntimeShadowNodeReference(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,27 @@ 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;
+
+  /*
+   * Transfer the runtime reference based on the fragment instructions.
+   */
+  void transferRuntimeShadowNodeReference(
+      const Shared& destinationShadowNode,
+      const ShadowNodeFragment& fragment) const;
+
 #pragma mark - DebugStringConvertible
 
 #if RN_DEBUG_STRING_CONVERTIBLE
@@ -245,10 +267,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/core/ShadowNodeFragment.h

@@ -27,6 +27,7 @@ struct ShadowNodeFragment {
   const ShadowNode::SharedListOfShared& children = childrenPlaceholder();
   const State::Shared& state = statePlaceholder();
   const ShadowNodeTraits traits = {};
+  const bool runtimeShadowNodeReference = true;
 
   /*
    * Placeholders.

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