Allow forwardRef to be composed with pure

This uses a different approach than the moddable approach because that
approach also started getting quite complex and affects tags and symbols.

This way we don't affect the perf of pure and when we drop forwardRef
we can just drop this path.

This also handles the composition of multiple pure wrappers with custom
comparisons.

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -196,6 +196,7 @@ function updateForwardRef(
   workInProgress: Fiber,
   type: any,
   nextProps: any,
+  updateExpirationTime: ExpirationTime,
   renderExpirationTime: ExpirationTime,
 ) {
   const render = type.render;
@@ -212,6 +213,23 @@ function updateForwardRef(
         renderExpirationTime,
       );
     }
+  } else if (
+    current !== null &&
+    type.compare !== undefined &&
+    (updateExpirationTime === NoWork ||
+      updateExpirationTime > renderExpirationTime)
+  ) {
+    const prevProps = current.memoizedProps;
+    // Default to shallow comparison
+    let compare = type.compare;
+    compare = compare !== null ? compare : shallowEqual;
+    if (compare(prevProps, nextProps) && current.ref === workInProgress.ref) {
+      return bailoutOnAlreadyFinishedWork(
+        current,
+        workInProgress,
+        renderExpirationTime,
+      );
+    }
   }
 
   let nextChildren;
@@ -769,6 +787,7 @@ function mountIndeterminateComponent(
           workInProgress,
           Component,
           resolvedProps,
+          updateExpirationTime,
           renderExpirationTime,
         );
         break;
@@ -1561,6 +1580,7 @@ function beginWork(
         workInProgress,
         type,
         workInProgress.pendingProps,
+        updateExpirationTime,
         renderExpirationTime,
       );
     }
@@ -1573,6 +1593,7 @@ function beginWork(
         workInProgress,
         Component,
         resolveDefaultProps(Component, unresolvedProps),
+        updateExpirationTime,
         renderExpirationTime,
       );
       workInProgress.memoizedProps = unresolvedProps;

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

@@ -226,4 +226,131 @@ describe('forwardRef', () => {
         '    in Foo (at **)',
     );
   });
+
+  it('should not bailout if forwardRef is not wrapped in pure', () => {
+    const Component = props => <div {...props} />;
+
+    let renderCount = 0;
+
+    const RefForwardingComponent = React.forwardRef((props, ref) => {
+      renderCount++;
+      return <Component {...props} forwardedRef={ref} />;
+    });
+
+    const ref = React.createRef();
+
+    ReactNoop.render(<RefForwardingComponent ref={ref} optional="foo" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(1);
+
+    ReactNoop.render(<RefForwardingComponent ref={ref} optional="foo" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(2);
+  });
+
+  it('should bailout if forwardRef is wrapped in pure', () => {
+    const Component = props => <div ref={props.forwardedRef} />;
+
+    let renderCount = 0;
+
+    const RefForwardingComponent = React.pure(
+      React.forwardRef((props, ref) => {
+        renderCount++;
+        return <Component {...props} forwardedRef={ref} />;
+      }),
+    );
+
+    const ref = React.createRef();
+
+    ReactNoop.render(<RefForwardingComponent ref={ref} optional="foo" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(1);
+
+    expect(ref.current.type).toBe('div');
+
+    ReactNoop.render(<RefForwardingComponent ref={ref} optional="foo" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(1);
+
+    const differentRef = React.createRef();
+
+    ReactNoop.render(
+      <RefForwardingComponent ref={differentRef} optional="foo" />,
+    );
+    ReactNoop.flush();
+    expect(renderCount).toBe(2);
+
+    expect(ref.current).toBe(null);
+    expect(differentRef.current.type).toBe('div');
+
+    ReactNoop.render(<RefForwardingComponent ref={ref} optional="bar" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(3);
+  });
+
+  it('should custom pure comparisons to compose', () => {
+    const Component = props => <div ref={props.forwardedRef} />;
+
+    let renderCount = 0;
+
+    const RefForwardingComponent = React.pure(
+      React.forwardRef((props, ref) => {
+        renderCount++;
+        return <Component {...props} forwardedRef={ref} />;
+      }),
+      (o, p) => o.a === p.a && o.b === p.b,
+    );
+
+    const ref = React.createRef();
+
+    ReactNoop.render(<RefForwardingComponent ref={ref} a="0" b="0" c="1" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(1);
+
+    expect(ref.current.type).toBe('div');
+
+    // Changing either a or b rerenders
+    ReactNoop.render(<RefForwardingComponent ref={ref} a="0" b="1" c="1" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(2);
+
+    // Changing c doesn't rerender
+    ReactNoop.render(<RefForwardingComponent ref={ref} a="0" b="1" c="2" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(2);
+
+    const ComposedPure = React.pure(
+      RefForwardingComponent,
+      (o, p) => o.a === p.a && o.c === p.c,
+    );
+
+    ReactNoop.render(<ComposedPure ref={ref} a="0" b="0" c="0" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(3);
+
+    // Changing just b no longer updates
+    ReactNoop.render(<ComposedPure ref={ref} a="0" b="1" c="0" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(3);
+
+    // Changing just a and c updates
+    ReactNoop.render(<ComposedPure ref={ref} a="2" b="2" c="2" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(4);
+
+    // Changing just c does not update
+    ReactNoop.render(<ComposedPure ref={ref} a="2" b="2" c="3" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(4);
+
+    // Changing ref still rerenders
+    const differentRef = React.createRef();
+
+    ReactNoop.render(<ComposedPure ref={differentRef} a="2" b="2" c="3" />);
+    ReactNoop.flush();
+    expect(renderCount).toBe(5);
+
+    expect(ref.current).toBe(null);
+    expect(differentRef.current.type).toBe('div');
+  });
 });

packages/react/src/forwardRef.js

@@ -9,9 +9,15 @@ import {REACT_FORWARD_REF_TYPE} from 'shared/ReactSymbols';
 
 import warningWithoutStack from 'shared/warningWithoutStack';
 
+export type ForwardRef<Props> = {
+  $$typeof: Symbol,
+  render: (props: Props, ref: React$Ref<ElementType>) => React$Node,
+  compare: void | null | ((oldProps: Props, newProps: Props) => boolean),
+};
+
 export default function forwardRef<Props, ElementType: React$ElementType>(
   render: (props: Props, ref: React$Ref<ElementType>) => React$Node,
-) {
+): ForwardRef<Props> {
   if (__DEV__) {
     if (typeof render !== 'function') {
       warningWithoutStack(
@@ -42,5 +48,6 @@ export default function forwardRef<Props, ElementType: React$ElementType>(
   return {
     $$typeof: REACT_FORWARD_REF_TYPE,
     render,
+    compare: undefined,
   };
 }

packages/react/src/pure.js

@@ -5,14 +5,37 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {REACT_PURE_TYPE} from 'shared/ReactSymbols';
+import type {ForwardRef} from './forwardRef';
+
+import {REACT_PURE_TYPE, REACT_FORWARD_REF_TYPE} from 'shared/ReactSymbols';
 
 import warningWithoutStack from 'shared/warningWithoutStack';
 
+function composeComparisonFunctions(inner, outer) {
+  return function(oldProps, newProps) {
+    // Either is allowed to block the update.
+    return outer(oldProps, newProps) || inner(oldProps, newProps);
+  };
+}
+
 export default function pure<Props>(
-  render: (props: Props) => React$Node,
+  render: (props: Props) => React$Node | ForwardRef,
   compare?: (oldProps: Props, newProps: Props) => boolean,
 ) {
+  if (
+    typeof render === 'object' &&
+    render.$$typeof === REACT_FORWARD_REF_TYPE
+  ) {
+    let forwardRef = (render: ForwardRef);
+    if (forwardRef.compare !== undefined && forwardRef.compare !== null) {
+      compare = composeComparisonFunctions(forwardRef.compare, compare);
+    }
+    return {
+      $$typeof: REACT_FORWARD_REF_TYPE,
+      render: forwardRef.render,
+      compare: compare === undefined ? null : compare,
+    };
+  }
   if (__DEV__) {
     if (typeof render !== 'function') {
       warningWithoutStack(