React Hooks 源码解读之 useMemo

「这是我参与11月更文挑战的第3天，活动详情查看：2021最后一次更文挑战」

react 版本：v17.0.3

1、Hook 入口

在 React Hooks 源码解读之Hook入口 一文中，我们介绍了 Hooks 的入口及hook处理函数的挂载，从 hook 处理函数的挂载关系我们可以得到这样的等式：

• 挂载阶段：

  useMemo = ReactCurrentDispatcher.current.useMemo = HooksDispatcherOnMount.useMemo = mountMemo;

• 更新阶段：

  useMemo = ReactCurrentDispatcher.current.useMemo = HooksDispatcherOnUpdate.useMemo = updateMemo;

因此，组件在挂载阶段，执行 useMemo，其实执行的是 mountMemo，而更新阶段，则执行的是updateMemo 。

2、挂载阶段

组件在挂载阶段，执行 useMemo，实际上执行的是 mountMemo，下面我们来看看 mountMemo 的实现：

2.1 mountMemo

// packages/react-reconciler/src/ReactFiberHooks.new.js

function mountMemo<T>(
  nextCreate: () => T,  // useMemo 的第一个参数：callback
  deps: Array<mixed> | void | null, // useMemo 的第二个参数：依赖项数组
): T {
  // 创建 hook 对象，将 hook 对象添加到 workInProgressHook 单向链表中，返回最新的 hook 链表
  const hook = mountWorkInProgressHook();
  // 根据 useMemo 的第二个参数初始化 依赖项数组
  const nextDeps = deps === undefined ? null : deps;
  // 执行 useMemo 的第一个参数 callback，计算需要缓存的值
  const nextValue = nextCreate();
  // 将计算后的值和相应的依赖项保存到 hook 对象的 memoizedState 属性上
  hook.memoizedState = [nextValue, nextDeps];
  // 返回缓存后的值
  return nextValue;
}

组件在挂载时，执行 useMemo，首先会创建一个 hook 对象，该 hook 对象将会被添加到 workInProgressHook 单向链表中：

const hook = mountWorkInProgressHook();

接着根据 useMemo 的第二个参数 deps 来初始化需要保存起来的 依赖项数组：

const nextDeps = deps === undefined ? null : deps;

然后执行 useMemo 的第一个参数 callback，计算需要缓存起来的值：

 const nextValue = nextCreate();

最后将计算的值和相应的依赖项保存到 hook 对象的 memoizedState 属性上，并返回缓存后的值：

// 将计算后的值和相应的依赖项保存到 hook 对象的 memoizedState 属性上
hook.memoizedState = [nextValue, nextDeps];
// 返回缓存后的值
return nextValue;

由此可知道，useMemo 的 memoizedState 存储的并不是一个具体的值，而是一个包含有缓存值和依赖项的数组。

2.2 mountWorkInProgressHook

在 mountMemo() 函数中，使用 mountWorkInProgressHook() 函数创建了一个新的 hook 对象，我们来看看它是如何被创建的：

// packages/react-reconciler/src/ReactFiberHooks.new.js

// 创建一个新的 hook 对象，并返回当前的 workInProgressHook 对象
// workInProgressHook 对象是全局对象，在 mountWorkInProgressHook 中首次初始化
function mountWorkInProgressHook(): Hook {
  const hook: Hook = {
    memoizedState: null,

    baseState: null,
    baseQueue: null,
    queue: null,

    next: null,
  };
  
   // Hooks are stored as a linked list on the fiber's memoizedState field
  // 将 新建的 hook 对象以链表的形式存储在当前的 fiber 节点memoizedState属性上

  // 只有在第一次打开页面的时候，workInProgressHook 为空
  if (workInProgressHook === null) {
    // This is the first hook in the list
    // 链表上的第一个 hook
    
    // currentlyRenderingFiber： The work-in-progress fiber. I've named it differently to distinguish it fromthe work-in-progress hook.
    
    currentlyRenderingFiber.memoizedState = workInProgressHook = hook;
  } else {
    // Append to the end of the list
    // 已经存在 workInProgressHook 对象，则将新创建的这个 Hook 接在 workInProgressHook 的尾部，形成链表
    workInProgressHook = workInProgressHook.next = hook;
  }
  return workInProgressHook;
}

可以看到，在新建一个 hook 对象时，如果全局的 workInProgressHook 对象不存在 (值为 null)，即组件在首次渲染时，将新建的 hook 对象赋值给 workInProgressHook 对象，也同时将 hook 对象赋值给 currentlyRenderingFiber 的 memoizedState 属性，如果 workInProgressHook 不为 null，则将 hook 对象接在 workInProgressHook 的尾部，从而形成一个单向链表。

3、更新阶段

组件在更新阶段，执行 useMemo，实际上执行的是 updateMemo，下面我们来看看 updateMemo 的实现。

3.1 updateMemo

// packages/react-reconciler/src/ReactFiberHooks.new.js

function updateMemo<T>(
  nextCreate: () => T, // useMemo 的第一个参数：callback
  deps: Array<mixed> | void | null, // useMemo 的第二个参数：依赖项数组
): T {
  // 找到该useMemo对应的hook 对象
  const hook = updateWorkInProgressHook();
  // 根据 useMemo 的第二个参数计算新的 依赖
  const nextDeps = deps === undefined ? null : deps;
  // 之前的 缓存值和依赖项 [nextValue, nextDeps]
  const prevState = hook.memoizedState;
  if (prevState !== null) {
    // Assume these are defined. If they're not, areHookInputsEqual will warn.
    if (nextDeps !== null) {
      const prevDeps: Array<mixed> | null = prevState[1];
      // 比较前后的依赖是否相同
      if (areHookInputsEqual(nextDeps, prevDeps)) {
        // 前后的依赖相同，说明不需要重新执行 nextCreate 计算新值，而是直接返回上一次计算的值
        return prevState[0];
      }
    }
  }

  // 前后的依赖数组不同或者没有提供依赖项数组，useMemo 在每次渲染时都会计算新的值

  // 执行 useMemo 的第一个参数 callback，计算新的值
  const nextValue = nextCreate();
  // 将计算后的新值和相应的依赖项保存到 hook 对象的 memoizedState 属性上
  hook.memoizedState = [nextValue, nextDeps];
  // // 返回计算后的新值
  return nextValue;
}

在更新阶段，执行 useMemo，首先会调用 updateWorkInProgressHook() ，获取当前正在执行 update 任务的fiber 节点上的hook 对象：

const hook = updateWorkInProgressHook();

接着获取新的依赖，方便与旧的依赖做比较，从而决定是否需要重新执行 useMemo 的第一个参数，计算新的值：

const nextDeps = deps === undefined ? null : deps;

然后从当前的 hook 对象中获取之前的依赖项数组，调用 areHookInputsEqual() 来比较前后的依赖项数组是否相同：

// 之前的 缓存值和依赖项 [nextValue, nextDeps]
const prevState = hook.memoizedState;
if (prevState !== null) {
  // Assume these are defined. If they're not, areHookInputsEqual will warn.
  if (nextDeps !== null) {
    const prevDeps: Array<mixed> | null = prevState[1];
    // 比较前后的依赖是否相同
    if (areHookInputsEqual(nextDeps, prevDeps)) {
      // 前后的依赖相同，说明不需要重新执行 nextCreate 计算新值，而是直接返回上一次计算的值
      return prevState[0];
    }
  }
}

如果 areHookInputsEqual() 返回的结果为 true，说明前后的 依赖项数组 是一样的，说明不需要重新执行 nextCreate 计算新值，而是直接返回上一次计算的值。

如果 areHookInputsEqual() 返回的结果为 false，则会执行下面的语句：

// 前后的依赖数组不同或者没有提供依赖项数组，useMemo 在每次渲染时都会计算新的值

  // 执行 useMemo 的第一个参数 callback，计算新的值
  const nextValue = nextCreate();
  // 将计算后的新值和相应的依赖项保存到 hook 对象的 memoizedState 属性上
  hook.memoizedState = [nextValue, nextDeps];
  // // 返回计算后的新值
  return nextValue;

只要是前后的依赖项数组不同或者没有提供依赖项数组，useMemo 在每次渲染时都会执行上面的代码计算新值。因此，我们在使用 useMemo 的时候，应该传入依赖项数组，有助于避免每次渲染时都进行高开销的计算。

3.2 updateWorkInProgressHook

在 updateMemo() 函数中，通过 updateWorkInProgressHook() 函数获取到了当前正在工作中的 Hook，即 workInProgressHook，我们来看看 updateWorkInProgressHook 的实现：

// packages/react-reconciler/src/ReactFiberHooks.new.js

function updateWorkInProgressHook(): Hook {
  // This function is used both for updates and for re-renders triggered by a
  // render phase update. It assumes there is either a current hook we can
  // clone, or a work-in-progress hook from a previous render pass that we can
  // use as a base. When we reach the end of the base list, we must switch to
  // the dispatcher used for mounts.

  // 获取 当前 hook 的下一个 hook
  let nextCurrentHook: null | Hook;
  if (currentHook === null) {
    const current = currentlyRenderingFiber.alternate;
    if (current !== null) {
      nextCurrentHook = current.memoizedState;
    } else {
      nextCurrentHook = null;
    }
  } else {
    nextCurrentHook = currentHook.next;
  }

  // 取下一个 hook 为当前的hook
  let nextWorkInProgressHook: null | Hook;
  if (workInProgressHook === null) {
    nextWorkInProgressHook = currentlyRenderingFiber.memoizedState;
  } else {
    nextWorkInProgressHook = workInProgressHook.next;
  }

  if (nextWorkInProgressHook !== null) {
    // There's already a work-in-progress. Reuse it.
    workInProgressHook = nextWorkInProgressHook;
    nextWorkInProgressHook = workInProgressHook.next;

    currentHook = nextCurrentHook;
  } else {
    // Clone from the current hook.

    // 拷贝当前的 hook，作为当前正在工作中的 workInProgressHook

    invariant(
      nextCurrentHook !== null,
      'Rendered more hooks than during the previous render.',
    );
    currentHook = nextCurrentHook;

    const newHook: Hook = {
      memoizedState: currentHook.memoizedState,

      baseState: currentHook.baseState,
      baseQueue: currentHook.baseQueue,
      queue: currentHook.queue,

      next: null,
    };

    if (workInProgressHook === null) {
      // This is the first hook in the list.
      currentlyRenderingFiber.memoizedState = workInProgressHook = newHook;
    } else {
      // Append to the end of the list.
      workInProgressHook = workInProgressHook.next = newHook;
    }
  }
  return workInProgressHook;
}

这里分两种情况：

• 如果是在 render 阶段，则会取下一个 hook 作为当前的hook，并返回 workInProgressHook；
• 如果是在 re-render 阶段，则在当前处理周期中，继续取当前的 workInProgressHook 做更新处理，最后再返回 workInProgressHook。

前后「依赖项数组」是否一样，是通过 areHookInputsEqual() 来判断的，我们来看看它的实现。

3.3 areHookInputsEqual

// packages/react-reconciler/src/ReactFiberHooks.new.js

function areHookInputsEqual(
  nextDeps: Array<mixed>,
  prevDeps: Array<mixed> | null,
) {
    
  // 删除了 dev 代码
    
  if (prevDeps === null) {
    
    // 删除了 dev 代码
    
    return false;
  }

  // 删除了 dev 代码
    
  // deps 是一个 Array，循环遍历去比较 array 中的每个 item
  for (let i = 0; i < prevDeps.length && i < nextDeps.length; i++) {
    // is比较函数是浅比较
    if (is(nextDeps[i], prevDeps[i])) {
      continue;
    }
    return false;
  }
  return true;
}

在 areHookInputsEqual() 函数中，循环遍历 deps ，调用 is 方法去比较依赖项数组中的每个依赖，值得注意的是，is 方法是浅比较，也就是说如果是深比较那一定会更新的。下面是 is 方法的源码：

// packages/shared/objectIs.js

function is(x: any, y: any) {
  return (
    (x === y && (x !== 0 || 1 / x === 1 / y)) || (x !== x && y !== y) // eslint-disable-line no-self-compare
  );
}

const objectIs: (x: any, y: any) => boolean =
  typeof Object.is === 'function' ? Object.is : is;

export default objectIs;

Object.is() 方法判断两个值是否为同一个值，若当前浏览器支持该方法，则调用该方法来判断前后两个依赖项是否相同，若不支持，则调用React自己实现 is 方法来比较。

4、useMemo 流程图

5、总结

我们知道 useMemo 将会返回计算后的 memoized 值，它仅会在某个依赖项改变时才重新计算 memoized 值。这种优化有助于避免在每次渲染时都进行高开销的计算。虽然 useMemo 可以对数据进行缓存，但是也不能因此而滥用，我们应该考虑哪些数据值得去缓存，因为对于 useMemo 来说，除了计算值耗时以外，对比依赖项的比较也是需要时间的，我们应该对此进行衡量，才能够更好的去用 hook，而不是为了优化而优化。