How does React.memo() work internally?

Episode 14 from React Internals Deep Dive, watch me explaining this on Youtube.

Based on [email protected], the implementation might have changed in newer version.

In How does React bailout work in reconciliation, we managed to understand how React stops reconciling when it finds there is nothing changed in the subtree.

Wait, sounds just like what React.memo() does, right?

1. Demo time
2. React.memo() creates a new element: REACT_MEMO_TYPE
3. reconciliation of MemoComponent
4. updateMemoComponent()
5. SimpleMemoComponent is the internal optimization
6. checkScheduledUpdateOrContext() is always called

This post is linked in Deeper Dive.

1. Demo time

Again let’s open the old demo, where C and D are both rerendered when the button is clicked.

Now let’s change the code a little bit, applying React.memo() to D.

Here is the new demo with memo, repeat the process, we can see now D is not rerendered, as we’d expected.

2. `React.memo()` creates a new element: REACT_MEMO_TYPE

First we set a debugger at where React.memo() is used, it leads us to the source code of memo()

It is pretty simple.

js
export function memo<Props>(
  type: React$ElementType,
  compare?: (oldProps: Props, newProps: Props) => boolean
) {
  const elementType = {
    $$typeof: REACT_MEMO_TYPE,
    type,
    compare: compare === undefined ? null : compare,
  };
  return elementType;
}

js
export function memo<Props>(
  type: React$ElementType,
  compare?: (oldProps: Props, newProps: Props) => boolean
) {
  const elementType = {
    $$typeof: REACT_MEMO_TYPE,
    type,
    compare: compare === undefined ? null : compare,
  };
  return elementType;
}

As the code explains itself, memo() create an element with

$$typeof to REACT_MEMO_TYPE
type set to our function passed in which is D()
a compare function.

TBH, I didn’t know that React.memo() accepts a second argument.

We now know that React.memo() creates a new fiber node that wraps things up. Looks straightforward, since if we are to add extra logic to avoid rerendering in D(), then we need some thing to wrap it up with that logic.

3. reconciliation of MemoComponent

Now again, let’s go to beginWork() as we explained in the last post, where fiber nodes are checked and updated.

We can easily find where REACT_MEMO_TYPE is used. source

js
case MemoComponent: {
  const type = workInProgress.type;
  const unresolvedProps = workInProgress.pendingProps;
  // Resolve outer props first, then resolve inner props.
  let resolvedProps = resolveDefaultProps(type, unresolvedProps);
  resolvedProps = resolveDefaultProps(type.type, resolvedProps);
  return updateMemoComponent(
    current,
    workInProgress,
    type,
    resolvedProps,
    renderLanes,
  );
}
case SimpleMemoComponent: {
  return updateSimpleMemoComponent(
    current,
    workInProgress,
    workInProgress.type,
    workInProgress.pendingProps,
    renderLanes,
  );
}

js
case MemoComponent: {
  const type = workInProgress.type;
  const unresolvedProps = workInProgress.pendingProps;
  // Resolve outer props first, then resolve inner props.
  let resolvedProps = resolveDefaultProps(type, unresolvedProps);
  resolvedProps = resolveDefaultProps(type.type, resolvedProps);
  return updateMemoComponent(
    current,
    workInProgress,
    type,
    resolvedProps,
    renderLanes,
  );
}
case SimpleMemoComponent: {
  return updateSimpleMemoComponent(
    current,
    workInProgress,
    workInProgress.type,
    workInProgress.pendingProps,
    renderLanes,
  );
}

There are actually two components, one is MemoComponent, and one is SimpleMemoComponent, we’ll understand why soon.

Hey, why is it MemoComponent, not REACT_MEMO_TYPE? REACT_MEMO_TYPE is $$typeof used for the element, MemoComponent is the function.

When the fiber is created from element, tag is set to MemoComponent here.

4. `updateMemoComponent()`

updateMemoComponent() is is the core of how React.memo() works. source

function updateMemoComponent(
  current: Fiber | null,
  workInProgress: Fiber,
  Component: any,
  nextProps: any,
  renderLanes: Lanes
): null | Fiber {
  if (current === null) {
if this is initial mount
    const type = Component.type;
    if (
      isSimpleFunctionComponent(type) &&
      Component.compare === null &&
      // SimpleMemoComponent codepath doesn't resolve outer props either.
      Component.defaultProps === undefined
    ) {
      let resolvedType = 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.
      workInProgress.tag = SimpleMemoComponent;
      workInProgress.type = resolvedType;
      return updateSimpleMemoComponent(
        current,
        workInProgress,
        resolvedType,
        nextProps,
        renderLanes
      );
go to an optimized branch
    }
    const child = createFiberFromTypeAndProps(
      Component.type,
      null,
      nextProps,
      workInProgress,
      workInProgress.mode,
      renderLanes
    );
    child.ref = workInProgress.ref;
    child.return = workInProgress;
    workInProgress.child = child;
    return child;
  }
Now below are re-render branch
  const currentChild = ((current.child: any): Fiber); // This is always exactly one child
  const hasScheduledUpdateOrContext = checkScheduledUpdateOrContext(
    current,
    renderLanes
  );
  if (!hasScheduledUpdateOrContext) {
    // 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) && current.ref === workInProgress.ref) {
      return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
    }
if props doesn't change they try early bailout
for bailout, refer to How does React bailout work in reconciliation?
  }
  // React DevTools reads this flag.
  workInProgress.flags |= PerformedWork;
  const newChild = createWorkInProgress(currentChild, nextProps);
  newChild.ref = workInProgress.ref;
  newChild.return = workInProgress;
  workInProgress.child = newChild;
  return newChild;
}

function updateMemoComponent(
  current: Fiber | null,
  workInProgress: Fiber,
  Component: any,
  nextProps: any,
  renderLanes: Lanes
): null | Fiber {
  if (current === null) {
if this is initial mount
    const type = Component.type;
    if (
      isSimpleFunctionComponent(type) &&
      Component.compare === null &&
      // SimpleMemoComponent codepath doesn't resolve outer props either.
      Component.defaultProps === undefined
    ) {
      let resolvedType = 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.
      workInProgress.tag = SimpleMemoComponent;
      workInProgress.type = resolvedType;
      return updateSimpleMemoComponent(
        current,
        workInProgress,
        resolvedType,
        nextProps,
        renderLanes
      );
go to an optimized branch
    }
    const child = createFiberFromTypeAndProps(
      Component.type,
      null,
      nextProps,
      workInProgress,
      workInProgress.mode,
      renderLanes
    );
    child.ref = workInProgress.ref;
    child.return = workInProgress;
    workInProgress.child = child;
    return child;
  }
Now below are re-render branch
  const currentChild = ((current.child: any): Fiber); // This is always exactly one child
  const hasScheduledUpdateOrContext = checkScheduledUpdateOrContext(
    current,
    renderLanes
  );
  if (!hasScheduledUpdateOrContext) {
    // 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) && current.ref === workInProgress.ref) {
      return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
    }
if props doesn't change they try early bailout
for bailout, refer to How does React bailout work in reconciliation?
  }
  // React DevTools reads this flag.
  workInProgress.flags |= PerformedWork;
  const newChild = createWorkInProgress(currentChild, nextProps);
  newChild.ref = workInProgress.ref;
  newChild.return = workInProgress;
  workInProgress.child = newChild;
  return newChild;
}

Despite the details, we can see the basic logic is

text
if (first time) {
  if (is simple) {
    update fiber tag to SimpleMemoComponent, so next time we go directly to updateSimpleMemoComponent()
    updateSimpleMemoComponent()
  } else {
    create child fibers
  }
} else {
  if (we can bailout) {
    try bailout
  } else {
    reconcile children
  }
}

text
if (first time) {
  if (is simple) {
    update fiber tag to SimpleMemoComponent, so next time we go directly to updateSimpleMemoComponent()
    updateSimpleMemoComponent()
  } else {
    create child fibers
  }
} else {
  if (we can bailout) {
    try bailout
  } else {
    reconcile children
  }
}

5. SimpleMemoComponent is the internal optimization

js
function updateSimpleMemoComponent(
  current: Fiber | null,
  workInProgress: Fiber,
  Component: any,
  nextProps: any,
  renderLanes: Lanes
): null | Fiber {
  if (current !== null) {
    const prevProps = current.memoizedProps;
    if (
      shallowEqual(prevProps, nextProps) &&
      current.ref === workInProgress.ref &&
      // Prevent bailout if the implementation changed due to hot reload.
      (__DEV__ ? workInProgress.type === current.type : true)
    ) {
      didReceiveUpdate = false;
      if (!checkScheduledUpdateOrContext(current, renderLanes)) {
        workInProgress.lanes = current.lanes;
        return bailoutOnAlreadyFinishedWork(
          current,
          workInProgress,
          renderLanes
        );
      } else if ((current.flags & ForceUpdateForLegacySuspense) !== NoFlags) {
        // This is a special case that only exists for legacy mode.
        // See https://github.com/facebook/react/pull/19216.
        didReceiveUpdate = true;
      }
    }
  }
  return updateFunctionComponent(
    current,
    workInProgress,
    Component,
    nextProps,
    renderLanes
  );
}

js
function updateSimpleMemoComponent(
  current: Fiber | null,
  workInProgress: Fiber,
  Component: any,
  nextProps: any,
  renderLanes: Lanes
): null | Fiber {
  if (current !== null) {
    const prevProps = current.memoizedProps;
    if (
      shallowEqual(prevProps, nextProps) &&
      current.ref === workInProgress.ref &&
      // Prevent bailout if the implementation changed due to hot reload.
      (__DEV__ ? workInProgress.type === current.type : true)
    ) {
      didReceiveUpdate = false;
      if (!checkScheduledUpdateOrContext(current, renderLanes)) {
        workInProgress.lanes = current.lanes;
        return bailoutOnAlreadyFinishedWork(
          current,
          workInProgress,
          renderLanes
        );
      } else if ((current.flags & ForceUpdateForLegacySuspense) !== NoFlags) {
        // This is a special case that only exists for legacy mode.
        // See https://github.com/facebook/react/pull/19216.
        didReceiveUpdate = true;
      }
    }
  }
  return updateFunctionComponent(
    current,
    workInProgress,
    Component,
    nextProps,
    renderLanes
  );
}

updateSimpleMemoComponent() is much simpler:

it just shallow compares the props and try to bailout
it doesn’t create new fiber, which means there is no D() in the fiber tree.

We can see the fiber tree from console. (watch my video here to see how)

Yeah, the memo element has type of D() but its child is going directly to div.

if we set pass down a compare function, then React couldn’t treat it as simple any more

js
const D = React.memo(_D, (a, b) => true);

js
const D = React.memo(_D, (a, b) => true);

Now if we inspect the fiber tree, things are different.

tag is 14, not simple now
type is not D any more
child is div but D

Yeah, in short, SimpleMemoComponent is an internal optimization for function components, which merges the memo fiber and wrapped fiber into one

It looks similar to View Flattening in React-native.

And the original pr for this optimization could be found here.

6. `checkScheduledUpdateOrContext()` is always called

Notice that though in updateMemoComponent() and updateSimpleMemoComponent(), the props are compared but checkScheduledUpdateOrContext() is always run, because the wrapped component could have being scheduled update by other functions, like some events or context.

So there are cases though props are equal but bailout does not happen.

Which means, as React hompeage says, React.memo() is only for performance optimization, Do not rely on it to “prevent” a render.

That’s it for React.memo(), hope it helps.

Want to know more about how React works internally?
Check out my series - React Internals Deep Dive!

How does React.memo() work internally?

1. Demo time

2. React.memo() creates a new element: REACT_MEMO_TYPE

3. reconciliation of MemoComponent

4. updateMemoComponent()

5. SimpleMemoComponent is the internal optimization

6. checkScheduledUpdateOrContext() is always called

2. `React.memo()` creates a new element: REACT_MEMO_TYPE

4. `updateMemoComponent()`

6. `checkScheduledUpdateOrContext()` is always called