React 实现 useMemo

本文将带大家实现 useMemo。

先看下如何使用。

function FunctionComponent() {
  const [count1, setCount1] = useState(1);
  const [count2, setCount2] = useState(1);
  
  // 昂贵运算
  const expensive = useMemo(() => {
      let sum = 0;
      for (let i = 0; i < count1; i++) {
          sum += i;
      }
      return sum
  }, [count1])

  return (
    <div className="border">
      <button
        onClick={() => {
          setCount1(count1 + 1);
        }}
      >
        {count1}
      </button>
      <button
        onClick={() => {
          setCount2(count2 + 1);
        }}
      >
        {count2}
      </button>
      <p>{expensive}</p>
    </div>
  );
}

ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render(
  (<FunctionComponent />) as any
);

实现 useMemo

// 检查 hook 依赖是否变化
export function areHookInputsEqual(
  nextDeps: Array<any>,
  prevDeps: Array<any> | null
): boolean {
  if (prevDeps === null) {
    return false;
  }

  for (let i = 0; i < prevDeps.length && i < nextDeps.length; i++) {
    if (Object.is(nextDeps[i], prevDeps[i])) {
      continue;
    }
    return false;
  }
  return true;
}

export function useMemo<T>(
  nextCreate: () => T,
  deps: Array<any> | void | null
): T {
  // 获取当前 hook
  const hook = updateWorkInProgressHook();

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

  const prevState = hook.memoizedState;
  // 检查依赖项是否发生变化
  if (prevState !== null) {
    if (nextDeps !== null) {
      const prevDeps = prevState[1];
      if (areHookInputsEqual(nextDeps, prevDeps)) {
        // 依赖项没有变化，返回上一次计算的结果，就是缓存的值
        return prevState[0];
      }
    }
  }

  const nextValue = nextCreate();

  // mount 阶段存储
  hook.memoizedState = [nextValue, nextDeps];

  return nextValue;
}