useState 简单实现

如何实现一个 useState

整体框架

使用 setInterval 来模拟组件的更新，需要编写 useState 来管理状态，在下一次渲染的时候能拿到最新的状态。

function App() {
    const [count, setCount] = useState(0);
    const [name, setName] = useState('John');

    setCount(count + 1);
    setName(name + count);

    console.log(count, name);
}
setInterval(() => {
    App();
}, 1000)

// 输出
// 0 John
// 1 John0
// 2 John01
// 3 John012
// ......

Fiber、Hook结构

// 因为函数组件里可以有多个 useState，所以需要一个链表来存储
// 链表的每个节点是一个 Hook 对象

// 每个 Hook 对象里存储了状态、更新队列
// 更新队列也是一个链表，每个节点是一个 Update 对象
// 调用 setState 的时候，会创建一个 Update 对象，并添加到当前 Hook 的更新队列中
// 在组件下一次渲染的时候，会遍历更新队列，计算新的状态

type Hook = {
    state: any;
    queue: Update | null;
    next: Hook | null;
}

type Update = {
    action: any;
    next: Update | null;
}

type FiberNode = {
    alternate: FiberNode | null;
    hook: Hook | null;
}

// 初始化 FiberNode
App.FiberNode = {
    alternate: null,
    hook: null
}

第一次渲染时，构造 Hook 链表关系

let currentHook: Hook | null = null; // 当前正在处理的 Hook

function mountState(fiberNode: FiberNode, initialValue: any): [any, (action: any) => void] {
    const hook: Hook = {
        state: initialValue,
        queue: null,
        next: null
    }
    if (!fiberNode.hook) {
        fiberNode.hook = hook; // 将 hook 添加到 FiberNode 的 hook 链表中, 记录第一个 Hook
    } else {
        currentHook!.next = hook;
    }
    currentHook = hook;
    const setter = (action: any) => {
        const update = {
            action: typeof action === 'function' ? action : () => action,
            next: null
        }
        if (hook.queue === null) {
            hook.queue = update;
        } else {
            let lastUpdate = hook.queue;
            while (lastUpdate.next !== null) {
                lastUpdate = lastUpdate.next;
            }
            lastUpdate.next = update;
        }
    }

    return [hook.state, setter];
}

下一次渲染时，计算新的状态

function updateState(): [any, (action: any) => void] {
    const hook = currentHook;
    const baseState = hook!.state;

    let firstUpdate = hook!.queue;
    let newState = baseState;

    // 遍历 queue 中的所有 update，计算新的 state
    while (firstUpdate) {
        const action = firstUpdate.action;
        newState = action(newState);
        firstUpdate = firstUpdate.next;
    }
    hook!.state = newState;

    const setter = (action: any) => {
        const update = {
            action: typeof action === 'function' ? action : () => action,
            next: null
        }
        if (hook!.queue === null) {
            hook!.queue = update;
        } else {
            let lastUpdate = hook!.queue;
            while (lastUpdate.next !== null) {
                lastUpdate = lastUpdate.next;
            }
            lastUpdate.next = update;
        }
    }

    currentHook = currentHook!.next; // 移动到下一个 Hook

    return [newState, setter];
}

useState 的实现

function useState(initialValue: any): [any, (action: any) => void] {
    const fiberNode = App.FiberNode;
    // 第一次 mount 的时候，构造 Hook
    if (fiberNode.alternate === null) {
        return mountState(fiberNode, initialValue);
    } else {
        return updateState();
    }
}

清除操作 + 完整实现

type Hook = {
    state: any;
    queue: Update | null;
    next: Hook | null;
}

type Update = {
    action: any;
    next: Update | null;
}

type FiberNode = {
    alternate: FiberNode | null;
    hook: Hook | null;
}


let currentHook: Hook | null = null;

/**
 * 
 * @param initialValue 
 * 
 */
function useState(initialValue: any) {
    const fiberNode = App.FiberNode;

    // 第一次 mount 的时候，构造 Hook
    if (fiberNode.alternate === null) {
        return mountState(fiberNode, initialValue);
    } else {
        return updateState();
    }
}

function mountState(fiberNode: FiberNode, initialValue: any): [any, (action: any) => void] {
    const hook: Hook = {
        state: initialValue,
        queue: null,
        next: null
    }
    if (!fiberNode.hook) {
        fiberNode.hook = hook;
    } else {
        currentHook!.next = hook;
    }
    currentHook = hook;
    const setter = (action: any) => {
        const update = {
            action: typeof action === 'function' ? action : () => action,
            next: null
        }
        if (hook.queue === null) {
            hook.queue = update;
        } else {
            let lastUpdate = hook.queue;
            while (lastUpdate.next !== null) {
                lastUpdate = lastUpdate.next;
            }
            lastUpdate.next = update;
        }
    }

    return [hook.state, setter];
}

function updateState(): [any, (action: any) => void] {
    const hook = currentHook;
    const baseState = hook!.state;

    let firstUpdate = hook!.queue;
    let newState = baseState;

    // 遍历 queue 中的所有 update，计算新的 state
    while (firstUpdate) {
        const action = firstUpdate.action;
        newState = action(newState);
        firstUpdate = firstUpdate.next;
    }
    hook!.state = newState;

    const setter = (action: any) => {
        const update = {
            action: typeof action === 'function' ? action : () => action,
            next: null
        }
        if (hook!.queue === null) {
            hook!.queue = update;
        } else {
            let lastUpdate = hook!.queue;
            while (lastUpdate.next !== null) {
                lastUpdate = lastUpdate.next;
            }
            lastUpdate.next = update;
        }
    }

    currentHook = currentHook!.next;

    return [newState, setter];
}

function App() {
    const [count, setCount] = useState(0);
    const [name, setName] = useState('John');

    setCount(count + 1);
    setName(name + count);

    console.log(count, name);
    
    if (App.FiberNode.alternate === null) { // 这里模拟 alternate 的切换
        App.FiberNode.alternate = {
            alternate: null,
            hook: null
        };
    }
}

App.FiberNode = {
    alternate: null,
    hook: null
} as FiberNode;


setInterval(() => {
    App();
    // reset 
    currentHook = App.FiberNode.hook; // 这里模拟 currentHook 的 reset
}, 1000)

打印：