实现useState
大家好，我是万维读客的讲师曹欢欢。上节课我们实现了函数式组件的调用，这节我们试一下在函数组件中增加useState功能，实现对应的useState函数。
useState介绍
大家开发组件经常需要处理状态，通过状态来驱动页面UI变化，react中状态主要是通过useState来保存数据。可以参考官方useState文档：https://zh-hans.react.dev/reference/react/useState，下面给个例子：
const [name, setName] = useState('Edward');
function handleClick() {
setName('Taylor');
// ...
一些限制：
- useState 是一个 Hook，因此你只能在 组件的顶层 或自己的 Hook 中调用它。你不能在循环或条件语句中调用它。如果你需要这样做，请提取一个新组件并将状态移入其中。
- 在严格模式中，React 将 两次调用初始化函数，以 帮你找到意外的不纯性。这只是开发时的行为，不影响生产。如果你的初始化函数是纯函数（本该是这样），就不应影响该行为。其中一个调用的结果将被忽略。
具体原理部分我们边写代码边介绍react是如何实现的。
编写测试用例
创建新的treact05文件夹，创建jsx.test.jsx，增加测试用例如下：
import { describe, it, expect } from 'vitest';
import * as Treact from './Treact';
describe('async render Function Component', () => {
it('render with act', async () => {
function App({title, children}){
return (
<div>
<h3>{title}</h3>
{children}
</div>
);
}
const container = document.createElement('div');
const root = Treact.createRoot(container);
await Treact.act(() => {
root.render(<App title="w3cdoc">
<App title="hello"></App>
</App>);
expect(container.innerHTML).toBe('');
});
console.log('container.innerHTML', container.innerHTML);
expect(container.innerHTML).toBe('<div><h3>w3cdoc</h3><div><h3>hello</h3></div></div>');
})
})
describe('fiber useState test', () => {
it('render useState', async () => {
const globalObj = {};
function App({ title, children }) {
const [count, setCount] = Treact.useState(0);
globalObj.count = count;
globalObj.setCount = setCount;
return (
<div>{count}</div>
);
}
const container = document.createElement('div');
const root = Treact.createRoot(container);
await Treact.act(() => {
root.render(<App />);
expect(container.innerHTML).toBe('');
});
await Treact.act(() => {
globalObj.setCount(count => count + 1);
});
await Treact.act(() => {
globalObj.setCount(globalObj.count + 1);
});
console.log('globalObj.count', globalObj.count);
expect(globalObj.count).toBe(1);
});
})
然后我们看控制台提示测试用例报错，找不到useState方法，下面我们就来补全这个方法。
实现useState
我们增加useState函数，补充基本路基如下：
...
render(element) {
// this.renderElement(element, this.container);
this._internalRoot.current = {
alternate: {
stateNode: this._internalRoot.containerInfo,
props: {
children: [element]
}
}
}
workInProgressRoot = this._internalRoot;
workInProgress = this._internalRoot.current.alternate;
// setTimeout(this.workloop.bind(this));
window.requestIdleCallback(workloop, { timeout: 100 });
}
}
...
export function useState(initialState) {
const hook = {
state: initialState,
queue: [],
}
//todo: 执行所有hooks
// ---
const setState = (action) => {
hook.queue.push(action);
// start re-render, 参考TreactRoot的render方法
workInProgressRoot.current.alternate = {
stateNode: workInProgressRoot.containerInfo,
props: workInProgressRoot.current.props,
alternate: workInProgressRoot.current, // 交替
}
workInProgress = workInProgressRoot.current.alternate;
window.requestIdleCallback(workloop, { timeout: 100 });
}
return [hook.state, setState];
}
上面我们导出了useState函数，返回对应的state和setState方法，setState方法里面是先保存要操作的动作action，然后触发重新渲染，这里是全部渲染，参考TreactRoot的render方法来触发workloop调用。我们这里修改了workloop方法，上节我们写在了TreactRoot类里面，这里我们直接移出来。
然后我们在实现执行hooks内容的部分代码，首先我们要找到hooks挂载的fiber节点，全局增加currentHookFiber来存储hook数据，增加currentHookFiberIndex来保存对应的指针，代码如下：
...
let currentHookFiber = null;
let currentHookFiberIndex = 0; // 如果有多个hooks函数
...
function performUnitOfWork(fiber) {
// console.log('fiber props', fiber.props, fiber.props == null ? JSON.stringify(fiber.return.props.children) : '');
const isFunctionComp = fiber.type instanceof Function;
if (isFunctionComp) {
currentHookFiber = fiber; // 函数式组件的当前的fiber节点
currentHookFiber.memorizedState = []; // 挂载hooks数据到memorizedState上
currentHookFiberIndex = 0; // 初始化，从第一个hooks函数开始处理
fiber.props.children = [fiber.type(fiber.props)];
} else {
...
export function useState(initialState) {
const oldHook = currentHookFiber.alternate?.memorizedState[currentHookFiberIndex];
const hook = {
state: oldHook ? oldHook.state : initialState, //获取上一次的数据
queue: [],
}
//获取事件
const actions = oldHook ? oldHook.queue : [];
actions.forEach(action => {
if(typeof action === 'function'){
hook.state = action(hook.state);
}else{
hook.state = action;
}
})
const setState = (action) => {
hook.queue.push(action);
// start re-render
workInProgressRoot.current.alternate = {
stateNode: workInProgressRoot.current.containerInfo,
props: workInProgressRoot.current.props,
alternate: workInProgressRoot.current,
}
workInProgress = workInProgressRoot.current.alternate;
window.requestIdleCallback(workloop, { timeout: 100 });
}
currentHookFiber.memorizedState.push(hook); // 更新hooks数据
currentHookFiberIndex ++; // 更新hooks函数指针
return [hook.state, setState];
}
代码到这里，大家应该能够理解hooks的渲染逻辑了，知道为什么hooks怎么执行函数了。
补全alternate
上面我们实现的useState代码都是从alternate影子节点获取数据的，大家知道为什么吗？现在我们补全一下alternate数据，代码如下：
// mount时oldFiber是空，update阶段有数据
let oldFiber = fiber.alternate?.child; //因为下面先处理的是child节点，所以这里先取child
fiber.props.children.forEach((child, idx) => {
let newFiber = null;
if (!oldFiber) {
// mount
newFiber = {
type: child.type,
stateNode: null,
props: child.props,
return: fiber,
alternate: null,
child: null,
sibling: null,
}
} else {
// update
newFiber = {
type: child.type,
stateNode: oldFiber.stateNode,
props: child.props,
return: fiber,
alternate: oldFiber, // 第一次是之前的child节点，第二次更新为sibling节点
child: null,
sibling: null,
}
}
if (idx == 0) {
fiber.child = newFiber;
} else {
preSibling.sibling = newFiber;
}
if(oldFiber){
// 第一次是之前的child节点，第二次更新为sibling节点
oldFiber = oldFiber.sibling;
}
preSibling = newFiber;
})
前面我们学习知道fiber树遍历是先处理的根节点，然后处理children子节点。根节点我们已经处理过了，现在需要再子节点遍历的时候增加对应的alternate属性。
然后查看我们测试用例执行情况，控制台输出两条测试用例都执行成功。
stdout | treact05/jsx.test.jsx > fiber setState test > render setState
globalObj.count 1
✓ treact05/jsx.test.jsx (2)
✓ async render Function Component (1)
✓ render with act
✓ fiber setState test (1)
✓ render setState
Test Files 1 passed (1)
Tests 2 passed (2)
Start at 13:48:27
Duration 23ms
PASS Waiting for file changes...
保证setState不变性
上面代码中我们每次执行useState都重新生成了一个setState函数，这其实是不必要的，这里我们可以用dispatch属性来保存下这个函数，优化代码如下：
const hook = {
state: oldHook ? oldHook.state : initialState,
queue: [],
dispatch: oldHook ? oldHook.dispatch : null,
}
...
const setState = oldHook?.dispatch ? oldHook.dispatch : (action) => {
检查测试用例通过，这里我们就基本完成了hooks中useState函数的代码部分。