From c19335bac9b0eaaba267b0f632e86adcab6b0db8 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Fri, 19 Oct 2018 14:24:33 -0700
Subject: [PATCH] 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.
---
 .../src/ReactFiberBeginWork.js                |  21 +++
 .../react/src/__tests__/forwardRef-test.js    | 127 ++++++++++++++++++
 packages/react/src/forwardRef.js              |   9 +-
 packages/react/src/pure.js                    |  27 +++-
 4 files changed, 181 insertions(+), 3 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 0cf92b4b2734..2450f3c82ec3 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/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;
diff --git a/packages/react/src/__tests__/forwardRef-test.js b/packages/react/src/__tests__/forwardRef-test.js
index c22424c00e18..e7ef1ffeabbb 100644
--- a/packages/react/src/__tests__/forwardRef-test.js
+++ b/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');
+  });
 });
diff --git a/packages/react/src/forwardRef.js b/packages/react/src/forwardRef.js
index 1018375e8c92..9c9f721d821b 100644
--- a/packages/react/src/forwardRef.js
+++ b/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,
   };
 }
diff --git a/packages/react/src/pure.js b/packages/react/src/pure.js
index 0431c05269ce..941fb50c1d93 100644
--- a/packages/react/src/pure.js
+++ b/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(