Go back to using ref callback instead of useImperativeHandle

We shouldn't create an object if the underlying node is null, as this
handle is meant to emulate the node but can't proxy methods if none
exist. Since useImperativeHandle isn't reset when the underlying node
changes, we go back to using ref callbacks. Rather than mutating the
underling node, we create a new object that wraps the relevant fields.

Author: necolas

packages/react-strict-dom/src/native/modules/useStrictDOMElement.js

@@ -7,9 +7,12 @@
  * @flow strict-local
  */
 
+import type { CallbackRef } from '../../types/react';
+
 import * as React from 'react';
 
 import { errorMsg } from '../../shared/logUtils';
+import { useElementCallback } from '../../shared/useElementCallback';
 import { useViewportScale } from './ContextViewportScale';
 
 type Options = {
@@ -25,36 +28,36 @@ function errorUnimplemented(name: string) {
   }
 }
 
-function proxy(nativeRef: React.RefObject<Node>, name: string) {
+function proxy(node: Node, name: string) {
   return (...args: Array<mixed>) => {
-    const node = nativeRef.current;
     if (node?.[name]) {
       return node[name](...args);
     }
     errorUnimplemented(name);
   };
 }
 
-export function useStrictDOMElement<T>(
-  ref: React.RefSetter<Node>,
-  { tagName }: Options
-): React.RefObject<T | null> {
-  const nativeRef = React.useRef<T | null>(null);
-  const { scale: viewportScale } = useViewportScale();
+const memoizedStrictRefs: WeakMap<Node, mixed> = new WeakMap();
 
-  React.useImperativeHandle(ref, () => {
-    return {
+function getOrCreateStrictRef(
+  node: Node,
+  tagName: string,
+  viewportScale: number
+) {
+  const ref = memoizedStrictRefs.get(node);
+  if (ref != null) {
+    return ref;
+  } else {
+    const strictRef = {
       get __nativeTag() {
-        const node = nativeRef.current as Node;
         return node?.__nativeTag;
       },
-      addEventListener: proxy(nativeRef, 'addEventListener'),
-      animate: proxy(nativeRef, 'animate'),
-      blur: proxy(nativeRef, 'blur'),
-      click: proxy(nativeRef, 'click'),
+      addEventListener: proxy(node, 'addEventListener'),
+      animate: proxy(node, 'animate'),
+      blur: proxy(node, 'blur'),
+      click: proxy(node, 'click'),
       get complete() {
         if (tagName === 'img') {
-          const node = nativeRef.current as Node;
           if (node?.complete == null) {
             // Assume images are never pre-loaded in React Native
             return false;
@@ -63,12 +66,11 @@ export function useStrictDOMElement<T>(
           }
         }
       },
-      contains: proxy(nativeRef, 'contains'),
-      dispatchEvent: proxy(nativeRef, 'dispatchEvent'),
-      focus: proxy(nativeRef, 'focus'),
-      getAttribute: proxy(nativeRef, 'getAttribute'),
+      contains: proxy(node, 'contains'),
+      dispatchEvent: proxy(node, 'dispatchEvent'),
+      focus: proxy(node, 'focus'),
+      getAttribute: proxy(node, 'getAttribute'),
       getBoundingClientRect() {
-        const node = nativeRef.current as Node;
         const getBoundingClientRect =
           node?.getBoundingClientRect ?? node?.unstable_getBoundingClientRect;
         if (getBoundingClientRect) {
@@ -85,21 +87,59 @@ export function useStrictDOMElement<T>(
         }
         return errorUnimplemented('getBoundingClientRect');
       },
-      getRootNode: proxy(nativeRef, 'getRootNode'),
-      hasPointerCapture: proxy(nativeRef, 'hasPointerCapture'),
+      getRootNode: proxy(node, 'getRootNode'),
+      hasPointerCapture: proxy(node, 'hasPointerCapture'),
       nodeName: tagName.toUpperCase(),
-      releasePointerCapture: proxy(nativeRef, 'releasePointerCapture'),
-      removeEventListener: proxy(nativeRef, 'removeEventListener'),
-      scroll: proxy(nativeRef, 'scroll'),
-      scrollBy: proxy(nativeRef, 'scrollBy'),
-      scrollIntoView: proxy(nativeRef, 'scrollIntoView'),
-      scrollTo: proxy(nativeRef, 'scrollTo'),
-      select: proxy(nativeRef, 'select'),
-      setSelectionRange: proxy(nativeRef, 'setSelectionRange'),
-      setPointerCapture: proxy(nativeRef, 'setPointerCapture'),
-      showPicker: proxy(nativeRef, 'showPicker')
+      releasePointerCapture: proxy(node, 'releasePointerCapture'),
+      removeEventListener: proxy(node, 'removeEventListener'),
+      scroll: proxy(node, 'scroll'),
+      scrollBy: proxy(node, 'scrollBy'),
+      scrollIntoView: proxy(node, 'scrollIntoView'),
+      scrollTo: proxy(node, 'scrollTo'),
+      select: proxy(node, 'select'),
+      setSelectionRange: proxy(node, 'setSelectionRange'),
+      setPointerCapture: proxy(node, 'setPointerCapture'),
+      showPicker: proxy(node, 'showPicker')
     };
-  }, [tagName, viewportScale]);
 
-  return nativeRef;
+    memoizedStrictRefs.set(node, strictRef);
+    return strictRef;
+  }
+}
+
+export function useStrictDOMElement<T>(
+  ref: React.RefSetter<Node>,
+  { tagName }: Options
+): CallbackRef<T> {
+  const { scale: viewportScale } = useViewportScale();
+
+  const elementCallback = useElementCallback(
+    React.useCallback(
+      // $FlowFixMe[unclear-type]
+      (node: any) => {
+        if (ref == null) {
+          return undefined;
+        } else {
+          const strictRef = getOrCreateStrictRef(node, tagName, viewportScale);
+          if (typeof ref === 'function') {
+            // $FlowFixMe[incompatible-type] - Flow does not understand ref cleanup.
+            const cleanup: void | (() => void) = ref(strictRef);
+            return typeof cleanup === 'function'
+              ? cleanup
+              : () => {
+                  ref(null);
+                };
+          } else {
+            ref.current = strictRef;
+            return () => {
+              ref.current = null;
+            };
+          }
+        }
+      },
+      [ref, tagName, viewportScale]
+    )
+  );
+
+  return elementCallback;
 }

packages/react-strict-dom/tests/__snapshots__/compat-test.native.js.snap-native

@@ -3,11 +3,7 @@
 exports[`<compat.native> "as" equals "image": as=image 1`] = `
 <Image
   accessibilityLabel="label"
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   srcSet="1x.img, 2x.img 2x"
   style={
     {
@@ -23,11 +19,7 @@ exports[`<compat.native> "as" equals "input": as=input 1`] = `
 <TextInput
   accessibilityLabel="label"
   placeholderTextColor="green"
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   secureTextEntry={true}
   style={
     {
@@ -42,11 +34,7 @@ exports[`<compat.native> "as" equals "text": as=text 1`] = `
 <Text
   accessibilityLabel="label"
   numberOfLines={3}
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   style={
     {
       "boxSizing": "content-box",
@@ -63,11 +51,7 @@ exports[`<compat.native> "as" equals "textarea": as=textarea 1`] = `
   accessibilityLabel="label"
   multiline={true}
   numberOfLines={3}
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   style={
     {
       "boxSizing": "content-box",
@@ -80,11 +64,7 @@ exports[`<compat.native> "as" equals "textarea": as=textarea 1`] = `
 exports[`<compat.native> "as" equals "view": as=view 1`] = `
 <View
   accessibilityLabel="label"
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   style={
     {
       "boxSizing": "content-box",
@@ -98,11 +78,7 @@ exports[`<compat.native> "as" equals "view": as=view 1`] = `
 exports[`<compat.native> default: default 1`] = `
 <Pressable
   accessibilityLabel="label"
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   style={
     {
       "boxSizing": "content-box",
@@ -116,11 +92,7 @@ exports[`<compat.native> default: default 1`] = `
 exports[`<compat.native> nested: nested 1`] = `
 <View
   accessibilityLabel="label"
-  ref={
-    {
-      "current": null,
-    }
-  }
+  ref={[Function]}
   style={
     {
       "boxSizing": "content-box",
@@ -129,11 +101,7 @@ exports[`<compat.native> nested: nested 1`] = `
   }
 >
   <Text
-    ref={
-      {
-        "current": null,
-      }
-    }
+    ref={[Function]}
     style={
       {
         "boxSizing": "content-box",