Fast-path SimpleMemoComponent with custom compare

Author: barryam3

packages/react-devtools-shared/src/__tests__/profilerChangeDescriptions-test.js

@@ -76,8 +76,7 @@ describe('Profiler change descriptions', () => {
         ▾ <App>
           ▾ <Context.Provider>
               <Child>
-            ▾ <Child> [Memo]
-                <Child>
+              <Child> [Memo]
             ▾ <RefForwardingComponent> [ForwardRef]
                 <Child>
     `);

packages/react-devtools-shared/src/__tests__/store-test.js

@@ -1893,6 +1893,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;
@@ -1922,6 +1927,7 @@ describe('Store', () => {
         <ForwardRefComponentWithAnonymousFunction />
         <ForwardRefComponentWithCustomDisplayName />
         <MemoComponent />
+        <MemoComponentWithCustomCompare />
         <MemoForwardRefComponent />
         <FakeHigherOrderComponent />
         <MemoizedFakeHigherOrderComponent />
@@ -1948,6 +1954,7 @@ describe('Store', () => {
               <MyComponent2>
             <Custom>
             <MyComponent4> [Memo]
+            <MyComponent5> [Memo]
           ▾ <MyComponent> [Memo]
               <MyComponent> [ForwardRef]
             <Baz> [withFoo][withBar]

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)