Fix incorrect wrapping of refs in Image component (#41679)

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

This fixes a bug in the original implementation of image attached callbacks (still experimental). The problem was that we were unconditionally caching the ref passed to the underlying image component, which meant that whenever users passed new ref setters we wouldn't call them again.

This fixes that by forcing the creation of a new ref value whenever a new ref is passed to the image component.

Changelog: [internal]

Reviewed By: jehartzog

Differential Revision: D51618512

fbshipit-source-id: ac15160c528563c2131e8b3444dea4a6096f20bc

packages/react-native/Libraries/Image/Image.android.js

@@ -15,10 +15,11 @@ import type {AbstractImageAndroid, ImageAndroid} from './ImageTypes.flow';
 import flattenStyle from '../StyleSheet/flattenStyle';
 import StyleSheet from '../StyleSheet/StyleSheet';
 import TextAncestor from '../Text/TextAncestor';
+import useMergeRefs from '../Utilities/useMergeRefs';
 import ImageAnalyticsTagContext from './ImageAnalyticsTagContext';
 import {
   unstable_getImageComponentDecorator,
-  useWrapRefWithImageAttachedCallbacks,
+  useRefWithImageAttachedCallbacks,
 } from './ImageInjection';
 import {getImageSourcesFromImageProps} from './ImageSourceUtils';
 import {convertObjectFitToResizeMode} from './ImageUtils';
@@ -199,7 +200,15 @@ let BaseImage: AbstractImageAndroid = React.forwardRef(
     const resizeMode =
       objectFit || props.resizeMode || style?.resizeMode || 'cover';
 
-    const actualRef = useWrapRefWithImageAttachedCallbacks(forwardedRef);
+    const imageAttachedCallbacksRef = useRefWithImageAttachedCallbacks();
+
+    const actualRef =
+      useMergeRefs<React.ElementRef<AbstractImageAndroid> | null>(
+        // $FlowFixMe[incompatible-call]
+        forwardedRef,
+        // $FlowFixMe[incompatible-call]
+        imageAttachedCallbacksRef,
+      );
 
     return (
       <ImageAnalyticsTagContext.Consumer>

packages/react-native/Libraries/Image/Image.ios.js

@@ -15,10 +15,11 @@ import type {AbstractImageIOS, ImageIOS} from './ImageTypes.flow';
 import {createRootTag} from '../ReactNative/RootTag';
 import flattenStyle from '../StyleSheet/flattenStyle';
 import StyleSheet from '../StyleSheet/StyleSheet';
+import useMergeRefs from '../Utilities/useMergeRefs';
 import ImageAnalyticsTagContext from './ImageAnalyticsTagContext';
 import {
   unstable_getImageComponentDecorator,
-  useWrapRefWithImageAttachedCallbacks,
+  useRefWithImageAttachedCallbacks,
 } from './ImageInjection';
 import {getImageSourcesFromImageProps} from './ImageSourceUtils';
 import {convertObjectFitToResizeMode} from './ImageUtils';
@@ -161,7 +162,14 @@ let BaseImage: AbstractImageIOS = React.forwardRef((props, forwardedRef) => {
   };
   const accessibilityLabel = props['aria-label'] ?? props.accessibilityLabel;
 
-  const actualRef = useWrapRefWithImageAttachedCallbacks(forwardedRef);
+  const imageAttachedCallbacksRef = useRefWithImageAttachedCallbacks();
+
+  const actualRef = useMergeRefs<React.ElementRef<AbstractImageIOS> | null>(
+    // $FlowFixMe[incompatible-call]
+    forwardedRef,
+    // $FlowFixMe[incompatible-call]
+    imageAttachedCallbacksRef,
+  );
 
   return (
     <ImageAnalyticsTagContext.Consumer>

packages/react-native/Libraries/Image/ImageInjection.js

@@ -52,19 +52,10 @@ export function unstable_unregisterImageAttachedCallback(
   imageAttachedCallbacks.delete(callback);
 }
 
-type ProxyRef = (ImageInstance | null) => void;
-
-export function useWrapRefWithImageAttachedCallbacks(
-  forwardedRef?: React.Ref<ImageComponent>,
-): ProxyRef {
+export function useRefWithImageAttachedCallbacks(): React.RefSetter<ImageInstance> {
   const pendingCleanupCallbacks = useRef<Array<() => void>>([]);
-  const proxyRef = useRef<ProxyRef>(node => {
-    if (typeof forwardedRef === 'function') {
-      forwardedRef(node);
-    } else if (typeof forwardedRef === 'object' && forwardedRef != null) {
-      forwardedRef.current = node;
-    }
 
+  const ref = useRef((node: ImageInstance | null) => {
     if (node == null) {
       if (pendingCleanupCallbacks.current.length > 0) {
         pendingCleanupCallbacks.current.forEach(cb => cb());
@@ -80,5 +71,5 @@ export function useWrapRefWithImageAttachedCallbacks(
     }
   });
 
-  return proxyRef.current;
+  return ref.current;
 }

packages/react-native/Libraries/Image/__tests__/Image-test.js

@@ -45,6 +45,67 @@ describe('<Image />', () => {
     expect(instance).toMatchSnapshot();
   });
 
+  it('should invoke original ref callbacks correctly when using image attached callbacks', () => {
+    jest.dontMock('../Image');
+
+    let imageInstanceFromCallback = null;
+    let imageInstanceFromRef1 = null;
+    let imageInstanceFromRef2 = null;
+
+    const callback = jest.fn((instance: ElementRef<typeof Image>) => {
+      imageInstanceFromCallback = instance;
+
+      return () => {
+        imageInstanceFromCallback = null;
+      };
+    });
+
+    ImageInjection.unstable_registerImageAttachedCallback(callback);
+
+    expect(imageInstanceFromCallback).toBe(null);
+
+    let testRenderer;
+
+    const ref1 = jest.fn(instance => {
+      imageInstanceFromRef1 = instance;
+    });
+
+    act(() => {
+      testRenderer = create(<Image source={{uri: 'foo-bar.jpg'}} ref={ref1} />);
+    });
+
+    expect(imageInstanceFromCallback).not.toBe(null);
+    expect(imageInstanceFromRef1).not.toBe(null);
+    expect(imageInstanceFromCallback).toBe(imageInstanceFromRef1);
+    expect(callback).toHaveBeenCalledTimes(1);
+    expect(ref1).toHaveBeenCalledTimes(1);
+
+    const ref2 = jest.fn(
+      (instance: React.ElementRef<typeof Image> | null): void => {
+        imageInstanceFromRef2 = instance;
+      },
+    );
+
+    act(() => {
+      testRenderer.update(<Image source={{uri: 'foo-bar.jpg'}} ref={ref2} />);
+    });
+
+    expect(imageInstanceFromCallback).not.toBe(null);
+    expect(imageInstanceFromRef1).toBe(null);
+    expect(imageInstanceFromRef2).not.toBe(null);
+    expect(imageInstanceFromCallback).toBe(imageInstanceFromRef2);
+    expect(callback).toHaveBeenCalledTimes(2);
+    expect(ref1).toHaveBeenCalledTimes(2);
+    expect(ref2).toHaveBeenCalledTimes(1);
+
+    act(() => {
+      testRenderer.update(<Image source={{uri: 'foo-bar.jpg'}} ref={ref2} />);
+    });
+
+    expect(callback).toHaveBeenCalledTimes(2);
+    expect(ref2).toHaveBeenCalledTimes(1);
+  });
+
   it('should call image attached callbacks (basic)', () => {
     jest.dontMock('../Image');