[React 源码] useMemo | useCallback [1.2k 字 - 阅读时长2min]

useMemo 原理

问题: 下面这段代码，从挂载到更新发生了什么？怎么挂载的？怎么更新的。

const reducer = (state, action) => {
  if (action.type === "add") return state + 1;
  else return state;
};

function Counter() {
    const [state, setState] = useState(0);
    const memoFn = useMemo(() => state ** 2, [state]);
  return (
    <div
      onClick={() => {
      setState((pre) => pre + 1)
      }}
    >
      {number}
    </div>
  );
}

这里，我们直接进入到 reconciler 阶段，默认已经通过深度优先调度到了 Counter 函数组件的 Fiber节点

---

useMemo mount 挂载阶段

第一：判断是函数节点的 tag 之后，调用 renderWithHooks.

/* 
workInProgress: 当前工作的 Fiber 节点
Componet：Counter 函数组件
_current: 老 Fiber 节点 也就是 workInProgress.alternate 
*/
let value = renderWithHooks(_current,workInProgress,Component);

第二：在 renderWithHooks 当中调用 Counter 函数

let children = Component();

第三： 调用 Counter 函数 的 useMemo 函数

export function useMemo(reducer, initialArg) {
  return ReactCurrentDispatcher.current.useMemo(create, deps);
}

第四：挂载阶段 ReactCurrentDispatcher.current.useMemo 实则是调用了 mountMemo,

第五：在 mountMemo 中调用，mountWorkInProgressHook 函数，创建 useMemo 的 Hook 对象，构建 fiber.memoizedState 也就是 Hook 链表。将我们传入的函数进行计算作为 nextValue，并且返回。组件当中可以消费计算值。

function mountMemo<T>(
  nextCreate: () => T,
  deps: Array<mixed> | void | null,
): T {
  const hook = mountWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  const nextValue = nextCreate();
  hook.memoizedState = [nextValue, nextDeps];
  return nextValue;
}

如果之后再有 useState useReducer，最终mout阶段的成果是

自此 useMemo 挂载阶段执行完毕

---

useMemo update 更新阶段

第一：更新阶段 ReactCurrentDispatcher.current.useMemo 实则是调用了 updateMemo函数

第二：在 updateMemo 中调用 updateWorkInProgressHook 函数，在此函数中最重要的就是通过 alternate 指针复用 currentFiber（老 Fiber） 的 memorizedState, 也就是 Hook 链表，并且按照严格的对应顺序来复用 currentFiber（老 Fiber） Hook 链表当中的 Hook（通过 currentHook 指针结合链表来实现)。

第三：通过 areHookInputsEqual 函数，来判断 deps 依赖是否发送了变化，如果没有发送变化，直接返回缓存的 memo 值，如果发生了变化，调用 nextCreate 函数，重新计算 memo 值，并返回。

function updateMemo<T>(
  nextCreate: () => T,
  deps: Array<mixed> | void | null,
): T {
  const hook = updateWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  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)) {
        return prevState[0];
      }
    }
  }
  if (shouldDoubleInvokeUserFnsInHooksDEV) {
    nextCreate();
  }
  const nextValue = nextCreate();
  hook.memoizedState = [nextValue, nextDeps];
  return nextValue;
}

areHookInputsEqual 函数，就是遍历新，旧依赖。通过 Obejct.is 来判断依赖是否相等。

function areHookInputsEqual(
  nextDeps: Array<mixed>,
  prevDeps: Array<mixed> | null
): boolean {
  for (let i = 0; i < prevDeps.length && i < nextDeps.length; i++) {
    if (is(nextDeps[i], prevDeps[i])) {
      continue;
    }
    return false;
  }

  return true;
}

自此 useMemo 更新完毕

---

useCallback 原理

问题: 下面这段代码，从挂载到更新发生了什么？怎么挂载的？怎么更新的。

const reducer = (state, action) => {
  if (action.type === "add") return state + 1;
  else return state;
};

function Counter() {
    const [state, setState] = useState(0);
    const memoFn = useCallback(() => state ** 2, [state]);
  return (
    <div
      onClick={() => {
      setState((pre) => pre + 1)
      }}
    >
      {number}
    </div>
  );
}

这里，我们直接进入到 reconciler 阶段，默认已经通过深度优先调度到了 Counter 函数组件的 Fiber节点

---

useCallback mount 挂载阶段

第一：判断是函数节点的 tag 之后，调用 renderWithHooks.

/* 
workInProgress: 当前工作的 Fiber 节点
Componet：Counter 函数组件
_current: 老 Fiber 节点 也就是 workInProgress.alternate 
*/
let value = renderWithHooks(_current,workInProgress,Component);

第二：在 renderWithHooks 当中调用 Counter 函数

let children = Component();

第三： 调用 Counter 函数 的 useCallback 函数

export function useMemo(reducer, initialArg) {
  return ReactCurrentDispatcher.current.useCallback(create, deps);
}

第四：挂载阶段 ReactCurrentDispatcher.current.useMemo 实则是调用了 mountCallback,

第五：在 mountCallback 中调用，mountWorkInProgressHook 函数，创建 useCallback 的 Hook 对象，构建 fiber.memoizedState 也就是 Hook 链表。将我们传入的函数返回。组件当中可以消费该函数。

function mountCallback<T>(callback: T, deps: Array<mixed> | void | null): T {
  const hook = mountWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  hook.memoizedState = [callback, nextDeps];
  return callback;
}

如果之后再有 useState useReducer，最终mout阶段的成果是

自此 useMemo 挂载阶段执行完毕

---

useCallback update 更新阶段

第一：更新阶段 ReactCurrentDispatcher.current.useMemo 实则是调用了 updateMemo函数

第二：在 updateMemo 中调用 updateWorkInProgressHook 函数，在此函数中最重要的就是通过 alternate 指针复用 currentFiber（老 Fiber） 的 memorizedState, 也就是 Hook 链表，并且按照严格的对应顺序来复用 currentFiber（老 Fiber） Hook 链表当中的 Hook（通过 currentHook 指针结合链表来实现)。

第三：通过 areHookInputsEqual 函数，来判断 deps 依赖是否发送了变化，如果依赖没有发送变化，直接返回缓存的 函数，如果依赖发生了变化，则返回新的函数。

function updateCallback<T>(callback: T, deps: Array<mixed> | void | null): T {
  const hook = updateWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  const prevState = hook.memoizedState;
  if (prevState !== null) {
    if (nextDeps !== null) {
      const prevDeps: Array<mixed> | null = prevState[1];
      if (areHookInputsEqual(nextDeps, prevDeps)) {
        return prevState[0];
      }
    }
  }
  hook.memoizedState = [callback, nextDeps];
  return callback;
}

areHookInputsEqual 函数，就是遍历新，旧依赖。通过 Obejct.is 来判断依赖是否相等。

function areHookInputsEqual(
  nextDeps: Array<mixed>,
  prevDeps: Array<mixed> | null
): boolean {
  for (let i = 0; i < prevDeps.length && i < nextDeps.length; i++) {
    if (is(nextDeps[i], prevDeps[i])) {
      continue;
    }
    return false;
  }

  return true;
}

自此 useCallback 更新完毕

---