Fix: Class components should "consume" ref prop (facebook#28719)

When a ref is passed to a class component, the class instance is
attached to the ref's current property automatically. This different
from function components, where you have to do something extra to attach
a ref to an instance, like passing the ref to `useImperativeHandle`.

Existing class component code is written with the assumption that a ref
will not be passed through as a prop. For example, class components that
act as indirections often spread `this.props` onto a child component. To
maintain this expectation, we should remove the ref from the props
object ("consume" it) before passing it to lifecycle methods. Without
this change, much existing code will break because the ref will attach
to the inner component instead of the outer one.

This is not an issue for function components because we used to warn if
you passed a ref to a function component. Instead, you had to use
`forwardRef`, which also implements this "consuming" behavior.

There are a few places in the reconciler where we modify the fiber's
internal props object before passing it to userspace. The trickiest one
is class components, because the props object gets exposed in many
different places, including as a property on the class instance.

This was already accounted for when we added support for setting default
props on a lazy wrapper (i.e. `React.lazy` that resolves to a class
component).

In all of these same places, we will also need to remove the ref prop
when `enableRefAsProp` is on.

Closes facebook#28602

---------

Co-authored-by: Jan Kassens <jan@kassens.net>

packages/react-dom/src/__tests__/ReactCompositeComponent-test.js

@@ -261,29 +261,17 @@ describe('ReactCompositeComponent', () => {
     await act(() => {
       root.render(<Component ref={refFn1} />);
     });
-    if (gate(flags => flags.enableRefAsProp)) {
-      expect(instance1.props).toEqual({prop: 'testKey', ref: refFn1});
-    } else {
-      expect(instance1.props).toEqual({prop: 'testKey'});
-    }
+    expect(instance1.props).toEqual({prop: 'testKey'});
 
     await act(() => {
       root.render(<Component ref={refFn2} prop={undefined} />);
     });
-    if (gate(flags => flags.enableRefAsProp)) {
-      expect(instance2.props).toEqual({prop: 'testKey', ref: refFn2});
-    } else {
-      expect(instance2.props).toEqual({prop: 'testKey'});
-    }
+    expect(instance2.props).toEqual({prop: 'testKey'});
 
     await act(() => {
       root.render(<Component ref={refFn3} prop={null} />);
     });
-    if (gate(flags => flags.enableRefAsProp)) {
-      expect(instance3.props).toEqual({prop: null, ref: refFn3});
-    } else {
-      expect(instance3.props).toEqual({prop: null});
-    }
+    expect(instance3.props).toEqual({prop: null});
   });
 
   it('should not mutate passed-in props object', async () => {

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -245,6 +245,7 @@ import {
   mountClassInstance,
   resumeMountClassInstance,
   updateClassInstance,
+  resolveClassComponentProps,
 } from './ReactFiberClassComponent';
 import {resolveDefaultProps} from './ReactFiberLazyComponent';
 import {
@@ -1763,9 +1764,9 @@ function mountLazyComponent(
   // Store the unwrapped component in the type.
   workInProgress.type = Component;
 
-  const resolvedProps = resolveDefaultProps(Component, props);
   if (typeof Component === 'function') {
     if (isFunctionClassComponent(Component)) {
+      const resolvedProps = resolveClassComponentProps(Component, props, false);
       workInProgress.tag = ClassComponent;
       if (__DEV__) {
         workInProgress.type = Component =
@@ -1779,6 +1780,7 @@ function mountLazyComponent(
         renderLanes,
       );
     } else {
+      const resolvedProps = resolveDefaultProps(Component, props);
       workInProgress.tag = FunctionComponent;
       if (__DEV__) {
         validateFunctionComponentInDev(workInProgress, Component);
@@ -1796,6 +1798,7 @@ function mountLazyComponent(
   } else if (Component !== undefined && Component !== null) {
     const $$typeof = Component.$$typeof;
     if ($$typeof === REACT_FORWARD_REF_TYPE) {
+      const resolvedProps = resolveDefaultProps(Component, props);
       workInProgress.tag = ForwardRef;
       if (__DEV__) {
         workInProgress.type = Component =
@@ -1809,6 +1812,7 @@ function mountLazyComponent(
         renderLanes,
       );
     } else if ($$typeof === REACT_MEMO_TYPE) {
+      const resolvedProps = resolveDefaultProps(Component, props);
       workInProgress.tag = MemoComponent;
       return updateMemoComponent(
         null,
@@ -3939,10 +3943,11 @@ function beginWork(
     case ClassComponent: {
       const Component = workInProgress.type;
       const unresolvedProps = workInProgress.pendingProps;
-      const resolvedProps =
-        workInProgress.elementType === Component
-          ? unresolvedProps
-          : resolveDefaultProps(Component, unresolvedProps);
+      const resolvedProps = resolveClassComponentProps(
+        Component,
+        unresolvedProps,
+        workInProgress.elementType === Component,
+      );
       return updateClassComponent(
         current,
         workInProgress,
@@ -4025,10 +4030,11 @@ function beginWork(
       }
       const Component = workInProgress.type;
       const unresolvedProps = workInProgress.pendingProps;
-      const resolvedProps =
-        workInProgress.elementType === Component
-          ? unresolvedProps
-          : resolveDefaultProps(Component, unresolvedProps);
+      const resolvedProps = resolveClassComponentProps(
+        Component,
+        unresolvedProps,
+        workInProgress.elementType === Component,
+      );
       return mountIncompleteClassComponent(
         current,
         workInProgress,
@@ -4043,10 +4049,11 @@ function beginWork(
       }
       const Component = workInProgress.type;
       const unresolvedProps = workInProgress.pendingProps;
-      const resolvedProps =
-        workInProgress.elementType === Component
-          ? unresolvedProps
-          : resolveDefaultProps(Component, unresolvedProps);
+      const resolvedProps = resolveClassComponentProps(
+        Component,
+        unresolvedProps,
+        workInProgress.elementType === Component,
+      );
       return mountIncompleteFunctionComponent(
         current,
         workInProgress,

packages/react-reconciler/src/ReactFiberClassComponent.js

@@ -23,6 +23,7 @@ import {
   enableDebugTracing,
   enableSchedulingProfiler,
   enableLazyContextPropagation,
+  enableRefAsProp,
 } from 'shared/ReactFeatureFlags';
 import ReactStrictModeWarnings from './ReactStrictModeWarnings';
 import {isMounted} from './ReactFiberTreeReflection';
@@ -34,7 +35,6 @@ import assign from 'shared/assign';
 import isArray from 'shared/isArray';
 import {REACT_CONTEXT_TYPE, REACT_CONSUMER_TYPE} from 'shared/ReactSymbols';
 
-import {resolveDefaultProps} from './ReactFiberLazyComponent';
 import {
   DebugTracingMode,
   NoMode,
@@ -904,7 +904,12 @@ function resumeMountClassInstance(
 ): boolean {
   const instance = workInProgress.stateNode;
 
-  const oldProps = workInProgress.memoizedProps;
+  const unresolvedOldProps = workInProgress.memoizedProps;
+  const oldProps = resolveClassComponentProps(
+    ctor,
+    unresolvedOldProps,
+    workInProgress.type === workInProgress.elementType,
+  );
   instance.props = oldProps;
 
   const oldContext = instance.context;
@@ -926,6 +931,13 @@ function resumeMountClassInstance(
     typeof getDerivedStateFromProps === 'function' ||
     typeof instance.getSnapshotBeforeUpdate === 'function';
 
+  // When comparing whether props changed, we should compare using the
+  // unresolved props object that is stored on the fiber, rather than the
+  // one that gets assigned to the instance, because that object may have been
+  // cloned to resolve default props and/or remove `ref`.
+  const unresolvedNewProps = workInProgress.pendingProps;
+  const didReceiveNewProps = unresolvedNewProps !== unresolvedOldProps;
+
   // Note: During these life-cycles, instance.props/instance.state are what
   // ever the previously attempted to render - not the "current". However,
   // during componentDidUpdate we pass the "current" props.
@@ -937,7 +949,7 @@ function resumeMountClassInstance(
     (typeof instance.UNSAFE_componentWillReceiveProps === 'function' ||
       typeof instance.componentWillReceiveProps === 'function')
   ) {
-    if (oldProps !== newProps || oldContext !== nextContext) {
+    if (didReceiveNewProps || oldContext !== nextContext) {
       callComponentWillReceiveProps(
         workInProgress,
         instance,
@@ -955,7 +967,7 @@ function resumeMountClassInstance(
   suspendIfUpdateReadFromEntangledAsyncAction();
   newState = workInProgress.memoizedState;
   if (
-    oldProps === newProps &&
+    !didReceiveNewProps &&
     oldState === newState &&
     !hasContextChanged() &&
     !checkHasForceUpdateAfterProcessing()
@@ -1052,10 +1064,11 @@ function updateClassInstance(
   cloneUpdateQueue(current, workInProgress);
 
   const unresolvedOldProps = workInProgress.memoizedProps;
-  const oldProps =
-    workInProgress.type === workInProgress.elementType
-      ? unresolvedOldProps
-      : resolveDefaultProps(workInProgress.type, unresolvedOldProps);
+  const oldProps = resolveClassComponentProps(
+    ctor,
+    unresolvedOldProps,
+    workInProgress.type === workInProgress.elementType,
+  );
   instance.props = oldProps;
   const unresolvedNewProps = workInProgress.pendingProps;
 
@@ -1225,6 +1238,42 @@ function updateClassInstance(
   return shouldUpdate;
 }
 
+export function resolveClassComponentProps(
+  Component: any,
+  baseProps: Object,
+  // Only resolve default props if this is a lazy component. Otherwise, they
+  // would have already been resolved by the JSX runtime.
+  // TODO: We're going to remove default prop resolution from the JSX runtime
+  // and keep it only for class components. As part of that change, we should
+  // remove this extra check.
+  alreadyResolvedDefaultProps: boolean,
+): Object {
+  let newProps = baseProps;
+
+  // Resolve default props. Taken from old JSX runtime, where this used to live.
+  const defaultProps = Component.defaultProps;
+  if (defaultProps && !alreadyResolvedDefaultProps) {
+    newProps = assign({}, newProps, baseProps);
+    for (const propName in defaultProps) {
+      if (newProps[propName] === undefined) {
+        newProps[propName] = defaultProps[propName];
+      }
+    }
+  }
+
+  if (enableRefAsProp) {
+    // Remove ref from the props object, if it exists.
+    if ('ref' in newProps) {
+      if (newProps === baseProps) {
+        newProps = assign({}, newProps);
+      }
+      delete newProps.ref;
+    }
+  }
+
+  return newProps;
+}
+
 export {
   constructClassInstance,
   mountClassInstance,

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -104,7 +104,7 @@ import {
   setCurrentFiber as setCurrentDebugFiberInDEV,
   getCurrentFiber as getCurrentDebugFiberInDEV,
 } from './ReactCurrentFiber';
-import {resolveDefaultProps} from './ReactFiberLazyComponent';
+import {resolveClassComponentProps} from './ReactFiberClassComponent';
 import {
   isCurrentUpdateNested,
   getCommitTime,
@@ -244,7 +244,11 @@ function shouldProfile(current: Fiber): boolean {
 }
 
 function callComponentWillUnmountWithTimer(current: Fiber, instance: any) {
-  instance.props = current.memoizedProps;
+  instance.props = resolveClassComponentProps(
+    current.type,
+    current.memoizedProps,
+    current.elementType === current.type,
+  );
   instance.state = current.memoizedState;
   if (shouldProfile(current)) {
     try {
@@ -471,7 +475,8 @@ function commitBeforeMutationEffectsOnFiber(finishedWork: Fiber) {
           // TODO: revisit this when we implement resuming.
           if (__DEV__) {
             if (
-              finishedWork.type === finishedWork.elementType &&
+              !finishedWork.type.defaultProps &&
+              !('ref' in finishedWork.memoizedProps) &&
               !didWarnAboutReassigningProps
             ) {
               if (instance.props !== finishedWork.memoizedProps) {
@@ -497,9 +502,11 @@ function commitBeforeMutationEffectsOnFiber(finishedWork: Fiber) {
             }
           }
           const snapshot = instance.getSnapshotBeforeUpdate(
-            finishedWork.elementType === finishedWork.type
-              ? prevProps
-              : resolveDefaultProps(finishedWork.type, prevProps),
+            resolveClassComponentProps(
+              finishedWork.type,
+              prevProps,
+              finishedWork.elementType === finishedWork.type,
+            ),
             prevState,
           );
           if (__DEV__) {
@@ -807,7 +814,8 @@ function commitClassLayoutLifecycles(
     // TODO: revisit this when we implement resuming.
     if (__DEV__) {
       if (
-        finishedWork.type === finishedWork.elementType &&
+        !finishedWork.type.defaultProps &&
+        !('ref' in finishedWork.memoizedProps) &&
         !didWarnAboutReassigningProps
       ) {
         if (instance.props !== finishedWork.memoizedProps) {
@@ -848,17 +856,19 @@ function commitClassLayoutLifecycles(
       }
     }
   } else {
-    const prevProps =
-      finishedWork.elementType === finishedWork.type
-        ? current.memoizedProps
-        : resolveDefaultProps(finishedWork.type, current.memoizedProps);
+    const prevProps = resolveClassComponentProps(
+      finishedWork.type,
+      current.memoizedProps,
+      finishedWork.elementType === finishedWork.type,
+    );
     const prevState = current.memoizedState;
     // We could update instance props and state here,
     // but instead we rely on them being set during last render.
     // TODO: revisit this when we implement resuming.
     if (__DEV__) {
       if (
-        finishedWork.type === finishedWork.elementType &&
+        !finishedWork.type.defaultProps &&
+        !('ref' in finishedWork.memoizedProps) &&
         !didWarnAboutReassigningProps
       ) {
         if (instance.props !== finishedWork.memoizedProps) {
@@ -918,7 +928,8 @@ function commitClassCallbacks(finishedWork: Fiber) {
     const instance = finishedWork.stateNode;
     if (__DEV__) {
       if (
-        finishedWork.type === finishedWork.elementType &&
+        !finishedWork.type.defaultProps &&
+        !('ref' in finishedWork.memoizedProps) &&
         !didWarnAboutReassigningProps
       ) {
         if (instance.props !== finishedWork.memoizedProps) {

packages/react-reconciler/src/ReactFiberLazyComponent.js

@@ -10,6 +10,10 @@
 import assign from 'shared/assign';
 
 export function resolveDefaultProps(Component: any, baseProps: Object): Object {
+  // TODO: Remove support for default props for everything except class
+  // components, including setting default props on a lazy wrapper around a
+  // class type.
+
   if (Component && Component.defaultProps) {
     // Resolve default props. Taken from ReactElement
     const props = assign({}, baseProps);