diff --git a/packages/react-devtools-shared/src/__tests__/profilerChangeDescriptions-test.js b/packages/react-devtools-shared/src/__tests__/profilerChangeDescriptions-test.js index 87d8132e50e..d650f82f423 100644 --- a/packages/react-devtools-shared/src/__tests__/profilerChangeDescriptions-test.js +++ b/packages/react-devtools-shared/src/__tests__/profilerChangeDescriptions-test.js @@ -76,8 +76,7 @@ describe('Profiler change descriptions', () => { ▾ - ▾ [Memo] - + [Memo] ▾ [ForwardRef] `); diff --git a/packages/react-devtools-shared/src/__tests__/store-test.js b/packages/react-devtools-shared/src/__tests__/store-test.js index 07e82040497..d1e6891712b 100644 --- a/packages/react-devtools-shared/src/__tests__/store-test.js +++ b/packages/react-devtools-shared/src/__tests__/store-test.js @@ -1887,6 +1887,11 @@ describe('Store', () => { ForwardRefComponentWithCustomDisplayName.displayName = 'Custom'; const MyComponent4 = (props, ref) => null; const MemoComponent = React.memo(MyComponent4); + const MyComponent5 = (props, ref) => null; + const MemoComponentWithCustomCompare = React.memo( + MyComponent5, + (a, b) => a === b, + ); const MemoForwardRefComponent = React.memo(ForwardRefComponent); const FakeHigherOrderComponent = () => null; @@ -1916,6 +1921,7 @@ describe('Store', () => { + @@ -1942,6 +1948,7 @@ describe('Store', () => { [Memo] + [Memo] ▾ [Memo] [ForwardRef] [withFoo][withBar] diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js index 33fadfd5396..c26231aeb9f 100644 --- a/packages/react-reconciler/src/ReactFiberBeginWork.js +++ b/packages/react-reconciler/src/ReactFiberBeginWork.js @@ -318,6 +318,12 @@ export const SelectiveHydrationException: mixed = new Error( let didReceiveUpdate: boolean = false; +// Map from SimpleMemoComponent function to compare a.k.a. arePropsEqual function. +const simpleMemoComponentToCompare = new WeakMap< + Function, + (objA: mixed, objB: mixed) => boolean, +>(); + let didWarnAboutBadClass; let didWarnAboutContextTypeOnFunctionComponent; let didWarnAboutContextTypes; @@ -478,16 +484,18 @@ function updateMemoComponent( ): null | Fiber { if (current === null) { const type = Component.type; - if (isSimpleFunctionComponent(type) && Component.compare === null) { + if (isSimpleFunctionComponent(type)) { let resolvedType = type; if (__DEV__) { resolvedType = resolveFunctionForHotReloading(type); } // If this is a plain function component without default props, - // and with only the default shallow comparison, we upgrade it - // to a SimpleMemoComponent to allow fast path updates. + // we upgrade it to a SimpleMemoComponent to allow fast path updates. workInProgress.tag = SimpleMemoComponent; workInProgress.type = resolvedType; + if (Component.compare) { + simpleMemoComponentToCompare.set(resolvedType, Component.compare); + } if (__DEV__) { validateFunctionComponentInDev(workInProgress, type); } @@ -548,9 +556,10 @@ function updateSimpleMemoComponent( // hasn't yet mounted. This happens when the inner render suspends. // We'll need to figure out if this is fine or can cause issues. if (current !== null) { + const compare = simpleMemoComponentToCompare.get(Component) ?? shallowEqual; const prevProps = current.memoizedProps; if ( - shallowEqual(prevProps, nextProps) && + compare(prevProps, nextProps) && current.ref === workInProgress.ref && // Prevent bailout if the implementation changed due to hot reload. (__DEV__ ? workInProgress.type === current.type : true)