【React Hooks系列】源码解析

案例

function App() {
    // 注意：这些 hooks 在每次 app render 的时候必须保证完全一样的顺序
    // 因为它们是保存在一条链表中的，而这条链表的顺序是不变的
    const [num, updateNum] = useState(0);
    const [num1, updateNum1] = useState(0);
    const [num2, updateNum2] = useState(0);
    return {
        onClick() {
            updateNum(num => num + 1);
        }
    }
}

实现一个useState

基于上面的案例实现一个 mini 版的 useState:

// 组件是否已挂载
let isMount = true;
// 当前正在处理的 hook
let workInProgressHook = null;

const fiber = {
    stateNode: App,
    // 是一条链表，保存的是该函数组件每个 hook 对应的数据
    memoizedState: null,
}

function useState(initialState) {
    // 由于一个组件中可能同时存在多个 useState
    // 所以需要获取当前的 useState 对应的 hook
    let hook;

    if (isMount) {
        // 组件第一次渲染，这三行代码都会走这块逻辑
        // const [num, updateNum] = useState(0);
        // const [num1, updateNum1] = useState(0);
        // const [num2, updateNum2] = useState(0);
        hook = {
            // 该 memoizedState 保存的就是 num，num1，num2
            memoizedState: initialState,
            // hook 是一条链表
            next: null,
            // 队列，保存当前 hook 对应的最新数据的变化，考虑下面这种情况：
            // onClick() {
            //  updateNum(num => num + 1);
            //  updateNum(num => num + 1);
            // }
            // 所以 queue 需要保存这些 action 即：num => num + 1、num => num + 1
            queue: {
                pending: null,
            }
        }
        if (!fiber.memoizedState) {
            // 组件第一次渲染时第一个 useState 会走这里
            fiber.memoizedState = hook;
        } else {
            // 组件第一次渲染时非第一个 useState 会走这里
            workInProgressHook.next = hook;
        }
        workInProgressHook = hook;
    } else {
        // 组件更新阶段，此时 memoizedState 已经是一条链表了
        // 因为 mount 时已经为每个 useState 创建了 hook，并将这些 hook 通过 next 指针连接成为一条链表
        hook = workInProgressHook;
        workInProgressHook = workInProgressHook.next;
    }

    // 基于已有的 state 计算新的 state
    let baseState = hook.memoizedState;

    if (hook.queue.pending) {
        // hook.queue.pending 存在代表本次更新有新的 update 需要被执行
        let firstUpdate = hook.queue.pending.next;
        do {
            const action = firstUpdate.action;
            // 此处该例中 action 是一个函数
            // updateNum(num => num + 1);
            baseState = action(baseState);
            firstUpdate = firstUpdate.next;
            // 当 firstUpdate 不等于第一个 update 时一直循环
        } while (firstUpdate !== hook.queue.pending.next)

        // 所有的 action 计算完成，将环状链表清空
        hook.queue.pending = null;
    }

    // hook.memoizedState 更新为最新的 state
    hook.memoizedState = baseState;

    return [baseState, dispatchAction.bind(null, hook.queue)];
}

// 对应 useState 的更新 num 的 updateNum 方法
// updateNum(num => num + 1);
// action即：num => num + 1
// queue即：queue: {
//    pending: null,
// }
function dispatchAction(queue, action) {
    // 考虑这种情况，所以 update 是一个链表
    // onClick() {
    //  updateNum(num => num + 1);
    //  updateNum(num => num + 1);
    // }
    const update = {
        action,
        next: null
    }

    // 代表当前的 hook 还没有要触发的更新
    // 注意：每个 hook 对应的 update 是一个环状链表，因为在真实的 React 中
    // 不仅要计算新的 state，且每个 update 每次更新是有优先级的
    // 比如点击触发的更新优先级高于网络请求再更新的优先级
    // 因此用户点击交互可以更快得到响应，所以用双向链表更容易操作
    // 在 update 的链表上，有些 update 需要被执行，有些不需要被执行
    if (queue.pending === null) {
        // 这种情况下第一个更新进入该逻辑
        // onClick() {
        //  updateNum(num => num + 1);
        //  updateNum(num => num + 1);
        // }
        // u0 -> u0
        update.next = update;
    } else {
        // 其余的更新进入这个逻辑
        // u1 -> u0 -> u1
        // queue.pending 保存的是最后一个 update（u1）
        // 那么 queue.pending.next 保存的就是这个环的第一个 update（u0）
        // 下面这行代码表示环的最后一个 update 指向环的第一个 update（u1 -> u0）
        update.next = queue.pending.next;
        // u0 -> u1
        queue.pending.next = update;
    }
    // queue.pending 保存的是最后一个 update
    queue.pending = update;

    // 触发更新
    schedule();
}

function schedule() {
    // 每次 schedule 触发一次新的更新时，都需要将 workInProgressHook 指向 fiber 的第一个 hook
    workInProgressHook = fiber.memoizedState;
    const app = fiber.stateNode();
    isMount = false;
    return app;
}

// 模拟点击
window.app = schedule();
app.onClick();

流程总结

1. 用户在代码中调用 updateNum(num => num + 1) 一次或者多次
2. 进入 dispatchAction，该方法会计算出一个环状链表用于存储一系列 action（存储在 hook 的 queue.pending 中），即:

    // action 指的是：num => num + 1、num => num + 1 ...
    onClick() {
        updateNum(num => num + 1);
        updateNum(num => num + 1);
    }

计算环状链表结束之后触发更新 schedule

1. 进入 schedule，从第一个 hook 开始更新组件，进入 useState
2. useState 中做了以下几件事情：
   • 挂载时：构建 fiber 的 memoizedState 为一个单向链表
   • 更新时：通过 workInProgressHook 取到当前需要更新的 hook（同一个组件可能同时存在多个 useState）
   • 取出该 hook 的环状链表存储的一系列 action（hook.queue.pending），遍历该环状链表，并基于 hook 的 memoizedState（上一次的 state）计算出新的state
   • 将 fiber 的 memoizedState 赋值为最新的 state
   • 将最新的 state 返回给用户

源码实现

React 源码中，不同阶段定义了不同的 hooks，效果与上例中我们自己定义的 isMount 相同。比如：

挂载阶段

const HooksDispatcherOnMount: Dispatcher = {
  readContext,
  use,
  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,
  useSyncExternalStore: mountSyncExternalStore,
  useId: mountId,
};

更新阶段

const HooksDispatcherOnUpdate: Dispatcher = {
  readContext,
  use,
  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,
  useSyncExternalStore: updateSyncExternalStore,
  useId: updateId,
};

mountState（useState 挂载阶段）

function mountState<S>(
  initialState: (() => S) | S,
): [S, Dispatch<BasicStateAction<S>>] {
  const hook = mountStateImpl(initialState);
  const queue = hook.queue;
  const dispatch: Dispatch<BasicStateAction<S>> = (dispatchSetState.bind(
    null,
    currentlyRenderingFiber,
    queue,
  ): any);
  queue.dispatch = dispatch;
  return [hook.memoizedState, dispatch];
}

function mountStateImpl<S>(initialState: (() => S) | S): Hook {
  const hook = mountWorkInProgressHook();
  if (typeof initialState === 'function') {
    // $FlowFixMe[incompatible-use]: Flow doesn't like mixed types
    initialState = initialState();
  }
  hook.memoizedState = hook.baseState = initialState;
  const queue: UpdateQueue<S, BasicStateAction<S>> = {
    pending: null,
    lanes: NoLanes,
    dispatch: null,
    lastRenderedReducer: basicStateReducer,
    lastRenderedState: (initialState: any),
  };
  hook.queue = queue;
  return hook;
}

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

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

    next: null,
  };

  if (workInProgressHook === null) {
    // This is the first hook in the list
    currentlyRenderingFiber.memoizedState = workInProgressHook = hook;
  } else {
    // Append to the end of the list
    workInProgressHook = workInProgressHook.next = hook;
  }
  return workInProgressHook;
}

updateState（useState 更新阶段）

useState 与 useReducer 的实现逻辑是完全一致的，唯一的区别是 useState 是一个预制了 reducer 的 useReducer，而 useReducer 需要我们自己定义一个 reducer。

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

// 对应的是 hook 更新方式的两种写法：
// updateNum(num => num + 1);
// updateNum(num + 1);
function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
  // $FlowFixMe[incompatible-use]: Flow doesn't like mixed types
  return typeof action === 'function' ? action(state) : action;
}

// updateWorkInProgressHook 的作用就是获取到当前 useState 对应的 hook
// updateReducerImpl 的作用是基于已有的 state 计算新的 state
function updateReducer<S, I, A>(
  reducer: (S, A) => S,
  initialArg: I,
  init?: I => S,
): [S, Dispatch<A>] {
  const hook = updateWorkInProgressHook();
  return updateReducerImpl(hook, ((currentHook: any): Hook), reducer);
}

mountEffect（useEffect 挂载阶段）

在 React 不同的生命周期，React 会从 hook 中取到对应的 create 方法，判断是否需要执行

function mountEffectImpl(
  fiberFlags: Flags,
  hookFlags: HookFlags,
  create: () => (() => void) | void,
  deps: Array<mixed> | void | null,
): void {
  // 获取当前 useEffect 对应的 hook
  const hook = mountWorkInProgressHook();
  const nextDeps = deps === undefined ? null : deps;
  currentlyRenderingFiber.flags |= fiberFlags;
  // hook.memoizedState 中保存 create 及 nextDeps
  // create 即回调函数，nextDeps 即依赖数组
  // useEffect(() => {xxx}, [xxx])
  hook.memoizedState = pushEffect(
    HookHasEffect | hookFlags,
    create,
    createEffectInstance(),
    nextDeps,
  );
}

mountRef（useRef 挂载阶段）

const xxxRef = useRef('hello');
// 在 React 的实现中，等同于：
// const { current: 'hello' } = useRef('hello');

function mountRef<T>(initialValue: T): {current: T} {
  const hook = mountWorkInProgressHook();
  const ref = {current: initialValue};
  hook.memoizedState = ref;
  return ref;
}

updateRef（useRef 更新阶段）

function updateRef<T>(initialValue: T): {current: T} {
  // 获取 ref 对应的 hook 
  const hook = updateWorkInProgressHook();
  // 返回对象 { current: xxx }
  return hook.memoizedState;
}

mountMemo（useMemo 挂载阶段）

// 在依赖项改变时返回重新计算后的值
const xxx = useMemo(() => 'aaa', [])

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;
}

updateMemo（useMemo 更新阶段）

updateMemo 中的 areHookInputsEqual 用的是浅比较，即 Object.is

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

mountCallback（useCallback 挂载阶段）

// fn 即缓存的函数
const fn = useCallback(() => console.log('hello'), [xxx])
// 使用 useCallback 时，App 组件接收到的 fn 函数是同一个引用，如果 App 组件有使用 PureComponent 或者使用 React.memo 包裹的话，App 组件不用重新 render
return <App fn={fn} />

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;
}

updateCallback（useCallback 更新阶段）

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 (nextDeps !== null) {
    const prevDeps: Array<mixed> | null = prevState[1];
    // 比较本次更新与上次的依赖项是否完全一致
    if (areHookInputsEqual(nextDeps, prevDeps)) {
    // 一致，返回缓存的回调函数
      return prevState[0];
    }
  }
  // 不一致，返回本次新的回调函数
  hook.memoizedState = [callback, nextDeps];
  return callback;
}