Add fast path for React.memo with custom compare

[everettbu]

Mirror of facebook/react#34540
Original author: barryam3

---

Summary

Fixes react/react#33793

This widens the SimpleMemoComponent fast path introduced in react/react#13903 to support memoized function components with custom props comparisons. The fast path eliminates the need for the wrapper MemoComponent to be tracked as a separate Fiber, reducing performance overhead.

In addition, React DevTools will show only one node for the component instead of two (because there is only one Fiber instead of two), which makes debugging easier by reducing visual and therefore mental clutter.

I considered several approaches as discussed in the issue before settling on this one as the best combination of simplicity, safety, and performance after benchmarking and prototyping.

How did you test this change?

Per the contributing instructions:

1. yarn test - two tests fail ("ReactFlightAsyncDebugInfo › can track async information when awaited" and "ReactFlightDOMEdge › should execute repeated host components only once") but these also fail at main.
2. yarn test --prod - 59 test failures but again this is the same as main. Here they are if you want to see: https://gist.github.com/barryam3/82007cf5ef6b1c27621e0395ba9a1e71
3. yarn prettier
4. yarn lint
5. yarn flow dom-node (I first tried yarn flow as recommended by the PR template, but it told me to pick a renderer, and dom-node is a good default)

Additionally I updated the DevTools unit tests to reflect the component tree in DevTools.

1. yarn run build-for-devtools
2. yarn run test-build-devtools - two tests fail ("ignoreList source map extension › for dev builds › should not ignore list anything" and "ignoreList source map extension › for production builds › should include every source") but these also fail at main

I made the equivalent change in my own application (thousands of components, hundreds memoized with custom arePropsEqual) via a pnpm patch and have not discovered any issues after running in production for about 2 months. I verified that this solves the "extra node" problem in DevTools. (In the interest of sharing any possible caveats, my app is still on React 18.2.0 and I chose a slightly simpler implementation for the patch of just setting a compare field on the function instead of using a Weak Map.)

[greptile-apps]

Greptile Overview

Greptile Summary

This PR extends the SimpleMemoComponent fast path to support custom compare functions, eliminating the wrapper fiber for memoized function components. The implementation uses a WeakMap to store custom compare functions keyed by the component function itself.

Key Changes:

• Added simpleMemoComponentToCompare WeakMap to track custom compare functions for SimpleMemoComponent instances
• Modified updateMemoComponent to upgrade all simple function components to SimpleMemoComponent (previously only those without custom compare)
• Modified updateSimpleMemoComponent to retrieve and use the custom compare function from the WeakMap, falling back to shallowEqual
• Updated DevTools tests to reflect the new single-fiber structure for memo components

Benefits:

• Reduces fiber tree complexity by eliminating wrapper nodes for memoized components
• Improves DevTools clarity by showing one node instead of two
• Performance improvement by using fast path for more memo components

Potential Issue:

• WeakMap collision: if the same function component is memoized multiple times with different compare functions (e.g., React.memo(MyFunc, compare1) and React.memo(MyFunc, compare2)), they will share the same compare function. While this pattern is uncommon, it represents a correctness issue.

Confidence Score: 3/5

• This PR is generally safe but has a logical edge case with WeakMap collisions that could cause incorrect behavior
• The implementation is well-structured and the fast path optimization is sound. The WeakMap collision issue with same-function-different-compare scenarios is a correctness bug, though it's likely rare in practice. The PR author's production testing is encouraging, but the edge case should be addressed before merge.
• packages/react-reconciler/src/ReactFiberBeginWork.js - WeakMap key collision needs resolution

Important Files Changed

File Analysis

Filename	Score	Overview
packages/react-reconciler/src/ReactFiberBeginWork.js	3/5	Extended SimpleMemoComponent fast path to support custom compare functions using a WeakMap. Potential collision issue with same function/different compares.
packages/react-devtools-shared/src/tests/profilerChangeDescriptions-test.js	5/5	Updated test snapshot to reflect single fiber node for memo components (removed nested structure).
packages/react-devtools-shared/src/tests/store-test.js	5/5	Added test case for memo component with custom compare function, correctly showing single fiber node in DevTools.

[greptile-apps]

_{3 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: WeakMap key collision: if the same function component is wrapped with React.memo using different compare functions (e.g., const A = React.memo(MyFunc, compare1) and const B = React.memo(MyFunc, compare2)), the second call will overwrite the first in the WeakMap. This causes both components to use the same compare function instead of their respective ones.

This is likely rare in practice, but consider using a different key strategy or documenting this limitation.

Prompt To Fix With AI

This is a comment left during a code review.
Path: packages/react-reconciler/src/ReactFiberBeginWork.js
Line: 322:325

Comment:
**logic:** WeakMap key collision: if the same function component is wrapped with `React.memo` using different `compare` functions (e.g., `const A = React.memo(MyFunc, compare1)` and `const B = React.memo(MyFunc, compare2)`), the second call will overwrite the first in the WeakMap. This causes both components to use the same compare function instead of their respective ones.

This is likely rare in practice, but consider using a different key strategy or documenting this limitation.

How can I resolve this? If you propose a fix, please make it concise.

[github-actions]

This pull request has been automatically marked as stale. If this pull request is still relevant, please leave any comment (for example, "bump"), and we'll keep it open. We are sorry that we haven't been able to prioritize reviewing it yet. Your contribution is very much appreciated.

[github-actions]

Closing this pull request after a prolonged period of inactivity. If this issue is still present in the latest release, please ask for this pull request to be reopened. Thank you!