80行代码实现一个简易版useState

[li-jia-nan]

废话不多说，直接上代码：

const App = () => {
    const [num, setNum] = useState(0);
    return <div onClick={() => setNum(num + 1)}>{num}</div>;
};

上面的App组件就是我们需要实现的功能，但是我们今天的重点在useState函数，不在render函数，所以为了简化起见，我们把jsx部分删掉，返回一个方法就好了：

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(num + 1);
        }
    };
};

然后就可以直接这样调用，模拟react的点击事件：

App().onClick();

在开始实现之前，我们先分析一下，我们需要实现哪些部分，可以看到，useState包含两部分：

• 第一部分是函数本身，调用之后会返回一个数组，数组的第0项是当前的状态（对应这里的num），数组的第1项是改变状态的方法（对应这里的setNum）

• 第二部分就是setNum方法的调用，调用方法时通过某种机制，让num的值发生改变

需要知道的是，在react中，setNum方法可以接收两种参数，一种是具体的值（也就要改变后的值），还有一种是接收一个函数，比如下面这样：

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

首先我们知道，每一个组件（不管是函数组件，还是类组件，或者原生dom的组件）在react中都有一个对应的fiber节点，所以第一步，这里我们先定义一个fiber对象，并且给它一个stateNode字段，这个字段保存的就是对应的组件本身，这个fiber对象跟下面的App组件是一一对应的（并且你需要知道，在react中，有非常多的fiber，每一个fiber都有一个与之对应的组件）。

const fiber = {
    stateNode: App,
};

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

定义好fiber对象之后，我们还需要让这个mini版的react能运行起来，所以我们还需要一个用来调度的方法，可以将它命名为schedule：

const fiber = {
    stateNode: App,
};

const schedule = () => {
    //…… 
};

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

众所周知，我们每次更新，都会触发一次调度，组件就会触发一次render，所以我们调度的方法本质就是执行了一遍App函数，所以schedule函数的内部是这样：

const fiber = {
    stateNode: App,
};

const schedule = () => {
    fiber.stateNode();
};

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

这样一来，每次schedule方法执行，就相当于触发了App组件的更新，接下来还需要一个变量，因为我们的组件在mount时和update时两种情况是不一样的，比如在react的类组件中，组件首次渲染时，会调用componentDidMount钩子函数，组件更新时，会调用componentDidUpdate钩子函数，所以我们需要区分组件的渲染是mount、还是update，那么在这里需要一个全局变量isMount用来作为标识，区分两种情况：

// 申明一个全局变量 isMount，用来区分 mount 和 update
let isMount = true;

const fiber = {
    stateNode: App,
};

const schedule = () => {
    fiber.stateNode();
    isMount = false;
    // 首次渲染之后，isMount 变成 false
};

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

显然，isMount的默认值应该是true，因为组件第一次渲染必然是mount的情况，当我们调用完schedule之后，需要把isMount改变为false

接下来的问题在于，我们开始定义的useState需要保存一个对应的值，这个num需要保存在哪里呢？

前面说过，每个函数组件都有一个对应的fiber，显然，我们的useState的数据也是保存在这个fiber中的，所以这里需要一个字段用来保存对应的数据，我们给它命名为memoizedState，初始值为null。

// 申明一个全局变量 isMount，用来区分 mount 和 update
let isMount = true;

const fiber = {
    stateNode: App,
    memoizedState: null, // 用来保存组件内部的状态
};

const schedule = () => {
    fiber.stateNode();
    isMount = false;
    // 首次渲染之后，isMount 变成 false
};

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

这时候新的问题出现了，就像下面这样，一个组件可以有多个hooks：

// ……
const [num1, setNum1] = useState(0);
const [num2, setNum2] = useState(0);
const [num3, setNum3] = useState(0);
// ……

• 这些不同的hook的状态怎么保存在同一个变量上呢？
• 组件在渲染的时候，如何让多个hook的调用顺序保持一致呢？

为了解决这两个问题，我们可以通过链表的结构来保存hook的数据，也就是说，fiber中的memoizedState保存的是一个链表，这个链表保存的就是当前组件的每一个useState的数据（也就是num1、num2、num3）。

既然memoizedState是一个链表，那么我们就需要一个变量（指针），用来指向当前正在处理的hook，我们将这个变量命名为workInProgressHook，初始值为null，然后在每次调用schedule方法的时候，我们需要让指针指向当前的hook保存的值，同时为了调用方便，我们将fiber.stateNode的结果返回：

let isMount = true; // 申明一个全局变量，用来区分 mount 和 update
let workInProgressHook = null; // 申明一个全局变量，作为链表的指针

const fiber = {
    stateNode: App, // stateNode 用来保存当前组件
    memoizedState: null, // 用来保存当前组件内部的状态
};

const schedule = () => {
    workInProgressHook = fiber.memoizedState; // 让指针指向当前的useState保存的值
    const app = fiber.stateNode(); // 执行组件的渲染函数，将结果保存在app里
    isMount = false; // 首次渲染之后，isMount 变成 false
    return app; // 将fiber.stateNode的结果返回
};

const App = () => {
    const [num, setNum] = useState(0);
    return {
        onClick() {
            setNum(n => n + 1);
        }
    };
};

接下来我们要实现useState方法：

useState接收一个参数，作为初始化状态：

const useState = initialState => {
    // ……
};

通过useState，我们要计算出当前的状态，和一个改变状态的方法并返回，那么我们首先要知道，我们的useState方法对应的是哪个hook（毕竟大家调用的都是用一个方法），所以我们首先要获取到当前的useState对应的hook，先初始化一个hook变量，然后我们需要区分组件是不是首次渲染，原因是：首次渲染的时候memoizedState保存的值是null，但是在update的时候，memoizedState的值不一定是null

const useState = initialState => {
    let hook;
    if (isMount) {
        // ……
    } else {
        // ……
    }
};

首次渲染时，我们需要创建一个hook对象，这个对象保存着新的memoizedState，这个新的memoizedState对应的是hook保存的当前状态，也就是函数的参数initialState，也就是num的值（我的天，我自己都快被绕晕了）

其次我们前面说过，hook是一条链表，所以还需要一个指针next，指向下一个hook，初始值为null，代码如下：

const useState = initialState => {
    let hook;
    if (isMount) {
        hook = {
            memoizedState: initialState,
            next: null,
        };
    } else {
        // ……
    }
};

另外，在首次渲染中我们需要判断：

• 如果memoizedState不存在，说明调用函数的是第一个hook，我们需要将fiber.memoizedState指向创建的hook

• 如果memoizedState存在，说明调用函数的不是第一个hook，我们需要将workInProgressHook的next指向我们创建的hook

然后我们需要将指向当前的指针workInProgressHook赋值为当前创建的hook，这样一来，就把我们刚创建的hook和之前创建的hook连接起来了，形成了一条链表：

const useState = initialState => {
    let hook;
    if (isMount) {
        hook = {
            memoizedState: initialState,
            next: null,
        };
        if (!fiber.memoizedState) {
            fiber.memoizedState = hook;
        } else {
            workInProgressHook.next = hook;
        }
        workInProgressHook = hook;
    } else {
        // ……
    }
};

上面分析完了mount的情况，然后我们分析update的情况：

• 在mount的时候，我们已经为每一个useState创建了一个hook，并且将这些hook通过next指针连接起来，所以在update时已经有一条链表了
• 然后我们只需要将hook赋值给workInProgressHook就好了，这样就能取到对应的hook
• 然后将workInProgressHook赋值给workInProgressHook.next，让它指向下一个useState

const useState = initialState => {
    let hook;
    if (isMount) {
        hook = {
            memoizedState: initialState,
            next: null,
        };
        if (!fiber.memoizedState) {
            fiber.memoizedState = hook;
        } else {
            workInProgressHook.next = hook;
        }
        workInProgressHook = hook;
    } else {
        hook = workInProgressHook;
        workInProgressHook = workInProgressHook.next;
    }
};

经过了上面的逻辑，我们已经取到了useState保存的对应数据，下一步要做的是，基于对应的数据计算新的状态（也就是实现setNum函数），我们可以给setNum函数命名为dispatchAction，接收一个参数：

const dispatchAction = action => {
    // ……
};

那么问题来了，我们怎么知道这个dispatchAction对应的是哪个useState方法呢？（因为大家都调用的是同一个dispatchAction）

为了将dispatchAction和useState一一对应起来，我们需要将useState的hook对应的数据传给dispatchAction

回到之前的useState函数，我们看到之前创建的hook只有memoizedState和next两个属性，memoizedState保存当前的状态，next是和下一个hook连接的指针，所以，我们还需要一个新的属性，用来保存改变后的状态，这个新属性我们命名为queue：

const useState = initialState => {
    let hook;
    if (isMount) {
        hook = {
            memoizedState: initialState,
            next: null,
            queue: {
                pending: null,
            }
        };
        if (!fiber.memoizedState) {
            fiber.memoizedState = hook;
        } else {
            workInProgressHook.next = hook;
        }
        workInProgressHook = hook;
    } else {
        hook = workInProgressHook;
        workInProgressHook = workInProgressHook.next;
    }
};

将新属性命名queue是因为这是一个队列，我们通过点击事件，触发了数据更新，多次调用会触发多次更新，同样，多次调用也会创建多个action，所以我们需要把它们连接起来，所以这里的pending变量用来保存当前的hook对应的数据将要发生的改变，再回到dispatchAction函数中，在dispatchAction中我们就需要接收这个值，形参我们可以命名为queue（对应useState函数中的queue）：

const dispatchAction = (queue, action) => {
    // ……
};

接下来我们需要在dispatchAction函数里面创建一个值，代表一次更新，我们命名为update，需要注意的是，因为更新是可以多次触发的，所以这个update也是一个链表，这个update中保存着action和next：

const dispatchAction = (queue, action) => {
    const update = {
        action,
        next: null,
    };
};

下一步要做是事情，就是单纯的链表操作：

判断useState函数中的queue.pending是否存在：

• 如果不存在，说明当前的hook上面没有需要触发的更新，我们创建的updata就是需要触发的第一个更新，需要将update.next指向自己
• 如果存在，说明当前的hook上面已经有了更新，我们需要把创建的update插入队列里
• 最后再将pending赋值为当前的update（也就是说每次执行dispatchAction创建的新的update就是这条链表的最后一个update）

操作完链表之后，下一步我们要做的就是在函数的最后调用schedule方法，也就是说，让dispatchAction函数触发一次更新

const dispatchAction = (queue, action) => {
    const update = {
        action,
        next: null,
    };
    if (queue.pending === null) {
        update.next = update;
    } else {
        update.next = queue.pending.next;
        queue.pending.next = update;
    }
    queue.pending = update;
    schedule();
};

接下来，我们再次回到之前的useState函数，现在useState中的hook的queue.pending上可能存在一条链表，我们需要通过这条链表计算新的state。

这里特别说明一下：下面的部分我自己也看不懂了，快被绕晕了，但是为了坚持写完这篇文章，还是硬着头皮写完了，参考卡颂老师在B站发布的视频《React Hooks的理念、实现、源码》，感兴趣的伙伴可以去看下

要计算新的state，首先需要拿到旧的state，也就是hook.memoizedState，我们声明一个新变量用来保存旧的state，命名为baseState，然后需要判断hook.queue.pending是否存在，如果存在，代表本次更新有新的update需要被执行，我们要拿到第一个update，然后遍历链表，取出对应的action，然后基于action计算出新的state，这里需要注意的是，因为我们的action是一个函数，所以需要将baseState传给action，返回新的值，赋给baseState，然后更新firstUpdate，将firstUpdate指向下一个update，循环结束之后，将hook.queue.pending赋值为null，清空链表

const useState = initialState => {
    let hook;
    if (isMount) {
        hook = {
            memoizedState: initialState,
            next: null,
            queue: {
                pending: null,
            }
        };
        if (!fiber.memoizedState) {
            fiber.memoizedState = hook;
        } else {
            workInProgressHook.next = hook;
        }
        workInProgressHook = hook;
    } else {
        hook = workInProgressHook;
        workInProgressHook = workInProgressHook.next;
    }
    let baseState = hook.memoizedState;
    if (hook.queue.pending) {
        let firstUpdate = hook.queue.pending.next;
        do {
            const action = firstUpdate.action;
            baseState = action(baseState);
            firstUpdate = firstUpdate.next;
        } while (firstUpdate !== hook.queue.pending.next);
        hook.queue.pending = null; // 循环结束，清空链表
    }
    hook.memoizedState = baseState;
    return [baseState, dispatchAction.bind(null, hook.queue)];
};

最后把hook.memoizedState赋值为新的baseState，然后返回一个数组，数组的第0项是新的baseState，第1项是dispatchAction，因为dispatchAction函数需要传入对应的queue，所以我们需要用bind将它的this指向null，bind的第二个参数就是hook.queue

最后一步，我们在App组件中打印出num，并且在F12的控制台里面输入schedule().onClick();就可以模拟点击事件了：

// 打开F12调用onClick方法，模拟点击事件
schedule().onClick();

完整代码如下：

let isMount = true; // 申明一个全局变量，用来区分 mount 和 update
let workInProgressHook = null; // 申明一个全局变量，作为链表的指针

const fiber = {
    stateNode: App, // stateNode 用来保存当前组件
    memoizedState: null, // 用来保存当前组件内部的状态
};

function useState(initialState) {
    let hook;
    if (isMount) {
        hook = {
            memoizedState: initialState,
            next: null,
            queue: {
                pending: null,
            }
        };
        if (!fiber.memoizedState) {
            fiber.memoizedState = hook;
        } else {
            workInProgressHook.next = hook;
        }
        workInProgressHook = hook;
    } else {
        hook = workInProgressHook;
        workInProgressHook = workInProgressHook.next;
    }
    let baseState = hook.memoizedState;
    if (hook.queue.pending) {
        let firstUpdate = hook.queue.pending.next;
        do {
            const action = firstUpdate.action;
            baseState = action(baseState);
            firstUpdate = firstUpdate.next;
        } while (firstUpdate !== hook.queue.pending.next);
        hook.queue.pending = null; // 循环结束，清空链表
    }
    hook.memoizedState = baseState;
    return [baseState, dispatchAction.bind(null, hook.queue)];
};

function dispatchAction(queue, action) {
    const update = {
        action,
        next: null,
    };
    if (queue.pending === null) {
        update.next = update;
    } else {
        update.next = queue.pending.next;
        queue.pending.next = update;
    }
    queue.pending = update;
    schedule();
};

function schedule() {
    workInProgressHook = fiber.memoizedState; // 让指针指向当前的useState保存的值
    const app = fiber.stateNode(); // 执行组件的渲染函数，将结果保存在app里
    isMount = false; // 首次渲染之后，isMount 变成 false
    return app; // 将fiber.stateNode的结果返回
};

function App() {
    const [num, setNum] = useState(0);
    console.log(num);
    return {
        onClick() {
            setNum(n => n + 1);
        },
    };
};

截止目前，我们用不到80行代码实现了useState的功能。