Neil.
· 约 1 分钟

React v17 源码系列：一文搞懂 useReducer

前瞻

本文将介绍 useReducer 的核心机制，在阅读之前，推荐先大致了解 React Reconcile 机制，想要从本文获益的话，至少应当理解 React Fiber, Work Loop, 单向、循环链表数据结构的概念，关于 React 源码解析可参考我的另一篇文章：React v17 源码解析

由于 React Hook 的复杂度较高，读一遍文章就想建立起结构化记忆是很困难的，所以我推荐先了解以下核心机制，待熟悉后，再参考我给出的完整代码，把以下概念串联起来：

  1. Mount 与 Update 的区别
  2. ReactCurrentDispatcher 运行时注入
  3. Hook 对象与两个单向链表：currentHook & workInProgressHook
  4. Fiber 与 Hook 的联系
  5. Update Queue 与两个循环链表：base & pending
  6. Lane 模型
  7. 更新触发：dispatch
  8. useReducerOnMount & useReducerOnUpdate

Mount VS Update

React Hook 在组件 Mount 和 Update 阶段的逻辑是不同的，我们暂且称之为 useReducerOnMount 和useReducerOnUpdate，为了避免一上来就 Deep Dive 吓退读者，让我们暂且忽略二者的细节吧，对这两个函数仅做简单介绍，后文我会给出完整代码以供参考。

Mount 阶段，useReducer 所做的事，相对来说很简单：

  1. 创建 Hook 对象
  2. 构建单向链表 workInProgressHook
  3. 计算初始状态
  4. bind dispatch 回调函数

Update 阶段，稍微复杂一些：

  1. 克隆 Hook 对象
  2. 构建单向链表 currentHook
  3. 构建单向链表 workInProgressHook
  4. 检查前后两次渲染 Hook 调用次数是否一致
  5. 处理 Update Queue

ReactCurrentDispatcher 运行时注入

讲完了 useReducerOnMount 和useReducerOnUpdate，引出我们的第一个问题，这二者是如何导出为一个 Hook 函数 useReducer 的呢？

答案就是 ReactCurrentDispatcher 机制。我们先来看下这段代码，这是一段省略了部分函数调用的 React useReducer 源码：

// Dispatcher 包含了所有 Hook 的实现细节
const ReactCurrentDispatcher: any = {
  current: null,
};

// 获取 Dispatcher
function resolveDispatcher() {
  return ReactCurrentDispatcher.current;
}

// 最终版
function useMockReducer(reducer, initialArgs, init?) {
  const dispatcher = resolveDispatcher();
  return dispatcher.useReducer(reducer, initialArgs, init);
}

有点感觉了吗？其实就是通过一个闭包全局变量 ReactCurrentDispatcher 实现的运行时注入。

具体来讲，在 reconcile 阶段，函数组件会通过 renderWithHook 函数进行 reconcile，ReactCurrentDispatcher.current 的引用替换就发生在此时此处：

// ...

if (current === null || current.memorizedState === null) {
  ReactCurrentDispatcher.current = HooksDispatcherOnMount;
} else {
  ReactCurrentDispatcher.current = HooksDispatcherOnUpdate;
}

// ...

const child = Component();

// ...

上述代码段执行后，ReactCurrentDispatcher.current 就指向了 Mount 或 Update 阶段 Hooks 的底层实现，调用过程：renderWithHook -> Component -> useReducer -> ReactCurrentDispatcher.current.useReducer

关于 reconcile 与 renderWithHook 可以查看我的另一篇文章：React v17 源码解析

Hook 对象与 currentHook & workInProgressHook

Fiber 与 Hook

Update Queue 与 queue.base & queue.pending

Lane 模型

触发组件更新：dispatchAction

Final Boss：useReducerOnMount & useReducerOnUpdate

完整 Mock 代码

// 当前正在渲染的 fiber 节点
let currentlyRenderingFiber;

// 当前正在处理的 Hook 节点
let workInProgressHook;

// 上次 render 后处理完毕的 Hook 节点
let currentHook;

// Hook 节点初始化
function mountWorkInProgressHook() {
  // 1. 初始化 hook 对象
  const hook = {
    memorizedState: null,
    next: null,
  };
  // 2. 维护 hooks 链表
  if (workInProgressHook === null) {
    // 2.1 首次执行，把当前 hook 节点，保存到当前正在渲染的 fiber 节点上
    currentlyRenderingFiber.memorizedState = hook;
  } else {
    // 2.2 非首次执行，维护 hooks 链表即可：当前 hook 节点 append 到 hooks 链表
    workInProgressHook.next = hook;
  }

  // 3. 将 hook 节点标记为正在处理
  workInProgressHook = hook;

  return hook;
}

// Hook 节点复用
function updateWorkInProgressHook() {
  // 1. 确定 currentHook 节点来源
  let nextCurrentHook;
  if (currentHook == null) {
    // 1.1 说明是首次 Update，直接复用 alternate fiber 上保存的 hook 节点（此处做了简化，实际还要判断 alternate 是否为空）
    nextCurrentHook = currentlyRenderingFiber.alternate.memorizedState;
  } else {
    // 1.2 说明不是首次 Update，那么直接复用 currentHook.next
    nextCurrentHook = currentHook.next;
  }

  // 2. 确定 workInProgressHook 节点来源
  let nextWorkInProgressHook;
  if (workInProgressHook == null) {
    // 2.1 说明这是组件中调用的首个 Hook（仔细体会跟 Case 1.1 的区别），直接复用 current fiber 上保存的 hook 节点（一般来说，就是 alternate 上的 hook 节点）
    nextWorkInProgressHook = currentlyRenderingFiber.memorizedState;
  } else {
    // 2.2 说明这并非组件中调用的首个 Hook，那么直接复用 workInProgressHook.next
    nextWorkInProgressHook = workInProgressHook.next;
  }

  // 3. 找到下个要处理的 Hook
  if (nextWorkInProgressHook !== null) {
    // 3.1 既然 nextWorkInProgress 有值，那就将 nextWorkInProgress 标记为处理中

    // 3.1.1 把 next 指向的 hook 标记为处理中
    workInProgressHook = nextWorkInProgressHook;
    // 3.1.2 再移动链表指针 next 到下个节点
    nextWorkInProgressHook = workInProgressHook.next;
  } else {
    // 3.2 说明 workInProgressHook 链表还没创建，此时克隆  currentHook 节点

    // 3.2.1 对于 Hook 来说，current 和 workInProgress 总是成对出现，此处既然 workInProgress 存在，
    // 那么 current 也必须存在，否则就说明这次将要处理的 hook 数量比上次 render 已处理完毕的 hook 数量要多，给出警告！
    if (nextCurrentHook === null) {
      throw new Error("Rendered more hooks than during the previous render.");
    }

    // 3.2.2 克隆 currentHook 节点
    const newHook = {
      memoizedState: currentHook.memoizedState,
      baseState: currentHook.baseState,
      baseQueue: currentHook.baseQueue,
      queue: currentHook.queue,
      next: null,
    };

    // 3.2.3 维护 Hook 链表结构
    if (workInProgressHook == null) {
      // 3.2.3.1 newHook 置为头结点
      currentlyRenderingFiber.memorizedState = newHook;
    } else {
      // 3.2.3.2 newHook 附加到链表尾
      workInProgressHook.next = newHook;
    }

    // 3.2.4 标记 newHook 为处理中
    workInProgressHook = newHook;
  }

  // 4. 移动 currentHook 链表指针到下个节点（此处简化，实际代码是在步骤 3 中完成的）
  currentHook = nextCurrentHook;

  // 5. 返回被标记为处理中的 hook 节点
  return workInProgressHook;
}

// 触发更新时会调用该函数，类组件的 setState 也会调用该函数，重新 render、commit
function scheduleUpdateOnFiber(fiber) {
  // ...
}

// 判断是否为交互事件中触发的更新
function isInterleavedUpdate(fiber): boolean {
  // ...
  return true;
}

// 交互事件中触发的更新会被维护在一个队列中
function pushInterleavedQueue(queue) {
  // ...
}

// dispatch 函数，每次执行生成一个 update 节点，这些节点会以循环链表形式，被维护在 Update Queue 的 pending 或 interleaved 属性上
function dispatchAction(fiber, queue, action) {
  // 1. 初始化 update 节点
  const update: any = {
    action,
    eagerReducer: null,
    eagerState: null,
    next: null,
  };

  // 2. 维护 updates 循环链表
  if (isInterleavedUpdate(fiber)) {
    // 2.1 交互事件中触发的更新，维护到 queue.interleaved，这是一个循环链表
    const interleaved = queue.interleaved;
    if (interleaved === null) {
      // This is the first update. Create a circular list.
      update.next = update;
      // At the end of the current render, this queue's interleaved updates will
      // be transfered to the pending queue.
      pushInterleavedQueue(queue);
    } else {
      update.next = interleaved.next;
      interleaved.next = update;
    }
    queue.interleaved = update;
  } else {
    // 2.2 普通更新，维护到 queue.pending，这是一个循环链表
    const pending = queue.pending;
    if (pending === null) {
      // This is the first update. Create a circular list.
      update.next = update;
    } else {
      update.next = pending.next;
      pending.next = update;
    }
    queue.pending = update;
  }

  // 3. 执行 reducer
  const lastRenderedReducer = queue.lastRenderedReducer;
  if (lastRenderedReducer !== null) {
    const currentState = queue.lastRenderedState;
    // 3.1 reducer 就是这里执行的
    const eagerState = lastRenderedReducer(currentState, action);
    // 3.2 初始化 update 节点的 eagerReducer 和 eagerState
    update.eagerReducer = lastRenderedReducer;
    update.eagerState = eagerState;
    // 3.3 优化：状态没有改变，立即返回，不会进入后续的 render、commit
    if (Object.is(currentState, eagerState)) {
      return;
    }
  }

  // 4. 更新状态，重新 render、commit
  scheduleUpdateOnFiber(fiber);
}

// Mount 时调用的 useReducer
function useReducerOnMount(reducer, initialArgs, init?) {
  // 1. hook 节点初始化
  const hook = mountWorkInProgressHook();
  // 2. 计算初始状态
  let memorizedState;
  if (init) {
    // 2.1 三参数用法
    memorizedState = init(initialArgs);
  } else {
    // 2.2 双参数用法，useState 本质就是执行的这条逻辑分支
    memorizedState = initialArgs;
  }
  hook.memorizedState = memorizedState;
  // 3. 初始化 Updates Queue
  const queue: any = {
    // 3.1 这个是普通 update 的循环链表
    pending: null,
    // 3.2 这个是交互事件触发的 update 的循环链表
    interleaved: null,
    // 3.3 这个就是传入的 reducer 函数
    lastRenderedReducer: reducer,
    // 3.4 这个就是当前状态
    lastRenderedState: memorizedState,
  };
  // 4. 初始化 dispatch 函数，dispatch 会保持对 currentlyRenderingFiber 和 queue 的引用
  queue.dispatch = dispatchAction.bind(null, currentlyRenderingFiber, queue);

  return [hook.memorizedState, queue.dispatch];
}

// Update 时调用的 useReducer
function useReducerOnUpdate(reducer, initialArgs, init?) {
  // 1. Hook 复用，无法复用时创建一个新的
  const hook = updateWorkInProgressHook();
  // 2. 更新 reducer
  hook.queue.lastRenderedReducer = reducer;

  // 3. 合并 baseQueue 和 pendingQueue
  if (hook.queue.pending !== null) {
    // ... 很复杂，考虑单独去写一篇《Update Queue 机制》
  }

  // 4. 遍历 baseQueue，进行状态计算
  if (currentHook.baseQueue !== null) {
    // ... 很复杂，考虑单独去写一篇《Update Queue 机制》
  }

  // 5. 返回的还是 Mount 时创建的 dispatch 回调，
  // 这也是为什么 setState/dispatch 可以安全地从 useEffect/useMemo/useCallback 的依赖列表中移除的原因
  return [hook.memorizedState, hook.queue.dispatch];
}

// Dispatcher 包含了所有 Hook 的实现细节
const ReactCurrentDispatcher: any = {
  current: null,
};

// 获取 Dispatcher
function resolveDispatcher() {
  return ReactCurrentDispatcher.current;
}

// 最终版
function useMockReducer(reducer, initialArgs, init?) {
  const dispatcher = resolveDispatcher();
  return dispatcher.useReducer(reducer, initialArgs, init);
}

// Hooks 宿主对象分为两种，一种是 Mount 时调用，一种 Update 时调用，他的作用就是统一函数名
const HooksDispatcherOnMount = {
  useReducer: useReducerOnMount,
  // ... other hooks implements
};
const HooksDispatcherOnUpdate = {
  useReducer: useReducerOnUpdate,
  // ... other hooks implements
};

// 当前正在处理的 work unit，这个概念不展开了，去看 React Work Loop 就明白了
let workInProgress;
// reconcile 函数组件时会调用该函数
function renderWithHooks(current: any, Component: any) {
  // ignore...

  // 1. 标记当前 work unit 为正在渲染的 fiber 节点，dispatchAction 将保持对该对象的引用（通过bind实现）
  currentlyRenderingFiber = workInProgress;
  // 2. ReactCurrentDispatcher 就在这里被赋值为 Hooks 宿主对象
  if (current === null || current.memorizedState === null) {
    ReactCurrentDispatcher.current = HooksDispatcherOnMount;
  } else {
    ReactCurrentDispatcher.current = HooksDispatcherOnUpdate;
  }

  // 3. 执行函数组件获取 React Element，我们的 useMockReducer 此时就会被执行
  const child = Component();

  // ignore...
}

// 调用起来和实际 useReducer 类似
function Foo() {
  // 首次调用 useMockReducer，将初始化一个 hooks 链表，新建 hook 节点作为头结点
  const [state, setState] = useMockReducer(
    (state, action) => {
      return {
        ...state,
        ...action,
      };
    },
    { count: 0 }
  );

  // 第二次调用 useMockReducer，再新建一个 hook 节点，添加到 hooks 链表尾
  const [state_, setState_] = useMockReducer(
    (state, action) => {
      return {
        ...state,
        ...action,
      };
    },
    { count: 0 }
  );

  // ...
}