useState和useReducer源码浅析

[zhangyu1818]

以下源码浅析的React版本为17.0.1，使用ReactDOM.render创建的同步应用，不含优先级相关。

流程简述

函数组件会调用renderWithHooks函数，这个函数主要会标记当前渲染的currentlyRenderingFiber节点，并判断该使用哪一个HooksDispatcher（React里Mount和Update所使用的Hooks不是同一个），接着执行此函数组件，分别处理hook函数，并得到children，将children返回到上层函数后，执行reconcileChildren生成child子节点。

renderWithHooks

renderWithHooks函数为函数组件的入口函数，无论是Mount时的mountIndeterminateComponent还是Update时的updateFunctionComponent都会进入这个函数来获取函数组件的children。

export function renderWithHooks<Props, SecondArg>(
  current: Fiber | null,
  workInProgress: Fiber,
  Component: (p: Props, arg: SecondArg) => any,
  props: Props,
  secondArg: SecondArg,
  nextRenderLanes: Lanes,
): any {
  // 省略代码...
  currentlyRenderingFiber = workInProgress;

  // 使用mount还是update的hook
  ReactCurrentDispatcher.current =
    current === null || current.memoizedState === null
      ? HooksDispatcherOnMount
      : HooksDispatcherOnUpdate;
  // 省略代码...
  
  // 执行函数后返回的children
  let children = Component(props, secondArg);

  // 省略代码...
  ReactCurrentDispatcher.current = ContextOnlyDispatcher;
  // 省略代码...
  
  return children;
}

renderWithHooks中会将当前的workInProgressFiber节点存在全局变量currentlyRenderingFiber中，这样方便后面获取Fiber信息，接着在调用函数获取children之前，先判断使用哪一个ReactCurrentDispatcher，最后返回children给上层函数处理。

需要注意的是Mount时和Update时用的不是同一个hook。

const HooksDispatcherOnMount: Dispatcher = {
  // ...
  useReducer: mountReducer,
  useState: mountState,
};
const HooksDispatcherOnUpdate: Dispatcher = {
  // ...
  useReducer: updateReducer,
  useState: updateState,
};

再来看下useState

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

// 返回的是调用函数组件之前赋值的ReactCurrentDispatcher.current
function resolveDispatcher() {
  const dispatcher = ReactCurrentDispatcher.current;
  return dispatcher;
}

所以后续会分为Mount和Update来分析

Hooks的数据结构

单个hook的数据结构如下

type Hook = {|
  memoizedState: any, // 对于useState和useReducer的值就是state的值
  baseState: any,
  baseQueue: Update<any, any> | null,
  queue: UpdateQueue<any, any> | null, // 更新队列
  next: Hook | null, // 下一个hook节点
|}

那函数组件如何找到对应的hooks信息呢？函数组件的hooks的信息储存在对应Fiber节点中的memoizedState字段中，代表函数组件里的第一个hook，hooks数据结构为单向链表，每一个节点可以通过next属性找到下一个hook。

function Container(){
  React.useState() // currentlyRenderingFiber.memoizedState (数据结构即是上面的Hook type)
  React.useState() // currentlyRenderingFiber.memoizedState.next (Hook.next)
 // ...
}

每一个hook的更新队列都会存在queue字段里，通过执行队列里的操作就可以得到最新的state值，结构如下

type UpdateQueue<S, A> = {|
  pending: Update<S, A> | null, // 最新的更新任务
  dispatch: (A => mixed) | null, // 也就是useState的第二个参数用来发起更新
  lastRenderedReducer: ((S, A) => S) | null, // 上一次的reducer
  lastRenderedState: S | null, // 上次的state值
|};

type Update<S, A> = {|
  lane: Lane,
  action: A,
  eagerReducer: ((S, A) => S) | null,
  eagerState: S | null, // 第一次调用dispatch时赋值
  next: Update<S, A>, // 下一个更新
  priority?: ReactPriorityLevel,
|};

hook的更新队列是一个单向环形链表，pending字段保存的是最新的update，通过update.next可以获取第一个加入更新队列的update。

update2(pending) -> update0 -> update1 -> update2

useState和useReducer

useState实际是一个自带了reducer的useReducer语法糖，所以需要放在一起分析。

mountState

function mountState<S>(
  initialState: (() => S) | S,
): [S, Dispatch<BasicStateAction<S>>] {
  // 创建一个新的hook结构，如果workInProgressHook已经存在就用next连接起来
  const hook = mountWorkInProgressHook();
  // useState的初始值
  if (typeof initialState === 'function') {
    initialState = initialState();
  }
  hook.memoizedState = hook.baseState = initialState;
  const queue = (hook.queue = {
    pending: null,
    dispatch: null,
    lastRenderedReducer: basicStateReducer, // 自带了一个basicStateReducer
    lastRenderedState: (initialState: any),
  });
  const dispatch: Dispatch<
    BasicStateAction<S>,
  > = (queue.dispatch = (dispatchAction.bind(
    null,
    currentlyRenderingFiber,
    queue,
  ): any));
  return [hook.memoizedState, dispatch];
}

// useState自带的reducer
function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
  return typeof action === 'function' ? action(state) : action;
}

mountState自带了一个basicStateReducer

mountReducer

function mountReducer<S, I, A>(
  reducer: (S, A) => S,
  initialArg: I,
  init?: I => S,
): [S, Dispatch<A>] {
  const hook = mountWorkInProgressHook();
  let initialState;
  // useReducer的初始值
  if (init !== undefined) {
    initialState = init(initialArg);
  } else {
    initialState = ((initialArg: any): S);
  }
  hook.memoizedState = hook.baseState = initialState;
  const queue = (hook.queue = {
    pending: null,
    dispatch: null,
    lastRenderedReducer: reducer, // 传入的reducer
    lastRenderedState: (initialState: any),
  });
  const dispatch: Dispatch<A> = (queue.dispatch = (dispatchAction.bind(
    null,
    currentlyRenderingFiber,
    queue,
  ): any));
  return [hook.memoizedState, dispatch];
}

mountState和mountReducer的区别在于函数参数和初始值赋值的不同，其他都是一样的。

进入函数首先会通过mountWorkInProgressHook来创建hook对象，接着初始化queue，通过bind将Fiber节点和queue传入dispatchAction函数实现部分参数，最后将值返回。

mountWorkInProgressHook

mountWorkInProgressHook函数功能比较简单，创建一个hook对象，将多个hook对象连接为单向链表。

function mountWorkInProgressHook(): Hook {
  // 创建新的hook对象
  const hook: Hook = {
    memoizedState: null,

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

    next: null,
  };

  if (workInProgressHook === null) {
    // workInProgressHook为null，则当前hook为函数里第一个hook，赋值给Fiber节点的memoizedState
    currentlyRenderingFiber.memoizedState = workInProgressHook = hook;
  } else {
    // 后续的hook就通过next来连接
    workInProgressHook = workInProgressHook.next = hook;
  }
  return workInProgressHook;
}

dispatchAction

dispatchAction方法为用来创建一个新的update，来发起更新

/* mountReducer ...
const dispatch: Dispatch<A> = (queue.dispatch = (dispatchAction.bind(
  null,
  currentlyRenderingFiber,
  queue,
): any));
*/

// 每次调用之前会bind上fiber和queue参数，只需要传入action
function dispatchAction<S, A>(
  fiber: Fiber,
  queue: UpdateQueue<S, A>,
  action: A,
) {
  // 省略代码...
  
  const update: Update<S, A> = {
    lane,
    action,
    eagerReducer: null,
    eagerState: null,
    next: (null: any),
  };

  // 环型链表
  // 将update添加到链表最后
  const pending = queue.pending;
  if (pending === null) {
    // pending为null的时候是第一次调用dispatchAction发起更新
    // 创建一个环形链表
    update.next = update;
  } else {
    // 将后续添加的update
    update.next = pending.next;
    pending.next = update;
  }
  // pending的值为最新的update
  queue.pending = update;

  const alternate = fiber.alternate;
  // 判断是否为render阶段的更新
  if (
    fiber === currentlyRenderingFiber ||
    (alternate !== null && alternate === currentlyRenderingFiber)
  ) {
    // 标记
    didScheduleRenderPhaseUpdateDuringThisPass = didScheduleRenderPhaseUpdate = true;
  } else {
    // 通过优先级判断是否为第一次更新(?)
    if (
      fiber.lanes === NoLanes &&
      (alternate === null || alternate.lanes === NoLanes)
    ) {
      // 当前没有更新队列，是第一次调用dispatchAction产生update,所以可以在进入render阶段之前就计算出state
      const lastRenderedReducer = queue.lastRenderedReducer;
      if (lastRenderedReducer !== null) {
        try {
          const currentState: S = (queue.lastRenderedState: any);
          const eagerState = lastRenderedReducer(currentState, action);
          // 在render阶段判断reducer如果没变化就直接取eagerState的值
          update.eagerReducer = lastRenderedReducer;
          update.eagerState = eagerState;
          if (is(eagerState, currentState)) {
            // 如果结果和当前state值一致，就不需要发起更新调度了，如果不一致，则可以在render阶段获取eagerState来直接取值而不需要再次计算
            return;
          }
        } catch (error) {
          // 错误会在render阶段抛出
        } 
      }
    }
    // 发起更新调度
    scheduleUpdateOnFiber(fiber, lane, eventTime);
  }
}

dispatchAction里会把更新连接进环形链表，如果不是render阶段的更新则会通过优先级判断Fiber节点上是否存在更新，如果不存在就会在dispatchAction里计算出新的state值，接着判断新旧值是否相同，相同就不需要发起更新调度了。

当同步调用多次dispatchAction就会产生多个update，会将他们组成环形链表。

function Updater(){
  const render = useState(0)[1];
  return <div onClick={()=>{
      render(prev => prev+1);
      render(prev => prev+1);
      render(prev => prev+1);
    }}>update</div>
}

这里的环形链表连接比较难理解

  const pending = queue.pending;
  if (pending === null) {
    update.next = update;
  } else {
    update.next = pending.next;
    pending.next = update;
  }
  queue.pending = update;

第一次执行的时候，pending === null，会创建一个自己连接自己的环形链表，这里表示为u0 -> u0。

第二次执行的时候，pending !== null，创建一个新的更新u1

• update.next = pending.next 即为u1.next = u0.next，上一次执行时u0.next -> u0 ，结果为u1.next = u0
• pending.next = update 这时候pending是u0，即为u0.next = u1
• queue.pending = update 即为queue.pending = u1，这时候最终结果为queue.pending(u1) -> u0 -> u1

第三次执行的时候，创建一个新的更新u2

• update.next = pending.next 即为u2.next = u1.next，上一次执行时queue.pending(u1) -> u0 -> u1，结果为u2.next = u0
• pending.next = update这时候pending是u1，即为u1.next = u2
• queue.pending = update 即为queue.pending = u2

这里仔细思考，实际上只改变了头和尾，中间的连接(u0 - > u1)没有改变，所以最后的结果为queue.pending(u2) -> u0 -> u1 -> u2

最终的环形链表的pending始终指向最新的update，而最新的update.next指向第一个更新u0

updateState

updateState调用的就完全是updateReducer了，只是传入了自带的reducer，所以updateState和updateReducer可谓是完全一致

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

updateReducer

mountReducer返回的值是initialState，updateReducer返回的值则是通过调用依次queue中的update计算后的state值。

function updateReducer<S, I, A>(
  reducer: (S, A) => S,
  initialArg: I,
  init?: I => S,
): [S, Dispatch<A>] {
  // update阶段对应的hook
  const hook = updateWorkInProgressHook();
  const queue = hook.queue;

  queue.lastRenderedReducer = reducer;

  // 与workInProgressHook对应的currentHook
  const current: Hook = (currentHook: any);

  let baseQueue = current.baseQueue;

  const pendingQueue = queue.pending;
  if (pendingQueue !== null) {
    if (baseQueue !== null) {
      // 省略代码...
    }
    // pendingQueue赋值给baseQueue
    current.baseQueue = baseQueue = pendingQueue;
    queue.pending = null;
  }

  if (baseQueue !== null) {
    // queue.pending为最新的update，next则为第一个update
    const first = baseQueue.next;
    let newState = current.baseState;

    let newBaseState = null;
    let newBaseQueueFirst = null;
    let newBaseQueueLast = null;
    let update = first;
    do {
      const updateLane = update.lane;
      if (!isSubsetOfLanes(renderLanes, updateLane)) {
        // 省略优先级相关代码...
      } else {
        if (newBaseQueueLast !== null) {
          // 省略代码...
        }

        // update.eagerReducer只有在第一次调用dispatchAction发起更新的时候才会赋值
        // 当reducer没有发生变化的时候
        if (update.eagerReducer === reducer) {
          newState = ((update.eagerState: any): S);
        } else {
          // 不是第一次调用dispatchAction就计算新的state
          const action = update.action;
          newState = reducer(newState, action);
        }
      }
      // 指向下一个更新
      update = update.next;
    } while (update !== null && update !== first);

    if (newBaseQueueLast === null) {
      newBaseState = newState;
    } else {
      newBaseQueueLast.next = (newBaseQueueFirst: any);
    }
    
    // 值不同才标识变化
    if (!is(newState, hook.memoizedState)) {
      markWorkInProgressReceivedUpdate();
    }

    // 赋值新的state
    hook.memoizedState = newState;
    hook.baseState = newBaseState;
    hook.baseQueue = newBaseQueueLast;

    queue.lastRenderedState = newState;
  }

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

updateReducer会依次将queue中的update放入reducer中计算，最后将新的state值赋值给hook.memoizedState并返回。

updateWorkInProgressHook

updateWorkInProgressHook和mountWorkInProgressHook功能相似，会返回一个浅拷贝的hook对象，更改currentHook和workInProgressHook的指向，同时连接一个新的workInProgressHook链表。

function updateWorkInProgressHook(): Hook {
  
  let nextCurrentHook: null | Hook;
  // 第一次进入fiber节点的时候执行hook没有hook指向
  if (currentHook === null) {
    // 找到alternate fiber节点的memoizedState hook对象
    const current = currentlyRenderingFiber.alternate;
    if (current !== null) {
      nextCurrentHook = current.memoizedState;
    } else {
      nextCurrentHook = null;
    }
  } else {
    // 后续执行hook的alternate hook对象
    nextCurrentHook = currentHook.next;
  }

  let nextWorkInProgressHook: null | Hook;
  // 省略代码...

  if (nextWorkInProgressHook !== null) {
    // 省略代码...
  } else {
    currentHook = nextCurrentHook;

    // 通过currentHook进行浅拷贝
    const newHook: Hook = {
      memoizedState: currentHook.memoizedState,

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

      next: null,
    };

    // 组装workInProgressHook链表
    if (workInProgressHook === null) {
      currentlyRenderingFiber.memoizedState = workInProgressHook = newHook;
    } else {
      workInProgressHook = workInProgressHook.next = newHook;
    }
  }
  return workInProgressHook;
}

总结

函数组件通过renderWithHooks可以确定当前的WorkInProgressFiber节点，通过是否存在currentFiber节点来判断当前为Mount还是Update，分别获取不同的ReactCurrentDispatcher，执行函数组件自己来获取children。

执行过程中，同时会执行到对应的hook函数，函数组件的hooks为单向链表存在Fiber节点的memoizedState字段上，通过hook.next可以顺序依次获取hook对象，每一个hook对象中存在memoizedState字段，对于useState和useReducer来说储存的即为state值本身，hook对象上存在queue代表当前hook的更新队列，为环形单向链表，queue.pending指向为最新的update，queue.pending.next执行为第一个update。

通过执行mountState和mountReducer来获取state初始值，通过执行updateState和updateReducer来计算queue中的update以获取最新的state值。

调用dispatchAction发起更新调度，同时在dispatchAction里会组装更新的queue环形单向链表，最后在render阶段会执行updateState和updateReducer来获取最新的state值。

如有错误，还望交流指正。