[RN] Add support for document instance in React Native (facebook#32260)

## Summary

We're adding support for `Document` instances in React Native (as
`ReactNativeDocument` instances) in
facebook/react-native#49012 , which requires the
React Fabric renderer to handle its lifecycle.

This modifies the renderer to create those document instances and
associate them with the React root, and provides a new method for React
Native to access them given its containerTag / rootTag.

## How did you test this change?

Tested e2e in facebook/react-native#49012
manually syncing these changes.

Author: rubennorte

packages/react-native-renderer/src/ReactFabric.js

@@ -41,7 +41,11 @@ import {
 import {getPublicInstanceFromInternalInstanceHandle} from './ReactFiberConfigFabric';
 
 // Module provided by RN:
-import {ReactFiberErrorDialog} from 'react-native/Libraries/ReactPrivate/ReactNativePrivateInterface';
+import {
+  ReactFiberErrorDialog,
+  createPublicRootInstance,
+  type PublicRootInstance,
+} from 'react-native/Libraries/ReactPrivate/ReactNativePrivateInterface';
 import {disableLegacyMode} from 'shared/ReactFeatureFlags';
 
 if (typeof ReactFiberErrorDialog.showErrorDialog !== 'function') {
@@ -126,10 +130,16 @@ function render(
       onRecoverableError = options.onRecoverableError;
     }
 
+    const publicRootInstance = createPublicRootInstance(containerTag);
+    const rootInstance = {
+      publicInstance: publicRootInstance,
+      containerTag,
+    };
+
     // TODO (bvaughn): If we decide to keep the wrapper component,
     // We could create a wrapper for containerTag as well to reduce special casing.
     root = createContainer(
-      containerTag,
+      rootInstance,
       concurrentRoot ? ConcurrentRoot : LegacyRoot,
       null,
       false,
@@ -140,6 +150,7 @@ function render(
       onRecoverableError,
       null,
     );
+
     roots.set(containerTag, root);
   }
   updateContainer(element, root, null, callback);
@@ -157,6 +168,9 @@ function stopSurface(containerTag: number) {
   if (root) {
     // TODO: Is it safe to reset this now or should I wait since this unmount could be deferred?
     updateContainer(null, root, null, () => {
+      // Remove the reference to the public instance to prevent memory leaks.
+      root.containerInfo.publicInstance = null;
+
       roots.delete(containerTag);
     });
   }
@@ -170,6 +184,16 @@ function createPortal(
   return createPortalImpl(children, containerTag, null, key);
 }
 
+function getPublicInstanceFromRootTag(
+  rootTag: number,
+): PublicRootInstance | null {
+  const root = roots.get(rootTag);
+  if (root) {
+    return root.containerInfo.publicInstance;
+  }
+  return null;
+}
+
 setBatchingImplementation(batchedUpdatesImpl, discreteUpdates);
 
 const roots = new Map<number, FiberRoot>();
@@ -195,6 +219,8 @@ export {
   // instance handles we use to dispatch events. This provides a way to access
   // the public instances we created from them (potentially created lazily).
   getPublicInstanceFromInternalInstanceHandle,
+  // Returns the document instance for that root tag.
+  getPublicInstanceFromRootTag,
   // DEV-only:
   isChildPublicInstance,
 };

packages/react-native-renderer/src/ReactFiberConfigFabric.js

@@ -31,6 +31,7 @@ import {
   createPublicTextInstance,
   type PublicInstance as ReactNativePublicInstance,
   type PublicTextInstance,
+  type PublicRootInstance,
 } from 'react-native/Libraries/ReactPrivate/ReactNativePrivateInterface';
 
 const {
@@ -108,7 +109,10 @@ export type TextInstance = {
 };
 export type HydratableInstance = Instance | TextInstance;
 export type PublicInstance = ReactNativePublicInstance;
-export type Container = number;
+export type Container = {
+  containerTag: number,
+  publicInstance: PublicRootInstance | null,
+};
 export type ChildSet = Object | Array<Node>;
 export type HostContext = $ReadOnly<{
   isInAParentText: boolean,
@@ -180,7 +184,7 @@ export function createInstance(
   const node = createNode(
     tag, // reactTag
     viewConfig.uiViewClassName, // viewName
-    rootContainerInstance, // rootTag
+    rootContainerInstance.containerTag, // rootTag
     updatePayload, // props
     internalInstanceHandle, // internalInstanceHandle
   );
@@ -189,6 +193,7 @@ export function createInstance(
     tag,
     viewConfig,
     internalInstanceHandle,
+    rootContainerInstance.publicInstance,
   );
 
   return {
@@ -221,7 +226,7 @@ export function createTextInstance(
   const node = createNode(
     tag, // reactTag
     'RCTRawText', // viewName
-    rootContainerInstance, // rootTag
+    rootContainerInstance.containerTag, // rootTag
     {text: text}, // props
     internalInstanceHandle, // instance handle
   );
@@ -501,7 +506,7 @@ export function finalizeContainerChildren(
   newChildren: ChildSet,
 ): void {
   if (!enableFabricCompleteRootInCommitPhase) {
-    completeRoot(container, newChildren);
+    completeRoot(container.containerTag, newChildren);
   }
 }
 
@@ -511,7 +516,7 @@ export function replaceContainerChildren(
 ): void {
   // Noop - children will be replaced in finalizeContainerChildren
   if (enableFabricCompleteRootInCommitPhase) {
-    completeRoot(container, newChildren);
+    completeRoot(container.containerTag, newChildren);
   }
 }
 

packages/react-native-renderer/src/ReactFiberConfigNative.js

@@ -15,6 +15,7 @@ import {
   ReactNativeViewConfigRegistry,
   UIManager,
   deepFreezeAndThrowOnMutationInDev,
+  type PublicRootInstance,
 } from 'react-native/Libraries/ReactPrivate/ReactNativePrivateInterface';
 
 import {create, diff} from './ReactNativeAttributePayload';
@@ -54,7 +55,10 @@ const {get: getViewConfigForType} = ReactNativeViewConfigRegistry;
 
 export type Type = string;
 export type Props = Object;
-export type Container = number;
+export type Container = {
+  containerTag: number,
+  publicInstance: PublicRootInstance | null,
+};
 export type Instance = ReactNativeFiberHostComponent;
 export type TextInstance = number;
 export type HydratableInstance = Instance | TextInstance;
@@ -143,7 +147,7 @@ export function createInstance(
   UIManager.createView(
     tag, // reactTag
     viewConfig.uiViewClassName, // viewName
-    rootContainerInstance, // rootTag
+    rootContainerInstance.containerTag, // rootTag
     updatePayload, // props
   );
 
@@ -176,7 +180,7 @@ export function createTextInstance(
   UIManager.createView(
     tag, // reactTag
     'RCTRawText', // viewName
-    rootContainerInstance, // rootTag
+    rootContainerInstance.containerTag, // rootTag
     {text: text}, // props
   );
 
@@ -349,7 +353,7 @@ export function appendChildToContainer(
 ): void {
   const childTag = typeof child === 'number' ? child : child._nativeTag;
   UIManager.setChildren(
-    parentInstance, // containerTag
+    parentInstance.containerTag, // containerTag
     [childTag], // reactTags
   );
 }
@@ -479,7 +483,7 @@ export function removeChildFromContainer(
 ): void {
   recursivelyUncacheFiberNode(child);
   UIManager.manageChildren(
-    parentInstance, // containerID
+    parentInstance.containerTag, // containerID
     [], // moveFromIndices
     [], // moveToIndices
     [], // addChildReactTags

packages/react-native-renderer/src/ReactNativeRenderer.js

@@ -11,6 +11,7 @@ import type {ReactPortal, ReactNodeList} from 'shared/ReactTypes';
 import type {ElementRef, ElementType, MixedElement} from 'react';
 import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
 import type {RenderRootOptions} from './ReactNativeTypes';
+import type {Container} from 'react-reconciler/src/ReactFiberConfig';
 
 import './ReactNativeInjection';
 
@@ -143,10 +144,16 @@ function render(
       onRecoverableError = options.onRecoverableError;
     }
 
+    const rootInstance: Container = {
+      containerTag,
+      // $FlowExpectedError[incompatible-type] the legacy renderer does not use public root instances
+      publicInstance: null,
+    };
+
     // TODO (bvaughn): If we decide to keep the wrapper component,
     // We could create a wrapper for containerTag as well to reduce special casing.
     root = createContainer(
-      containerTag,
+      rootInstance,
       LegacyRoot,
       null,
       false,

packages/react-native-renderer/src/ReactNativeTypes.js

@@ -231,6 +231,7 @@ export opaque type Node = mixed;
 export opaque type InternalInstanceHandle = mixed;
 type PublicInstance = mixed;
 type PublicTextInstance = mixed;
+export opaque type PublicRootInstance = mixed;
 
 export type ReactFabricType = {
   findHostInstance_DEPRECATED<TElementType: ElementType>(

packages/react-native-renderer/src/__mocks__/react-native/Libraries/ReactPrivate/ReactNativePrivateInterface.js

@@ -9,6 +9,7 @@
 
 export opaque type PublicInstance = mixed;
 export opaque type PublicTextInstance = mixed;
+export opaque type PublicRootInstance = mixed;
 
 module.exports = {
   get BatchedBridge() {
@@ -59,4 +60,7 @@ module.exports = {
   get createPublicTextInstance() {
     return require('./createPublicTextInstance').default;
   },
+  get createPublicRootInstance() {
+    return require('./createPublicRootInstance').default;
+  },
 };

packages/react-native-renderer/src/__mocks__/react-native/Libraries/ReactPrivate/createPublicInstance.js

@@ -7,15 +7,20 @@
  * @flow strict
  */
 
-import type {PublicInstance} from './ReactNativePrivateInterface';
+import type {
+  PublicInstance,
+  PublicRootInstance,
+} from './ReactNativePrivateInterface';
 
 export default function createPublicInstance(
   tag: number,
   viewConfig: mixed,
   internalInstanceHandle: mixed,
+  rootPublicInstance: PublicRootInstance | null,
 ): PublicInstance {
   return {
     __nativeTag: tag,
     __internalInstanceHandle: internalInstanceHandle,
+    __rootPublicInstance: rootPublicInstance,
   };
 }

packages/react-native-renderer/src/__mocks__/react-native/Libraries/ReactPrivate/createPublicRootInstance.js

@@ -0,0 +1,16 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow strict
+ */
+
+import type {PublicRootInstance} from './ReactNativePrivateInterface';
+
+export default function createPublicRootInstance(
+  rootTag: number,
+): PublicRootInstance {
+  return null;
+}

scripts/flow/react-native-host-hooks.js

@@ -143,6 +143,7 @@ declare module 'react-native/Libraries/ReactPrivate/ReactNativePrivateInterface'
   };
   declare export opaque type PublicInstance;
   declare export opaque type PublicTextInstance;
+  declare export opaque type PublicRootInstance;
   declare export function getNodeFromPublicInstance(
     publicInstance: PublicInstance,
   ): Object;
@@ -153,7 +154,11 @@ declare module 'react-native/Libraries/ReactPrivate/ReactNativePrivateInterface'
     tag: number,
     viewConfig: __ViewConfig,
     internalInstanceHandle: mixed,
+    publicRootInstance: PublicRootInstance | null,
   ): PublicInstance;
+  declare export function createPublicRootInstance(
+    rootTag: number,
+  ): PublicRootInstance;
   declare export function createPublicTextInstance(
     internalInstanceHandle: mixed,
   ): PublicTextInstance;