forwardRef() components should not re-render on deep setState() (#12690)

* Add a failing test for forwardRef memoization

* Memoize forwardRef props and bail out on strict equality

* Bail out only when ref matches the current ref

Author: gaearon

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -169,12 +169,20 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
 
   function updateForwardRef(current, workInProgress) {
     const render = workInProgress.type.render;
-    const nextChildren = render(
-      workInProgress.pendingProps,
-      workInProgress.ref,
-    );
+    const nextProps = workInProgress.pendingProps;
+    const ref = workInProgress.ref;
+    if (hasLegacyContextChanged()) {
+      // Normally we can bail out on props equality but if context has changed
+      // we don't do the bailout and we have to reuse existing props instead.
+    } else if (workInProgress.memoizedProps === nextProps) {
+      const currentRef = current !== null ? current.ref : null;
+      if (ref === currentRef) {
+        return bailoutOnAlreadyFinishedWork(current, workInProgress);
+      }
+    }
+    const nextChildren = render(nextProps, ref);
     reconcileChildren(current, workInProgress, nextChildren);
-    memoizeProps(workInProgress, nextChildren);
+    memoizeProps(workInProgress, nextProps);
     return workInProgress.child;
   }
 

packages/react/src/__tests__/forwardRef-test.internal.js

@@ -232,6 +232,39 @@ describe('forwardRef', () => {
     expect(ref.current).toBe(null);
   });
 
+  it('should not re-run the render callback on a deep setState', () => {
+    let inst;
+
+    class Inner extends React.Component {
+      render() {
+        ReactNoop.yield('Inner');
+        inst = this;
+        return <div ref={this.props.forwardedRef} />;
+      }
+    }
+
+    function Middle(props) {
+      ReactNoop.yield('Middle');
+      return <Inner {...props} />;
+    }
+
+    const Forward = React.forwardRef((props, ref) => {
+      ReactNoop.yield('Forward');
+      return <Middle {...props} forwardedRef={ref} />;
+    });
+
+    function App() {
+      ReactNoop.yield('App');
+      return <Forward />;
+    }
+
+    ReactNoop.render(<App />);
+    expect(ReactNoop.flush()).toEqual(['App', 'Forward', 'Middle', 'Inner']);
+
+    inst.setState({});
+    expect(ReactNoop.flush()).toEqual(['Inner']);
+  });
+
   it('should warn if not provided a callback during creation', () => {
     expect(() => React.forwardRef(undefined)).toWarnDev(
       'forwardRef requires a render function but was given undefined.',