React.pure automatically forwards ref

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -236,6 +236,7 @@ function updatePureComponent(
   renderExpirationTime: ExpirationTime,
 ) {
   const render = Component.render;
+  const ref = workInProgress.ref;
 
   if (
     current !== null &&
@@ -246,7 +247,7 @@ function updatePureComponent(
     // Default to shallow comparison
     let compare = Component.compare;
     compare = compare !== null ? compare : shallowEqual;
-    if (compare(prevProps, nextProps)) {
+    if (workInProgress.ref === current.ref && compare(prevProps, nextProps)) {
       return bailoutOnAlreadyFinishedWork(
         current,
         workInProgress,
@@ -261,10 +262,10 @@ function updatePureComponent(
   if (__DEV__) {
     ReactCurrentOwner.current = workInProgress;
     ReactCurrentFiber.setCurrentPhase('render');
-    nextChildren = render(nextProps);
+    nextChildren = render(nextProps, ref);
     ReactCurrentFiber.setCurrentPhase(null);
   } else {
-    nextChildren = render(nextProps);
+    nextChildren = render(nextProps, ref);
   }
 
   // React DevTools reads this flag.

packages/react-reconciler/src/__tests__/ReactPure-test.internal.js

@@ -35,14 +35,14 @@ describe('pure', () => {
   }
 
   // Tests should run against both the lazy and non-lazy versions of `pure`.
-  // To make the tests work for both versions, we wrap the non-lazy verion in
+  // To make the tests work for both versions, we wrap the non-lazy version in
   // a lazy function component.
   sharedTests('normal', (...args) => {
     const Pure = React.pure(...args);
-    function Indirection(props) {
-      return <Pure {...props} />;
+    function Indirection(props, ref) {
+      return <Pure {...props} ref={ref} />;
     }
-    return Promise.resolve(Indirection);
+    return Promise.resolve(React.forwardRef(Indirection));
   });
   sharedTests('lazy', (...args) => Promise.resolve(React.pure(...args)));
 
@@ -84,110 +84,169 @@ describe('pure', () => {
         expect(ReactNoop.flush()).toEqual([1]);
         expect(ReactNoop.getChildren()).toEqual([span(1)]);
       });
-    });
 
-    it("does not bail out if there's a context change", async () => {
-      const {unstable_Suspense: Suspense} = React;
-
-      const CountContext = React.createContext(0);
-
-      function Counter(props) {
-        const count = CountContext.unstable_read();
-        return <Text text={`${props.label}: ${count}`} />;
-      }
-      Counter = pure(Counter);
-
-      class Parent extends React.Component {
-        state = {count: 0};
-        render() {
-          return (
-            <Suspense>
-              <CountContext.Provider value={this.state.count}>
-                <Counter label="Count" />
-              </CountContext.Provider>
-            </Suspense>
+      it("does not bail out if there's a context change", async () => {
+        const {unstable_Suspense: Suspense} = React;
+
+        const CountContext = React.createContext(0);
+
+        function Counter(props) {
+          const count = CountContext.unstable_read();
+          return <Text text={`${props.label}: ${count}`} />;
+        }
+        Counter = pure(Counter);
+
+        class Parent extends React.Component {
+          state = {count: 0};
+          render() {
+            return (
+              <Suspense>
+                <CountContext.Provider value={this.state.count}>
+                  <Counter label="Count" />
+                </CountContext.Provider>
+              </Suspense>
+            );
+          }
+        }
+
+        const parent = React.createRef(null);
+        ReactNoop.render(<Parent ref={parent} />);
+        expect(ReactNoop.flush()).toEqual([]);
+        await Promise.resolve();
+        expect(ReactNoop.flush()).toEqual(['Count: 0']);
+        expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
+
+        // Should bail out because props have not changed
+        ReactNoop.render(<Parent ref={parent} />);
+        expect(ReactNoop.flush()).toEqual([]);
+        expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
+
+        // Should update because there was a context change
+        parent.current.setState({count: 1});
+        expect(ReactNoop.flush()).toEqual(['Count: 1']);
+        expect(ReactNoop.getChildren()).toEqual([span('Count: 1')]);
+      });
+
+      it('accepts custom comparison function', async () => {
+        const {unstable_Suspense: Suspense} = React;
+
+        function Counter({count}) {
+          return <Text text={count} />;
+        }
+        Counter = pure(Counter, (oldProps, newProps) => {
+          ReactNoop.yield(
+            `Old count: ${oldProps.count}, New count: ${newProps.count}`,
           );
+          return oldProps.count === newProps.count;
+        });
+
+        ReactNoop.render(
+          <Suspense>
+            <Counter count={0} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual([]);
+        await Promise.resolve();
+        expect(ReactNoop.flush()).toEqual([0]);
+        expect(ReactNoop.getChildren()).toEqual([span(0)]);
+
+        // Should bail out because props have not changed
+        ReactNoop.render(
+          <Suspense>
+            <Counter count={0} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual(['Old count: 0, New count: 0']);
+        expect(ReactNoop.getChildren()).toEqual([span(0)]);
+
+        // Should update because count prop changed
+        ReactNoop.render(
+          <Suspense>
+            <Counter count={1} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual(['Old count: 0, New count: 1', 1]);
+        expect(ReactNoop.getChildren()).toEqual([span(1)]);
+      });
+
+      it('warns for class components', () => {
+        class SomeClass extends React.Component {
+          render() {
+            return null;
+          }
         }
-      }
-
-      const parent = React.createRef(null);
-      ReactNoop.render(<Parent ref={parent} />);
-      expect(ReactNoop.flush()).toEqual([]);
-      await Promise.resolve();
-      expect(ReactNoop.flush()).toEqual(['Count: 0']);
-      expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
-
-      // Should bail out because props have not changed
-      ReactNoop.render(<Parent ref={parent} />);
-      expect(ReactNoop.flush()).toEqual([]);
-      expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
-
-      // Should update because there was a context change
-      parent.current.setState({count: 1});
-      expect(ReactNoop.flush()).toEqual(['Count: 1']);
-      expect(ReactNoop.getChildren()).toEqual([span('Count: 1')]);
-    });
+        expect(() => pure(SomeClass)).toWarnDev(
+          'pure: The first argument must be a function component.',
+          {withoutStack: true},
+        );
+      });
 
-    it('accepts custom comparison function', async () => {
-      const {unstable_Suspense: Suspense} = React;
+      it('warns if first argument is not a function', () => {
+        expect(() => pure()).toWarnDev(
+          'pure: The first argument must be a function component. Instead ' +
+            'received: undefined',
+          {withoutStack: true},
+        );
+      });
 
-      function Counter({count}) {
-        return <Text text={count} />;
-      }
-      Counter = pure(Counter, (oldProps, newProps) => {
-        ReactNoop.yield(
-          `Old count: ${oldProps.count}, New count: ${newProps.count}`,
+      it('forwards ref', async () => {
+        const {unstable_Suspense: Suspense} = React;
+        const Transparent = pure((props, ref) => {
+          return <div ref={ref} />;
+        });
+        const divRef = React.createRef();
+
+        ReactNoop.render(
+          <Suspense>
+            <Transparent ref={divRef} />
+          </Suspense>,
         );
-        return oldProps.count === newProps.count;
+        ReactNoop.flush();
+        await Promise.resolve();
+        ReactNoop.flush();
+        expect(divRef.current.type).toBe('div');
       });
 
-      ReactNoop.render(
-        <Suspense>
-          <Counter count={0} />
-        </Suspense>,
-      );
-      expect(ReactNoop.flush()).toEqual([]);
-      await Promise.resolve();
-      expect(ReactNoop.flush()).toEqual([0]);
-      expect(ReactNoop.getChildren()).toEqual([span(0)]);
-
-      // Should bail out because props have not changed
-      ReactNoop.render(
-        <Suspense>
-          <Counter count={0} />
-        </Suspense>,
-      );
-      expect(ReactNoop.flush()).toEqual(['Old count: 0, New count: 0']);
-      expect(ReactNoop.getChildren()).toEqual([span(0)]);
-
-      // Should update because count prop changed
-      ReactNoop.render(
-        <Suspense>
-          <Counter count={1} />
-        </Suspense>,
-      );
-      expect(ReactNoop.flush()).toEqual(['Old count: 0, New count: 1', 1]);
-      expect(ReactNoop.getChildren()).toEqual([span(1)]);
-    });
+      it('updates if only ref changes', async () => {
+        const {unstable_Suspense: Suspense} = React;
+        const Transparent = pure((props, ref) => {
+          return [<Text key="text" text="Text" />, <div key="div" ref={ref} />];
+        });
 
-    it('warns for class components', () => {
-      class SomeClass extends React.Component {
-        render() {
-          return null;
-        }
-      }
-      expect(() => pure(SomeClass)).toWarnDev(
-        'pure: The first argument must be a function component.',
-        {withoutStack: true},
-      );
-    });
+        const divRef = React.createRef();
+        const divRef2 = React.createRef();
+
+        ReactNoop.render(
+          <Suspense>
+            <Transparent ref={divRef} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual([]);
+        await Promise.resolve();
+        expect(ReactNoop.flush()).toEqual(['Text']);
+        expect(divRef.current.type).toBe('div');
+        expect(divRef2.current).toBe(null);
+
+        // Should re-render (new ref)
+        ReactNoop.render(
+          <Suspense>
+            <Transparent ref={divRef2} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual(['Text']);
+        expect(divRef.current).toBe(null);
+        expect(divRef2.current.type).toBe('div');
 
-    it('warns if first argument is not a function', () => {
-      expect(() => pure()).toWarnDev(
-        'pure: The first argument must be a function component. Instead ' +
-          'received: undefined',
-        {withoutStack: true},
-      );
+        // Should not re-render (same ref)
+        ReactNoop.render(
+          <Suspense>
+            <Transparent ref={divRef2} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual([]);
+        expect(divRef.current).toBe(null);
+        expect(divRef2.current.type).toBe('div');
+      });
     });
   }
 });