烧脑预警，这波心智负担有点重，深度探讨 useState 的实现原理

在前面的一篇文章中，我们介绍了 Fiber 的详细属性所代表的含义。在函数式组件中，其中与 hook 相关的属性为 memoizedState。

Fiber = {
  memoizedState: Hook
}

Fiber.memoizedState 是一个链表的起点，该链表的节点信息为。

export type Hook = {
  memoizedState: any,
  baseState: any,
  baseQueue: Update | null,
  queue: any,
  next: Hook | null,
}

useState 调用分为两个阶段，一个是初始化阶段，一个是更新阶段。当我们在 beginWork 中调用 renderWithHooks 时，通过判断 Fiber.memozedState 是否有值来分辨当前执行属于初始阶段还是更新阶段。

ReactCurrentDispatcher.current =
  current === null || current.memoizedState === null
    ? HooksDispatcherOnMount
    : HooksDispatcherOnUpdate;

在 react 模块中，我们可以看到 useState 的源码非常简单。

export function useState(
  initialState: (() => S) | S,
): [S, Dispatch] {
  const dispatcher = resolveDispatcher();
  return dispatcher.useState(initialState);
}

这里的 dispatcher，其实就是我们在 react-reconciler 中判断好的 ReactCurrentDispatcher.currenthook 的初始化方法挂载在 HooksDispatcherOnMount 上。

const HooksDispatcherOnMount: Dispatcher = {
  readContext,

  useCallback: mountCallback,
  useContext: readContext,
  useEffect: mountEffect,
  useImperativeHandle: mountImperativeHandle,
  useLayoutEffect: mountLayoutEffect,
  useInsertionEffect: mountInsertionEffect,
  useMemo: mountMemo,
  useReducer: mountReducer,
  useRef: mountRef,
  useState: mountState,
  useDebugValue: mountDebugValue,
  useDeferredValue: mountDeferredValue,
  useTransition: mountTransition,
  useMutableSource: mountMutableSource,
  useSyncExternalStore: mountSyncExternalStore,
  useId: mountId,

  unstable_isNewReconciler: enableNewReconciler,
};

hook 的更新方法挂载在 HooksDispatcherOnUpdate 上。

const HooksDispatcherOnUpdate: Dispatcher = {
  readContext,

  useCallback: updateCallback,
  useContext: readContext,
  useEffect: updateEffect,
  useImperativeHandle: updateImperativeHandle,
  useInsertionEffect: updateInsertionEffect,
  useLayoutEffect: updateLayoutEffect,
  useMemo: updateMemo,
  useReducer: updateReducer,
  useRef: updateRef,
  useState: updateState,
  useDebugValue: updateDebugValue,
  useDeferredValue: updateDeferredValue,
  useTransition: updateTransition,
  useMutableSource: updateMutableSource,
  useSyncExternalStore: updateSyncExternalStore,
  useId: updateId,

  unstable_isNewReconciler: enableNewReconciler,
};

因此，在初始化时，useState 调用的是 mountState，在更新时，useState 调用的是 updateState

一、mountState

mountState 的源码如下：

function mountState(
  initialState: (() => S) | S,
): [S, Dispatch] {
  const hook = mountWorkInProgressHook();
  if (typeof initialState === 'function') {
    initialState = initialState();
  }
  hook.memoizedState = hook.baseState = initialState;
  const queue: UpdateQueue = {
    pending: null,
    lanes: NoLanes,
    dispatch: null,
    lastRenderedReducer: basicStateReducer,
    lastRenderedState: (initialState: any),
  };
  hook.queue = queue;
  const dispatch: Dispatch = (queue.dispatch = (dispatchSetState.bind(
    null,
    currentlyRenderingFiber,
    queue,
  ): any));
  return [hook.memoizedState, dispatch];
}

理解这个源码的关键在第一行代码。

const hook = mountWorkInProgressHook();

react 在 ReactFiberHooks.new.js 模块全局中创建了如下三个变量。

let currentlyRenderingFiber: Fiber = (null: any);
let currentHook: Hook | null = null;
let workInProgressHook: Hook | null = null;

currentlyRenderingFiber 表示当前正在 render 中的 Fiber 节点。currentHook 表示当前 Fiber 的链表。

workInProgressHook 表示当前正在构建中的新链表。

mountWorkInProgressHook 方法会创建当前这个 mountState 执行所产生的 hook 链表节点。

function mountWorkInProgressHook(): Hook {
  const hook: Hook = {
    memoizedState: null,

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

    next: null,
  };

  if (workInProgressHook === null) {
    // 作为第一个节点
    currentlyRenderingFiber.memoizedState = workInProgressHook = hook;
  } else {
    // 添加到链表的下一个节点
    workInProgressHook = workInProgressHook.next = hook;
  }
  // 返回当前节点
  return workInProgressHook;
}

hook 节点的 queue 表示一个新的链表结构，用于存储针对同一个 state 的多次 update 操作。，.pending 指向下一个 update 链表节点。此时因为是初始化操作，因此值为 null，此时我们会先创建一个 queue。

const queue: UpdateQueue = {
  pending: null,
  lanes: NoLanes,
  dispatch: null,
  lastRenderedReducer: basicStateReducer,
  lastRenderedState: (initialState: any),
};
hook.queue = queue;

此时，dispatch 还没有赋值。在接下来我们调用了 dispatchSetState，我们待会儿来详细介绍这个方法，他会帮助 queue.pending 完善链表结构或者进入调度阶段，并返回了当前 hook 需要的 dispatch 方法。

const dispatch: Dispatch = (queue.dispatch = (dispatchSetState.bind(
  null,
  currentlyRenderingFiber,
  queue,
): any));

最后将初始化之后的缓存值和操作方法通过数组的方式返回。

return [hook.memoizedState, dispatch];

二、updateState

更新时，将会调用 updateState 方法，他的代码非常简单，就是直接调用了一下 updateReducer。

function updateState(
  initialState: (() => S) | S,
): [S, Dispatch] {
  return updateReducer(basicStateReducer, (initialState: any));
}

这里的需要注意的是有一个模块中的全局方法 basicStateReducer，该方法执行会结合传入的 action 返回最新的 state 值。

function basicStateReducer(state: S, action: BasicStateAction): S {
  // $FlowFixMe: Flow doesn't like mixed types
  return typeof action === 'function' ? action(state) : action;
}

代码中区分的情况是 useState 与 useReducer 的不同。useState 传入的是值，而 useReducer 传入的是函数

三、updateReducer

updateReducer 的代码量稍微多了一些，不过他的主要逻辑是计算出最新的 state 值。

当我们使用 setState 多次调用 dispatch 之后， 在 Hook 节点的 hook.queue 上会保存一个循环链表用于存储上一次的每次调用传入的 state 值，updateReducer 的主要逻辑就是遍历该循环链表，并计算出最新值。

此时首先会将 queue.pending 的链表赋值给 hook.baseQueue，然后置空 queue.pending。

const pendingQueue = queue.pending;
current.baseQueue = baseQueue = pendingQueue;
queue.pending = null;

然后通过 while 循环遍历 hook.baseQueue 通过 reducer 计算出最新的 state 值。

// 简化版代码
const first = baseQueue.next;

if (first !== null) {
  let newState = current.baseState;
  let update = first;
  do {
    // 执行每一次更新，去更新状态
    const action = update.action;
    newState = reducer(newState, action);
    update = update.next;
  } while (update !== null && update !== first);

  hook.memoizedState = newState;
}

最后再返回。

const dispatch: Dispatch = (queue.dispatch: any);
return [hook.memoizedState, dispatch];

四、dispatchSetState

当我们调用 setState 时，最终调用的是 dispatchSetState 方法。

setLoading -> dispatch -> dispatchSetState

该方法有两个逻辑，一个是同步调用，一个是并发模式下的异步调用。

同步调用时，主要的目的在于创建 hook.queue.pending 指向的环形链表。

首先我们要创建一个链表节点，该节点我们称之为 update。

const lane = requestUpdateLane(fiber);

const update: Update = {
  lane,
  action,
  hasEagerState: false,
  eagerState: null,
  next: (null: any),
};

然后会判断是否在 render 的时候调用了该方法。

if (isRenderPhaseUpdate(fiber)) {
  enqueueRenderPhaseUpdate(queue, update);
} else {

isRenderPhaseUpdate 用于判断当前是否是在 render 时调用，他的逻辑也非常简单。

function isRenderPhaseUpdate(fiber: Fiber) {
  const alternate = fiber.alternate;
  return (
    fiber === currentlyRenderingFiber ||
    (alternate !== null && alternate === currentlyRenderingFiber)
  );
}

这里需要重点关注是 enqueueRenderPhaseUpdate 是如何创建环形链表的。他的代码如下：

function enqueueRenderPhaseUpdate(
  queue: UpdateQueue,
  update: Update,
) {
  didScheduleRenderPhaseUpdateDuringThisPass = didScheduleRenderPhaseUpdate = true;
  const pending = queue.pending;
  if (pending === null) {
    update.next = update;
  } else {
    update.next = pending.next;
    pending.next = update;
  }
  queue.pending = update;
}

我们用图示来表达一下这个逻辑，光看代码可能理解起来比较困难。

当只有一个 update 节点时。

新增一个：

再新增一个：

在后续的逻辑中，会面临的一种情况是当渲染正在发生时，收到了来自并发事件的更新，我们需要等待直到当前渲染结束或中断再将其加入到 Fiber/Hook 队列。因此React 需要一个数组来存储这些更新，代码逻辑如下：

const concurrentQueues: Array = [];
let concurrentQueuesIndex = 0;

function enqueueUpdate(
  fiber: Fiber,
  queue: ConcurrentQueue | null,
  update: ConcurrentUpdate | null,
  lane: Lane,
) {
  concurrentQueues[concurrentQueuesIndex++] = fiber;
  concurrentQueues[concurrentQueuesIndex++] = queue;
  concurrentQueues[concurrentQueuesIndex++] = update;
  concurrentQueues[concurrentQueuesIndex++] = lane;

  concurrentlyUpdatedLanes = mergeLanes(concurrentlyUpdatedLanes, lane);

  fiber.lanes = mergeLanes(fiber.lanes, lane);
  const alternate = fiber.alternate;
  if (alternate !== null) {
    alternate.lanes = mergeLanes(alternate.lanes, lane);
  }
}

在这个基础之上，React 就有机会处理那些不会立即导致重新渲染的更新进入队列。如果后续有更高优先级的更新出现，将会重新对其进行排序。

export function enqueueConcurrentHookUpdateAndEagerlyBailout(
  fiber: Fiber,
  queue: HookQueue,
  update: HookUpdate,
): void {
  // This function is used to queue an update that doesn't need a rerender. The
  // only reason we queue it is in case there's a subsequent higher priority
  // update that causes it to be rebased.
  const lane = NoLane;
  const concurrentQueue: ConcurrentQueue = (queue: any);
  const concurrentUpdate: ConcurrentUpdate = (update: any);
  enqueueUpdate(fiber, concurrentQueue, concurrentUpdate, lane);
}

dispatchSetState 的逻辑中，符合条件就会执行该函数。

if (is(eagerState, currentState)) {
  // Fast path. We can bail out without scheduling React to re-render.
  // It's still possible that we'll need to rebase this update later,
  // if the component re-renders for a different reason and by that
  // time the reducer has changed.
  // TODO: Do we still need to entangle transitions in this case?
  enqueueConcurrentHookUpdateAndEagerlyBailout(fiber, queue, update);
  return;
}

很显然，这就是并发更新的逻辑，代码会最终调用 scheduleUpdateOnFiber，该方法是由 react-reconciler 提供，他后续会将任务带入到 scheduler 中调度。

// 与 enqueueConcurrentHookUpdateAndEagerlyBailout 方法逻辑
// 但会返回 root 节点
const root = enqueueConcurrentHookUpdate(fiber, queue, update, lane);

const eventTime = requestEventTime();
scheduleUpdateOnFiber(root, fiber, lane, eventTime);
entangleTransitionUpdate(root, queue, lane);

五、总结与思考

这就是 useState 的实现原理。其中包含了大量的逻辑操作，可能跟我们在使用时所想的那样有点不太一样。这里大量借助了闭包和链表结构来完成整个构想。

这个逻辑里面也会有大量的探讨存在于大厂面试的过程中。例如

• 为什么不能把 hook 的写法放到 if 判断中去。
• setState 的合并操作是如何做到的。
• hook 链表和 queue.pending 的环状链表都应该如何理解？
• setState 之后，为什么无法直接拿到最新值，彻底消化了之后这些问题都能很好的得到解答。