React18 源码解析之 useState 的原理

我们解析的源码是 React18.1.0 版本，请注意版本号。React 源码学习的 GitHub 仓库地址：https://github.com/wenzi0github/react。

useState()是我们最常见的几个 hooks 之一，今天我们来了解下他的用法和源码实现。

1. useState 的使用 #

我们先来 useState() 的用法，我们知道 useState() 返回的第 2 个参数是 dispatch(即 set 方法)，用来更新 useState() 的状态值。dispatch 的参数，既可以传入普通数据，也可以传入有返回值的函数：

import { useState } from 'react';

function App() {
  const [count, setCount] = useState(0);
  const [userInfo, setUserInfo] = useState({ name: 'wenzi', age: 24 });

  const handleClickByVal = () => {
    setCount(count + 1);
  };
  const handleClickByCallback = () => {
    setCount(count => count + 1);
  };
  const handleUpdateUserInfo = () => {
    setCount({ ...userInfo, age: userInfo.age + 1 });
  };

  return (
    <div className="App">

      <p>{count}p>

      <p>

        <button onClick={handleClickByVal}>add by valbutton>

      p>

      <p>

        <button onClick={handleClickByCallback}>add by callbackbutton>

      p>

      <p>

        <button onClick={handleUpdateUserInfo}>update userInfobutton>

      p>

    div>
  );
}

用过 useState() 的同学，还知道它还有如下的几个特点：

1. setState()的参数，既可以传入普通数据，也可以传入 callback；在以 callback 的方式传入时，callback 里的参数就是截止到当前最新的 state，使用的是执行 callback()后的返回值；
2. 传入的数据若是 object 类型，并不会自动和之前的数据进行合并，如上面的userInfo，我们需要手动合并后，再调用 set 方法；

1.1 传参的区别 #

useState()在初始时，或调用 dispatch()时，都有两种传参方式：一种是直接传入数据；一种是以函数 callback 的形式传入，state 的值就是该函数执行后的结果。

function App() {
  // 初始时传入一个callback，现在count的值就是他的返回值，即 Date.now()
  const [count, setCount] = useState(() => {
    return Date.now();
  });
}

这里我们主要关注的是多次调用 setState()时，不同的传参方式，他使用的 state 是不一样的。如

直接使用变量：

// 直接使用变量
function AppData() {
  const [count, setCount] = useState(0);

  const handleClick = () => {
    setCount(count + 1);
    setCount(count + 1);
    setCount(count + 1);
  };

  return (
    <div className="App">

      <button onClick={handleClick}>click me, {count}button>

    div>
  );
}

使用 callback 中的变量：

// 使用callback中的变量
function AppCallback() {
  const [count, setCount] = useState(0);

  const handleClick = () => {
    setCount(count => count + 1);
    setCount(count => count + 1);
    setCount(count => count + 1);
  };

  return (
    <div className="App">

      <button onClick={handleClick}>click me, {count}button>

    div>
  );
}

点击一次按钮后，这两个组件最终展示的 count 值是不一样的， 中展示的是 1，中展示的是 3。

为什么会出现这种现象呢？这是因为，在执行setCount(count + 1)时，变量 count 在函数组件的当前生命周期内，它永远是 0，因此即使调用再多的次数也没用。这里我们简化一下，就方便理解了。

function App() {
  const count = 0; // count是一个固定值

  setCount(count + 1);
  setCount(count + 1);
  setCount(count + 1);

  setTimeout(() => {
    setCount(count + 1);
  }, 1000);
}

对同一次的渲染来说，count 是一个固定值，无论在哪里使用这个值，都是固定的。setCount(count+1)的作用仅仅是把要更新的最新数据记录在了 React 内部，然后等待下次的渲染更新。

而setCount(count => count + 1)则不一样，callback 中的 prevState 则是执行到当前语句之前最新的那个 state。因此在执行第 2 条语句前，count 已经变成了 1；同理第 3 条语句。

我们稍后会从源码的层面分析下这种现象。

1.2 获取 setState()更新后的值 #

很多同学在初次使用useState()时，经常会在调用 setState()后，马上就使用更新后的数据。

function App() {
  const [count, setCount] = useState(0);

  const getList = () => {
    // console.log(count);
    fetch('https://www.xiabingbao.com', {
      method: 'POST',
      body: JSON.stringify({ count }),
    });
  };

  const handleClick = () => {
    setCount(count + 1);
    console.log(count);

    // 本意是想用更新后的最新count来调用 getList()
    getList();
  };
}

其实我们通过上面第 1.1 节的了解，已经知道此时输出的 count 还是之前的数值 0。那怎么才能使用最新的数据，来做后续的操作呢？

1. 先计算出最新值，然后同步传给 setCount()和 getList()；
2. 用 useEffect()来监听 count 的变化；

1.2.1 先计算出最新的值 #

我们可以把更新操作放在前面，先得到结果，然后再同步传给 setCount()和 getList()。

const handleClick = () => {
  const newCount = count + 1;
  setCount(newCount);
  getList(newCount);
};

这就得要求我们把函数 getList()改造为传参的形式。

1.2.2 用 useEffect()来监听 count 的变化 #

既然不确定什么时候回拿到最新的值，那我们就监听他的变化，等它更新了之后再进行后续的请求。

function App() {
  const [count, setCount] = useState(0);

  useEffect(() => {
    getList();
  }, [count]);

  const handleClick = () => {
    setCount(count + 1);
  };
}

1.3 object 类型的数据不能自动合并 #

之前在类组件中的 state，我们可以只传入需要改动的字段，React 会帮助我们合并：

class App {
  state = {
    name: 'wenzi',
    age: 24,
  };

  handleClick() {
    this.setState({ age: this.state.age + 1 }); // 只传入有改动的字段即可
  }
}

但在函组件的useState()中，传入什么数据，就使用什么数据。若还需要使用之前的部分数据时，就需要我们自己来合并数据了，然后再传给 setState()。

function App() {
  const [userInfo, setUserInfo] = useState({ name: 'wenzi', age: 24 });

  setCount({ ...userInfo, age: userInfo.age + 1 }); // 直接使用state
  setCount(userInfo => ({ ...userInfo, age: userInfo.age + 1 })); // 用callback的方式使用state
}

在类组件中，所有的状态都必须挂载在state上。在函数组件中，我们可以根据情况进行更细粒度的拆分，如 count 何 userInfo 的拆分；如果觉得 userInfo 不够精细，还可以把其中的 name 和 age 再拆分，单独进行控制。

function App() {
  const [name, setName] = useState('wenzi');
  const [age, setAge] = useState(24);
}

React 官方更推荐精细化地拆分控制，一方面是控制起来更方便，若 state 比较复杂，那在每次调用 setState()时，都要手动合并数据（当然，您可以自己实现一个自动合并数据的 hook）。另一方面在后期的维护和扩展上更容易，不必考虑其他属性的影响。

1.4 typescript 的使用 #

在 typescript 环境中，useState()是支持泛型的，state 的类型默认就是初始数据的类型，如：

function App() {
  const [name, setName] = useState('wenzi'); // name 是 string 类型
  const [age, setAge] = useState(24); // age 是 number 类型

  // userInfo 是有多个属性的类型，且已明确了属性，有且只有name和age两个属性，并且这两个属性的类型分别是string和number
  const [userInfo, setUserInfo] = useState({ name: 'wenzi', age: 24 });
}

一些相对复杂的数据类型，或者多种数据类型的组合，我们可以显式地设置 state 的类型。

enum SEX_TYPE {
  MALE = 0,
  FEMALE = 1,
}

interface UserInfoType {
  name: string;
  age: number;
  score?: number;
}

function App() {
  const [name, setName] = useState<string | null>(null); // name 是 string 类型 或 null，并且初始为null
  const [sex, setSex] = useState<SEX_TYPE>(SEX_TYPE.MALE); // sex是枚举类型

  // 显式地明确 userInfo 的各个属性，score可选
  const [userInfo, setUserInfo] = useState<UserInfoType>({ name: 'wenzi', age: 24 });

  // 更复杂的ts类型
  const [userInfo, setUserInfo] = useState<Required<Pick<UserInfoType, 'score'>>>({ score: 96 });
}

在 ts 中，明确各个变量参数的类型，一个原因是为了避免对其随意的赋值，再一个原因，从类型定义上我们就能知道这个变量的具体类型，或他的属性是什么。

我们在上面已经了解了 useState() 不少的使用方式，这里我们通过源码的角度，来看看为什么出现上面的这些现象。

2 hook 的初始挂载 #

useState() 这个 hook 的大致结构：

在第一次初始声明 useState()，state 的值就是传入的值，若不传入，则是 undefined。我们再来看下 hook 的结构：

const hook: Hook = {
  memoizedState: null, // 这个hook目前在函数组件中显示的值，初始时，即为传入的数据（若传入的是函数，则为函数执行后的结果）

  /**

   * 该hook所有的set操作开始执行时的初始值，初始挂载时，该值与 memoizedState 相同；

   * 在中间更新过程中，若存在低优先级的set操作，则 baseState 此时为执行到目前set的值

   **/
  baseState: null,

  /**

   * 执行set操作的链表，这里包含了上次遗留下来的所有set操作，和本次将要执行的所有set操作

   **/
  baseQueue: null,

  // 所有的set操作，都会挂载到 queue.pendig 上
  queue: null,

  // 指向到下一个hook的指针
  next: null,
};

注意：我们之前在讲解 hooks 挂载的时候，也讲到过 memoizedState 属性。这两个 memoizedState 属性是不一样的。fiber.memoizedState 是用来挂载 hook 节点链表的；而现在讲解的 hook.memoizedState 是用来挂载该 hook 的数值的。

function mountState(initialState: (() => S) | S): [S, Dispatch<BasicStateAction>] {
  /**

   * 创建一个hook节点，并将其挂载到 currentlyRenderingFiber 链表的最后

   * @type {Hook}

   */
  const hook = mountWorkInProgressHook();
  if (typeof initialState === 'function') {
    // 若传入的是函数，则使用执行该函数后得到的结果
    initialState = initialState();
  }
  /**

   * 设置该 hook 的初始值

   * memoizedState 用来存储当前hook要显示的数据

   * baseState 用来存储执行setState()的初始数据

   **/
  hook.memoizedState = hook.baseState = initialState;

  // 为该 hook 添加一个 queue 结构，用来存放所有的 setState() 操作
  const queue = {
    pending: null,
    interleaved: null,
    lanes: NoLanes,
    dispatch: null,
    lastRenderedReducer: basicStateReducer, // 上次render后使用的reducer
    lastRenderedState: initialState, // 上次render后的state
  };
  hook.queue = queue;

  /**

   * 这里用到了 bind() 的偏函数的特性，我们稍后会在下面进行讲解，

   *

   */
  const dispatch = (queue.dispatch = dispatchSetState.bind(null, currentlyRenderingFiber, queue));
  return [hook.memoizedState, dispatch]; // useState() 返回的数据
}

/**

 * 对当前的 state 执行的基本操作，若传入的不是函数类型，则直接返回该值，

 * 若传入的是函数类型，返回执行该函数的结果

 * @param {S} state 当前节点的state

 * @param {BasicStateAction} action 接下来要对该state执行的操作

 * @returns {S}

 */
function basicStateReducer(state, action) {
  return typeof action === 'function' ? action(state) : action;
}

setCount(count + 1); // action 是数值
setCount(count => {
  // action是函数，参数为当前的 count
  console.log('dispatch setCount');
  return count + 1;
});

function dispatchSetState(fiber: Fiber, queue, action) {}

/**

 * 这里已经提前把当前的 fiber 节点和 hook 的 queue 结构传进去了，

 * 就只留一个 action 参数给dispatch。

 */
const dispatch = (queue.dispatch = dispatchSetState.bind(null, currentlyRenderingFiber, queue));

/**

 * 设置学生的某学科的分数

 *

 * @param nick 学生姓名

 * @param subject 学科

 * @param score 分数

 */
const setStudentInfo = (nick, subject, score) => {
  console.log(nick, subject, score);
};

// 设置jack的分数
// 已预设了1个参数，剩余的两个参数供新函数设置
const setJackInfo = setStudentInfo.bind(null, 'Jack');
setJackInfo('math', 89); // Jack math 89
setJackInfo('computer', 92); // Jack computer 92

// 已预设了2个参数，剩余的一个参数供新函数设置
const setTomEnglishScore = setStudentInfo.bind(null, 'Tom', 'english');
setTomEnglishScore(97); // Tom english 97

/**

 * 派生一个 setState(action) 方法，并将传入的 action 存放起来

 * 同一个 useState() 的 setState(action) 方法可能会执行多次，这里会把参数里的 action 均会放到queue.pending的链表中

 * @param {Fiber} fiber 当前的fiber节点

 * @param {UpdateQueue} queue

 * @param {A} action 即执行setState()传入的数据，可能是数据，也能是方法，setState(1) 或 setState(prevState => prevState+1);

 */
function dispatchSetState(fiber: Fiber, queue: UpdateQueue, action: A) {}

function dispatchSetState(fiber: Fiber, queue: UpdateQueue, action: A) {
  /**

   * 获取当前 fiber 更新的优先级，

   * 当前 action 要执行的优先级，就是触发当前fiber更新更新的优先级

   */
  const lane = requestUpdateLane(fiber);

  /**

   * 将 action 操作封装成一个 update节点，用于后续构建链表使用

   */
  const update: Update = {
    lane, // 该节点的优先级，即当前fiber的优先级
    action, // 操作，可能直接是数值，也可能是函数
    hasEagerState: false, // 是否是急切状态
    eagerState: null, // 提前计算出结果，便于在render()之前判断是否要触发更新
    next: (null: any), // 指向到下一个节点的指针
  };

  if (isRenderPhaseUpdate(fiber)) {
    /**

     * 是否是渲染阶段的更新，若是，则拼接到 queue.pending 的后面

     */
    enqueueRenderPhaseUpdate(queue, update);
  } else {
    /**

     * 正常执行

     * 将 update 形成单向环形链表，并放到 queue.pending 里

     * 即 hook.queue.pending 里，存放着 update 的数据

     * queue.pending指向到update链表的最后一个元素，next即是第1个元素

     * 示意图： https://p3-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/31b3aa9d0f5d4284af1db2c73ea37b9a~tplv-k3u1fbpfcp-zoom-in-crop-mark:1304:0:0:0.awebp

     */
    enqueueUpdate(fiber, queue, update, lane);

    const alternate = fiber.alternate;
    if (fiber.lanes === NoLanes && (alternate === null || alternate.lanes === NoLanes)) {
      /**

       * 当前组件不存在更新，那么首次触发状态更新时，就能立刻计算出最新状态，进而与当前状态比较。

       * 如果两者一致，则省去了后续render的过程。

       * 可以直接执行当前的action，用来提前判断是否需要当前的函数组件fiber节点

       * 若新的state与现在的state一样，我们可以直接提前退出，

       * 若不相同，则标记该fiber节点是需要更新的；同时计算后的state可以直接用于后面的更新流程，不用再重新计算一次。

       * 根据这文档， https://www.51cto.com/article/703718.html

       * 比如从0更新到1，此后每次的更新都是1，即使是相同的值，也会再次重新渲染一次，因为两棵树上的fiber节点，

       * 在一次更新后，只会有一个fiber节点会消除更新标记，

       * 再更新一次，另一个对应的节点才会消除更新标记；再下一次，就会进入到当前的流程，然后直接return

       */
      const lastRenderedReducer = queue.lastRenderedReducer; // 上次render后的reducer，在mount时即 basicStateReducer
      if (lastRenderedReducer !== null) {
        let prevDispatcher;

        const currentState: S = (queue.lastRenderedState: any); // 上次render后的state，mount时为传入的initialState
        const eagerState = lastRenderedReducer(currentState, action);

        update.hasEagerState = true; // 表示该节点的数据已计算过了
        update.eagerState = eagerState; // 存储计算出来后的数据
        if (is(eagerState, currentState)) {
          // 若这次得到的state与上次的一样，则不再重新渲染
          return;
        }
      }
    }

    const eventTime = requestEventTime();

    /**

     * 将当前的优先级lane和触发时间给到 fiber 和 fiber.alternate，

     * 并以 fiber 的父级节点往上到root所有的节点，将 lane 添加他们的 childLanes 属性中，表示该节点的子节点有更新，

     * 在 commit 阶段就会更新该 fiber 节点

     * 这里面还存在一个任务优先级的调度，我们暂时先不考虑

     */
    const root = scheduleUpdateOnFiber(fiber, lane, eventTime);
    if (root !== null) {
      entangleTransitionUpdate(root, queue, lane);
    }
  }

  markUpdateInDevTools(fiber, lane, action);
}

/**

 * useState()的更新阶段

 * 传入要更新的值initialState，并返回新的[state, setState]

 * @param initialState

 * @returns {[(*|S), Dispatch]}

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

function updateReducer(reducer: (S, A) => S, initialArg: I, init?: I => S): [S, Dispatch] {
  if (pendingQueue !== null) {
    // 若上次更新时，有遗留下来的低优先级任务；同时当前也有要更新的任务，
    // 则将当前跟新的任务拼接到上次遗留任务的后面
    // 然后放到 baseQueue 中
  }

  if (baseQueue !== null) {
    // 当前次的更新时，更新链表不为空，那就得检查是否有可以在本地更新时要执行的任务

    do {
      if (!isSubsetOfLanes(renderLanes, updateLane)) {
        // 当前任务不满足优先级，存储起来，方便下次更新时使用
      } else {
        // 若任务优先级足够，则执行该任务；
        // 但若此时已经有低优先级的任务，为保证下次更新跳过这些任务，
        // 也会将这些任务存储起来
      }
    } while (update !== null && update !== first);

    if (newBaseQueueLast === null) {
      // 所有的任务都符合优先级，都执行完了，则下次更新时的初始值，就是上面do-while后得到的 newState 的值。
    } else {
      // 若有低优先级的任务，则将链表的最后一个节点的next指向到头结点，形成单向环形链表
    }

    if (!is(newState, hook.memoizedState)) {
      // 若新产生的 newState 跟之前的值不一样，则标记该fiber节点需要更新
      markWorkInProgressReceivedUpdate();
    }

    hook.memoizedState = newState; // 整个update链表执行完，得到的newState，用于本次渲染时使用
    hook.baseState = newBaseState; // 下次执行链表时的初始值
    hook.baseQueue = newBaseQueueLast; // 新的update链表，可能为空

    queue.lastRenderedState = newState; // 将本次的state存储为上次rendered后的值
  }

  /**

   * 返回最新的state

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

function App() {
  const [count, setCount] = useState(0);
  const [random, setRandom] = useState(0);

  const handleClick = () => {
    setCount(count => count + 1);
    setCount(count => count + 1);
    setRandom(Math.random());
    setRandom(Math.random());
  };

  console.log('refresh', Math.random());

  return (
    <div>

      <p>

        App, {count}, {random}

      p>

      <button onClick={handleClick}>click mebutton>

    div>
  );
}

function App(props) {
  const [count, setCount] = useState(props.count);

  useEffect(() => {
    // props 中的 count 属性发生变动时，重新赋值
    setCount(props.count);
  }, [props.count]);
}

function App() {
  let [count, setCount] = useState(0);
  const [user, setUser] = useState({ age: 20 });
  const [random, setRandom] = useState(0);

  // 直接修改state的值
  const handleStateDirectly = () => {
    count = 2;
    user.age = 23;

    // 值可以改变，但视图并不会更新
    console.log(count, user);
  };

  // 修改random，用来刷新组件
  const handeRandom = () => {
    setRandom(Math.random());
  };

  return (
    <div className="App">

      <p>{count}p>

      <p>{user.age}p>

      <p>

        <button onClick={handleStateDirectly}>change state directlybutton>

      p>

      <p>

        <button onClick={handeRandom}>set randombutton>

      p>

    div>
  );
}