Allow arbitrary types to be wrapped in pure (#13903)

* Allow arbitrary types to be wrapped in pure

This creates an outer fiber that container the pure check and an inner
fiber that represents which ever type of component.

* Add optimized fast path for simple pure function components

Special cased when there are no defaultProps and it's a simple function
component instead of class. This doesn't require an extra fiber.

We could make it so that this also works with custom comparer but that
means we have to go through one extra indirection to get to it.
Maybe it's worth it, donno.

Author: sebmarkbage

packages/react-reconciler/src/ReactFiber.js

@@ -299,10 +299,15 @@ function shouldConstruct(Component: Function) {
   return !!(prototype && prototype.isReactComponent);
 }
 
-export function resolveLazyComponentTag(
-  fiber: Fiber,
-  Component: Function,
-): WorkTag {
+export function isSimpleFunctionComponent(type: any) {
+  return (
+    typeof type === 'function' &&
+    !shouldConstruct(type) &&
+    type.defaultProps === undefined
+  );
+}
+
+export function resolveLazyComponentTag(Component: Function): WorkTag {
   if (typeof Component === 'function') {
     return shouldConstruct(Component) ? ClassComponent : FunctionComponent;
   } else if (Component !== undefined && Component !== null) {
@@ -406,20 +411,15 @@ export function createHostRootFiber(isConcurrent: boolean): Fiber {
   return createFiber(HostRoot, null, null, mode);
 }
 
-function createFiberFromElementWithoutDebugInfo(
-  element: ReactElement,
+export function createFiberFromTypeAndProps(
+  type: any, // React$ElementType
+  key: null | string,
+  pendingProps: any,
+  owner: null | Fiber,
   mode: TypeOfMode,
   expirationTime: ExpirationTime,
 ): Fiber {
-  let owner = null;
-  if (__DEV__) {
-    owner = element._owner;
-  }
-
   let fiber;
-  const type = element.type;
-  const key = element.key;
-  const pendingProps = element.props;
 
   let fiberTag = IndeterminateComponent;
   // The resolved type is set if we know what the final type will be. I.e. it's not lazy.
@@ -522,8 +522,18 @@ export function createFiberFromElement(
   mode: TypeOfMode,
   expirationTime: ExpirationTime,
 ): Fiber {
-  const fiber = createFiberFromElementWithoutDebugInfo(
-    element,
+  let owner = null;
+  if (__DEV__) {
+    owner = element._owner;
+  }
+  const type = element.type;
+  const key = element.key;
+  const pendingProps = element.props;
+  const fiber = createFiberFromTypeAndProps(
+    type,
+    key,
+    pendingProps,
+    owner,
     mode,
     expirationTime,
   );

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -31,6 +31,7 @@ import {
   Profiler,
   SuspenseComponent,
   PureComponent,
+  SimplePureComponent,
   LazyComponent,
 } from 'shared/ReactWorkTags';
 import {
@@ -103,8 +104,10 @@ import {
 import {readLazyComponentType} from './ReactFiberLazyComponent';
 import {
   resolveLazyComponentTag,
+  createFiberFromTypeAndProps,
   createFiberFromFragment,
   createWorkInProgress,
+  isSimpleFunctionComponent,
 } from './ReactFiber';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
@@ -235,48 +238,99 @@ function updatePureComponent(
   nextProps: any,
   updateExpirationTime,
   renderExpirationTime: ExpirationTime,
-) {
-  const render = Component.render;
-
+): null | Fiber {
+  if (current === null) {
+    let type = Component.type;
+    if (isSimpleFunctionComponent(type) && Component.compare === null) {
+      // If this is a plain function component without default props,
+      // and with only the default shallow comparison, we upgrade it
+      // to a SimplePureComponent to allow fast path updates.
+      workInProgress.tag = SimplePureComponent;
+      workInProgress.type = type;
+      return updateSimplePureComponent(
+        current,
+        workInProgress,
+        type,
+        nextProps,
+        updateExpirationTime,
+        renderExpirationTime,
+      );
+    }
+    let child = createFiberFromTypeAndProps(
+      Component.type,
+      null,
+      nextProps,
+      null,
+      workInProgress.mode,
+      renderExpirationTime,
+    );
+    child.ref = workInProgress.ref;
+    child.return = workInProgress;
+    workInProgress.child = child;
+    return child;
+  }
+  let currentChild = ((current.child: any): Fiber); // This is always exactly one child
   if (
-    current !== null &&
-    (updateExpirationTime === NoWork ||
-      updateExpirationTime > renderExpirationTime)
+    updateExpirationTime === NoWork ||
+    updateExpirationTime > renderExpirationTime
   ) {
-    const prevProps = current.memoizedProps;
+    // This will be the props with resolved defaultProps,
+    // unlike current.memoizedProps which will be the unresolved ones.
+    const prevProps = currentChild.memoizedProps;
     // Default to shallow comparison
     let compare = Component.compare;
     compare = compare !== null ? compare : shallowEqual;
-    if (compare(prevProps, nextProps)) {
+    if (compare(prevProps, nextProps) && current.ref === workInProgress.ref) {
       return bailoutOnAlreadyFinishedWork(
         current,
         workInProgress,
         renderExpirationTime,
       );
     }
   }
+  let newChild = createWorkInProgress(
+    currentChild,
+    nextProps,
+    renderExpirationTime,
+  );
+  newChild.ref = workInProgress.ref;
+  newChild.return = workInProgress;
+  workInProgress.child = newChild;
+  return newChild;
+}
 
-  // The rest is a fork of updateFunctionComponent
-  let nextChildren;
-  prepareToReadContext(workInProgress, renderExpirationTime);
-  if (__DEV__) {
-    ReactCurrentOwner.current = workInProgress;
-    ReactCurrentFiber.setCurrentPhase('render');
-    nextChildren = render(nextProps);
-    ReactCurrentFiber.setCurrentPhase(null);
-  } else {
-    nextChildren = render(nextProps);
+function updateSimplePureComponent(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  Component: any,
+  nextProps: any,
+  updateExpirationTime,
+  renderExpirationTime: ExpirationTime,
+): null | Fiber {
+  if (
+    current !== null &&
+    (updateExpirationTime === NoWork ||
+      updateExpirationTime > renderExpirationTime)
+  ) {
+    const prevProps = current.memoizedProps;
+    if (
+      shallowEqual(prevProps, nextProps) &&
+      current.ref === workInProgress.ref
+    ) {
+      return bailoutOnAlreadyFinishedWork(
+        current,
+        workInProgress,
+        renderExpirationTime,
+      );
+    }
   }
-
-  // React DevTools reads this flag.
-  workInProgress.effectTag |= PerformedWork;
-  reconcileChildren(
+  return updateFunctionComponent(
     current,
     workInProgress,
-    nextChildren,
+    Component,
+    nextProps,
     renderExpirationTime,
   );
-  return workInProgress.child;
 }
 
 function updateFragment(
@@ -725,10 +779,7 @@ function mountLazyComponent(
   let Component = readLazyComponentType(elementType);
   // Store the unwrapped component in the type.
   workInProgress.type = Component;
-  const resolvedTag = (workInProgress.tag = resolveLazyComponentTag(
-    workInProgress,
-    Component,
-  ));
+  const resolvedTag = (workInProgress.tag = resolveLazyComponentTag(Component));
   startWorkTimer(workInProgress);
   const resolvedProps = resolveDefaultProps(Component, props);
   let child;
@@ -768,7 +819,7 @@ function mountLazyComponent(
         null,
         workInProgress,
         Component,
-        resolvedProps,
+        resolveDefaultProps(Component.type, resolvedProps), // The inner type can have defaults too
         updateExpirationTime,
         renderExpirationTime,
       );
@@ -1564,10 +1615,7 @@ function beginWork(
     case PureComponent: {
       const type = workInProgress.type;
       const unresolvedProps = workInProgress.pendingProps;
-      const resolvedProps =
-        workInProgress.elementType === type
-          ? unresolvedProps
-          : resolveDefaultProps(type, unresolvedProps);
+      const resolvedProps = resolveDefaultProps(type.type, unresolvedProps);
       return updatePureComponent(
         current,
         workInProgress,
@@ -1577,6 +1625,16 @@ function beginWork(
         renderExpirationTime,
       );
     }
+    case SimplePureComponent: {
+      return updateSimplePureComponent(
+        current,
+        workInProgress,
+        workInProgress.type,
+        workInProgress.pendingProps,
+        updateExpirationTime,
+        renderExpirationTime,
+      );
+    }
     default:
       invariant(
         false,

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -35,6 +35,7 @@ import {
   Profiler,
   SuspenseComponent,
   PureComponent,
+  SimplePureComponent,
   LazyComponent,
 } from 'shared/ReactWorkTags';
 import {Placement, Ref, Update} from 'shared/ReactSideEffectTags';
@@ -539,6 +540,7 @@ function completeWork(
       break;
     case LazyComponent:
       break;
+    case SimplePureComponent:
     case FunctionComponent:
       break;
     case ClassComponent: {

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

@@ -183,21 +183,49 @@ describe('pure', () => {
         expect(ReactNoop.getChildren()).toEqual([span(1)]);
       });
 
-      it('warns for class components', () => {
-        class SomeClass extends React.Component {
+      it('supports non-pure class components', async () => {
+        const {unstable_Suspense: Suspense} = React;
+
+        class CounterInner extends React.Component {
+          static defaultProps = {suffix: '!'};
           render() {
-            return null;
+            return <Text text={this.props.count + '' + this.props.suffix} />;
           }
         }
-        expect(() => pure(SomeClass)).toWarnDev(
-          'pure: The first argument must be a function component.',
-          {withoutStack: true},
+        const Counter = pure(CounterInner);
+
+        ReactNoop.render(
+          <Suspense fallback={<Text text="Loading..." />}>
+            <Counter count={0} />
+          </Suspense>,
+        );
+        expect(ReactNoop.flush()).toEqual(['Loading...']);
+        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([]);
+        expect(ReactNoop.getChildren()).toEqual([span('0!')]);
+
+        // Should update because count prop changed
+        ReactNoop.render(
+          <Suspense>
+            <Counter count={1} />
+          </Suspense>,
         );
+        expect(ReactNoop.flush()).toEqual(['1!']);
+        expect(ReactNoop.getChildren()).toEqual([span('1!')]);
       });
 
-      it('warns if first argument is not a function', () => {
+      it('warns if first argument is undefined', () => {
         expect(() => pure()).toWarnDev(
-          'pure: The first argument must be a function component. Instead ' +
+          'pure: The first argument must be a component. Instead ' +
             'received: undefined',
           {withoutStack: true},
         );

packages/react-reconciler/src/__tests__/__snapshots__/ReactIncrementalPerf-test.internal.js.snap

@@ -415,7 +415,7 @@ exports[`ReactDebugFiberPerf supports pure 1`] = `
 
 ⚛ (React Tree Reconciliation: Completed Root)
   ⚛ Parent [mount]
-    ⚛ Pure(Foo) [mount]
+    ⚛ Foo [mount]
 
 ⚛ (Committing Changes)
   ⚛ (Committing Snapshot Effects: 0 Total)

packages/react-test-renderer/src/ReactTestRenderer.js

@@ -28,6 +28,7 @@ import {
   ForwardRef,
   Profiler,
   PureComponent,
+  SimplePureComponent,
 } from 'shared/ReactWorkTags';
 import invariant from 'shared/invariant';
 import ReactVersion from 'shared/ReactVersion';
@@ -166,6 +167,7 @@ function toTree(node: ?Fiber) {
         rendered: childrenToTree(node.child),
       };
     case FunctionComponent:
+    case SimplePureComponent:
       return {
         nodeType: 'component',
         type: node.type,