Inside React useState

Intro

Ever wonder how useState works under the hood? Especially how a call to setState can trigger a UI rendering? Let's take a peek into a simple implementation of useState.

The simple implementation is taken from step 8 of "Build your own React" tutorial by Rodrigo Pombo.

Annex on the Fiber tree

Understanding the Fiber tree data structure and the two Fiber trees used in React lifecycles, current and wip, are important to understand the code of useState. I wrote extensively about the Fiber tree here and the two Fiber trees as part of the React lifecycle here.

Here are a few important notes.

Fiber tree is a tree data structure, where each node represents either a DOM node (h1, div, p, etc.) or a functional component. The entire Fiber tree can be thought of as a representation of the virtual DOM.

The current tree reflects the DOM that is currently visible on the browser. When state is changed and a new render is triggered, the wip tree can be thought of as a draft of the new DOM that is about to be rendered on the browser. Changes are made to this wip throughout a render (thus, the name work-in-progress).

There are two properties of a Fiber node that we care about.

• The alternate field is a pointer to the corresponding node in the previous render. Okay, what does "corresponding" mean here?
• The hooks field points to an array of hooks associated with that node (so only nodes that represent functional components can have non-empty hooks)

useState under the hood

function useState(initial: any) {
    // Initiate the hook for the current rendering
    const oldHook =
        wipFiber?.alternate &&
        wipFiber.alternate.hooks &&
        wipFiber.alternate.hooks[hookIndex];

    const hook = {
        state: oldHook ? oldHook.state : initial,
        queue: [],
    };

    // apply all the actions saved from last rendering
    const actions = oldHook ? oldHook.queue : [];
    actions.forEach((action: Function) => {
        hook.state = action(hook.state);
    });

    function setState(action: Function) {
        // all the actions are pushed into the queue
        hook.queue.push(action);

        // set wipRoot to initiate the rendering
        wipRoot = {
            type: currentRoot.type,
            props: currentRoot?.props,
            alternate: currentRoot,
        };

        nextUnitOfWork = wipRoot;
        deletions = [];
    }

    return [hook.state, setState];
}

Initiate the hook

In the first render, the state of the hook will be initiated with initial value. In subsequent renders, the state of the hook will be initiated as the state of the corresponding hook in the previous render.

While the first sentence is straightforward, the second is a bit hard to see from the code. Let's zoom in...

    const oldHook =
        wipFiber?.alternate &&
        wipFiber.alternate.hooks &&
        wipFiber.alternate.hooks[hookIndex];

From the wipFiber tree, wipFiber.alternate accesses the corresponding fiber node in the current tree , and wipFiber.alternate.hooks access the hooks of that node.

The index of a hook in a functional component doesn't change through renders, we can use a global variable hookIndex to track and retrieve the hook corresponding to the current hook in the previous render.

Defining setState

Now, let's look at the setState function. When setState is called with a callback function action, the action is added to a queue that is part of the hook. This queue will be unpacked in the next render, and each actions will be applied on the state.

When setState is called, it triggers a render by setting the root of the wip tree from null to the root Fiber node. However, the render is not triggered immediately. There is another function that, whenever the browser runs out of high-priority tasks, will check if wip is null. If wip is not null, it will actually begin re-rendering.

Apply all the actions on state

Functional component are essentially functions that will be run in each render. Every time a functional component is run, all of its useState will be triggered, which means new hooks and new setState functions will be created. In addition...

    // apply all the actions saved from last rendering
    const actions = oldHook ? oldHook.queue : [];
    actions.forEach((action: Function) => {
        hook.state = action(hook.state);
    });

All the actions that were saved in the oldHook.queue in the previous render (remember in the setState of the previous render where we push the action into the hook.queue?) will be applied on the old state to retrieve the new state.

Final words

And that, folks, is the anatomy of the useState function. If you want to see how useState fits into the bigger picture, definitely check out Rodrigo's tutorial.