Build your own React – Episode III

Cover photo by Suzanne D. Williams – three pupas

Welcome to the Babbel Bytes blog series on ‘Build your own (simple) React from scratch!’. In this series, we are taking a deep dive into the intricacies of React and building our own version of it based on a workshop originally created for React Day Berlin 2022.

In the previous articles, we implemented our very own dom-handlers in order to render JSX to the DOM, and we took a theoretical look at the idea of a Virtual DOM as a means to keep track of updates and avoid expensive computation in the real DOM.

In this latest installment of the series, we will start making our components a bit more useful by implementing our first hook: useState.

Hooks 101

Hooks have changed the game for functional components in React apps, allowing them to be stateful and perform side effects (a change aside from the main purpose of a function, for example tracking an event) after rendering or on prop changes.

To enable the use cases mentioned above, we decided to focus on the two most used hooks within our own React:

• useState to allow components to hold state that can be updated and can trigger further updates
• useEffect to run side effects after components are rendered to the DOM

Anatomy of a hook

Before starting to think about how we could implement a hook ourselves, let’s first think a bit about how any hook normally works.

We will take useState as an example but the elements and logic described apply to all hooks.

// We always import a react hook as a named export from react
import { useState } from 'react';

// Using a hook outside of a component is forbidden (see Rules of hooks)
useState(test); // THIS IS ILLEGAL

const Counter = () => {
    // we must use the hook in a react component
    const [count, setCount] = useState(0);
    // hooks have no identification, react relies on order of calling
    const [anotherState, setAnotherState] = useState(0);
    if (anotherState > 0) {
        // which explains the rule of hooks that hooks can't be assign conditionally
        const [anIllegalState, setAnIllegalState] = useState(0);
    }

    return (
        <div>
            {count}
            <button
               onClick={() => setCount(count => count + 1)}
            >
              +
            </button>
        </div>
    );
};

The fundamentals of using React hooks are:

1. Exporting them as a named export from React
2. Must be called inside a React functional component
3. Cannot be called conditionally, as React uses the order of calling to keep track of them. (if (condition) { useState } 🙅)

Now let’s take our Counter above and use it in an App component to think further about our architecture:

const App = () => {
    return (
        <section>
            <div>
                <div>Workshops I attended this year</div>
                <Counter />
            </div>
            <div>
                <div>Workshops I gave this year</div>
                <Counter />
            </div>
        </section>
    )
};

This App has two categories with a counter, the workshops attended and the workshop given.

Most likely, those two counts are quite different, so while they both use the same component and the same hook useState to keep track of their states, each state should be independent!

This adds one more piece of the puzzle: the state (as in the generic concept of state, not React’s one) a hook is in, is bound to the instance of the component where it is defined.

So to recap, we can say that a hook is bound to an instance of the Component it is defined in and identified by the order of the hook calls within that component e.g. the first hook defined will have an index of 0.

Now that we have a clearer picture of how a hook works, let’s close the final gaps on useState specifically.

useState allows a developer to keep track of a state local to an instance of a component, and to re-render said component after the state updates.

To do so, the developer uses the useState API which is as follows:

useState takes one argument, the initial state, which will be the first value this state will be set to when the component first renders

It returns an array where the first item is the current value of the state, and the second value is a function to update the state.

After every state update, React will make sure to re-run the component’s render function, so that it can reflect updates on its state.

const App = () => {
    const [currentStateValue, setStateValueThenRerender] = useState(initialStateValueSetOnFirstRender);
};

That last bit is very important, when a component contains a state, what it renders is derived from the value that state holds.

Let’s implement the useState hook

Now that we have covered the theory of how hooks work, and more specifically useState, it’s time to get our hands dirty and implement it.

But before we jump into it, let’s try to visualize where we are starting, and where we are trying to go.

So far, our React rendering has been taking static JSX and rendering it to the DOM.

We can see the flow of rendering goes left to right, app developers provide JSX, the React core (our index.js) transforms the JSX containing components into a renderable VDOM, that renderable VDOM is then handled by the DOM handlers and a DOM tree is created to replicate it in memory (for more about the VDOM, refer to article 2 of the series). The final step is to insert that DOM tree into the browser’s DOM, thus displaying the app to users.

Now that we have state, we would like to update what the DOM is rendering based on user input.

If we add this to the diagram above, we would want an action on the very right to be able to re-trigger the flow from the renderable VDOM, updating what that renderable VDOM will look like, then flushing those new elements to the actual DOM.

Now if we zoom in a bit, at the component level, the cycle of initial render then state update would look like this

We can see here, after a user interaction affecting a state, we must re-render the component.

During that re-render, the state value will have been updated to the new value, so the result of rendering will match the new state.

The final thing to keep in mind is, React’s core is unaware of what platform it will render to, it’s only providing a structure to create components without taking care of how they render to the platform. (This is why a framework like react-native can work, it still uses React core, but instead of using react-dom and the DOM handlers, uses native source code to render Android/iOS UI elements, you can read more about custom renderers here)

Because of that indirection, our diagram above is slightly wrong.

Technically, our app developers will interact with the DOM handlers, and the DOM handlers will call React core to get the renderable VDOM.

The call stack looks something like this.

So our goal will be to add within our Core a mechanism to keep and update state, and each state update must re-trigger the component’s instance render function, create a new renderable VDOM and finally flush all those updates to the DOM.

For that purpose, we will create a subscription system between DOM handlers and Core, so that DOM handlers can be made aware of updates to the renderable VDOM and re-render every time an update happens, making our architecture look like this

On this new branch, we have set up the starting point for hooks to work as we described above.

The first very big difference compared to how we have worked so far is that our React will now need to re-render after each state update, but so far we have only been rendering to the DOM statically once.

To achieve this, we created a subscription model in the dom-handlers, so that our react render is able to call the DOM renderer several times.

If you head to our dom-handlers, you will see we are now calling a function startRenderSubscription to subscribe to DOM changes, and on each call we will re-render our DOM.

This new function is defined in our package’s index file and looks like this

// this function will call the updateCallback on every state change
// so the DOM is re-rendered
export const startRenderSubscription = (element, updateCallback) => {
  // 1
  let vdom = {
    previous: {},
    current: {},
  };
  // 2
  const update = hooks => {
    // 2.1
    const renderableVDOM = rootRender(element, hooks, vdom);

    // 2.2
    vdom.previous = vdom.current;
    vdom.current = [];

    // 2.3
    updateCallback(renderableVDOM);
  };
  // 3
  const hooks = createHooks(update);
  // 4
  update(hooks);
};

In a nutshell, it does the following:

1. Keeps track of the virtual DOM over time by keeping a previous and current version
2. Creates an update function (let’s break this one down after)
3. Create our hooks structure
4. Calls the update function with the hooks we just created for the first render

The update function, which is where most of the code happens, does the following:

1. Uses rootRender to get the renderableVDOM given the provided element, the hooks and the previous and current VDOM
2. Updates the VDOM references to be ready for the next render
3. Calls the updateCallback with the renderableVDOM to refresh the UI the user sees.

So the changes we made to this are mostly the fact we can now call the updateCallback to re-render the UI, and we have prepared hooks and VDOM structures to be able to keep track of them over time.

Now one more tricky part of how hooks work with React is the way their API is designed.

While each hook is bound to a specific instance of a component, the developer using the hook does not do anything specific for that binding to happen.

They just import the hook from react and use them within a component.

As creators of a new React library, this means we will need to find a stratagem to bind each hook call to a specific component correctly, somehow hooking into the functions we are exporting and adapting what they do when they are called.

Let’s continue digging into the existing changes, we are now passing hooks and VDOM to our rootRender to be able to track changes.

Concretely, looking at the renderComponentElement function, this entails the following pieces of extra logic:

Keeping track of elements in our current VDOM

 setCurrentVDOMElement(
    VDOMPointer,
    createVDOMElement(element, VDOMPointer),
    VDOM,
  );

Calling a registerHooks function for functional components, this is how we will take care of the binding of hooks to a specific component

hooks.registerHooks(VDOMPointer, isFirstRender);

And to derive whether a component is rendered for the first time for isFirstRender, we can use our previous VDOM:

// Access the previous element (or undefined if not found) 
const previousDOMElement = (getVDOMElement(VDOMPointer, VDOM.previous) || {}).element;

// If we have a previous element, verify the types are matching
const isFirstRender =
  previousDOMElement === undefined ||
  previousDOMElement.type !== element.type;

This new code will allow us to keep track of our hooks as per the needs we have expressed before: registerHooks is called with a component’s instance (referenced by its unique ID: the VDOMPointer) and whether it’s rendered for the first time (to enable us to reset hooks when needed).

And during the creation of our hooks, we provide an update function which will enable us to re-render after a hook update.

Now let’s think a bit more about how we will keep track of each hooks.

As we mentioned earlier, we want each component’s instance to have its own hooks’ states kept separately.

So far, we have used VDOMPointer as unique IDs for each of the elements in our tree, including components, so we can safely reuse this as a unique ID for a component’s instance.

If doing so, we could create a map where the keys are VDOMPointers and the values are hooks state.

interface HooksMap {
 [VDOMPointer]: HooksState
};

And as we mentioned earlier, we will need to bind components to a specific component’s instance with the registerHooks function.

For this, we will need to be able to somehow update the behavior of calling the useState function so that it will be aware of which component’s instance it is bound to.

To do this, we decided to make the function swappable with another function at runtime.

We created an object on which we can create our global function, and we will update the function in that object to replace the useState hook

let globalHooksReplacer = {};

export const useState = (...args) => globalHooksReplacer.useState(...args);

Now the role of registerHooks will be to make sure that the globalHooksReplacer.useState function is updated for the current component’s instance.

So our register hooks function, at it’s simplest, would look like this

const registerHooks = (VDOMPointer, isFirstRender) => {
  if (isFirstRender) {
    resetHooks(VDOMPointer);
  }
  globalHooksReplacer.useState = setupUseState(VDOMPointer);
}

We technically have a few more dependencies for this function we need to take care of, it needs to have access to the hooks map (so it can reset hooks and set their state), and to the update function so that the UI can be re-rendered after a state update.

For the sake of allowing us to split responsibilities, we have used higher-order functions to manage dependencies of the different parts.

So our final registerHooks function is implemented like this:

const makeRegisterHooks =
  (hooksMap, makeUseState) => (VDOMPointer, isFirstRender) => {
    if (isFirstRender) {
      hooksMap[VDOMPointer] = {};
    }
    const useState = makeUseState(VDOMPointer, isFirstRender);
    globalHooksReplacer.useState = useState;
};

We receive the hooks map and a higher order function makeUseState to make the useState function as our first arguments (our highest dependencies).

And based on this, we return a function which will be useful to us, which receives the VDOMPointer and isFirstRender arguments to set up the hooks for the component’s instance matching the VDOMPointer.

In the function’s body, we reset the hooks map for the component’s instance by setting it to an empty object.

We then create the appropriate useState by calling our provided makeUseState function.

Finally, we can replace our global useState by the useState we just created, which is bound to our component’s instance.

Now as we saw above, useState also has more dependencies than just the VDOMPointer and isFirstRender.

It also needs to be able to trigger a re-render (after a state update) and need to access the hooks map to retrieve the current state.

We set it up with a higher order function to, similar to makeRegisterHooks, but we have two set of dependencies we will want to resolve at different times:

1. The update and hooks map we will resolve when setting up the hooks initially
2. The VDOMPointer and isFirstRender we will resolve once we are about to render a component instance

For that reason, our function to create state is actually a higher order higher order function, as in a higher order function which returns a higher order function. 🤯

Let’s have a look at the signature of that monster!

const createMakeUseState =
  // First set of dependencies to be resolved earlier
  (onUpdate, hooksMap) => 
  // Second set of dependencies to be resolved later
  (VDOMPointer, isFirstRender) => {
    // Actual useState function developers will call in their components
    return initialState => {
         return [state, setState];
      }
    };
  };

Now that we have thought about the structure of our hooks, the last thing we have to do is to wire it all up!

In our index, you might remember us calling a function createHooks, this is the function that will set up everything and connect the functions we created above:

export const createHooks = onUpdate => {
  // 1
  const hooksMap = {};
  // 2
  const hooks = { current: null };
  // 3
  const boundOnUpdate = () => onUpdate(hooks.current);
  // 4
  const makeUseState = createMakeUseState(boundOnUpdate, hooksMap);
  // 5
  const registerHooks = makeRegisterHooks(hooksMap, makeUseState);
  // 6
  hooks.current = { registerHooks };
  // 7
  return hooks.current;
};

Let’s break down each part

1. We create the initial hooksMap to keep track of the hooks’ states of each component
2. We create a variable similar to a React ref to keep our hooks object (this is useful because we need to pass the hooks to onUpdate, but we need to use onUpdate in order to create our hooks object), so using that ref like structure, we are able to already use the variable but initialize it later
3. We bind the current hooks to the onUpdate function to simplify the calls to it.
4. We create makeUseState by providing its dependencies: the boundOnUpdate function and the hooks map.
5. We create registerHooks by providing its dependencies: the hooks map and the makeUseState
6. We replace the current value of our hooks by an object containing our registerHooks function
7. We return the hooks object

Phewww, that was a lot of setup! 🥵

But now we have a structure that should enable us to appropriately track our hooks, the last thing we have to do now is to implement the logic of useState!
We have already created a few things when it comes to the useState:

const createMakeUseState =
  (onUpdate, hooksMap) => (VDOMPointer, isFirstRender) => {
    // Reminder: this is the function that will execute when a developer calls `useState` in their component
    return initialState => {
      // 1
      if (isFirstRender) {
        // 2
        const computedInitialState =
          typeof initialState === 'function' ? initialState() : initialState;

        // 3
        const setState = newStateOrCb => {
          // We are missing code here for setState to work
          // 4
          const newStateFn =
            typeof newStateOrCb === 'function'
              ? newStateOrCb
              : () => newStateOrCb;
          const currentState = newStateFn(previousState);

        };

        // 5 (we are missing the setState function in the return)
        return [computedInitialState, () => 'replace this function'];
      }
    };
  };

1. On the first render, useState should return the initial state
2. The initial state can be a value or a function. The function version is used for React’s lazy initial state which can be used to improve performance if the initial state requires expensive computation. This means to get our initial state value, if initialState is a function, we should call this function once on the first render.
3. We create the setState function which developers will use in their components
4. setState supports its call value to be a value (the new state value) or a function for functional updates, in this case we will need to call that function with the previousState to get the new state

If you are coding along, it is now your time to start with the first task of this episode: filling in the blanks in the createMakeUseState function so that it will work correctly for component’s instances which call useState once.

To help you a little bit, here is a diagram recapitulating what the data structure we are using looks like

You will need to:

1. Think about where we want to store/retrieve the state value based on what we mentioned so far
2. Complete the setState function so that it actually updates the state and re-renders
3. Add a return statement for cases other than the first render

If you succeed, you should see the todo app updating when you change states. 🎉

It will have some weird behaviors for now, which we will explain later
As a first step, let’s retrieve the hooks structure for our current component:

const createMakeUseState =
 (onUpdate, hooksMap) => (VDOMPointer, isFirstRender) => {
   let currentHook = hooksMap[VDOMPointer];

Basically, we use the hooksMap as our data structure, and it is organized by VDOMPointer, so to retrieve the hooks of the current component, we just retrieve the hooks at VDOMPointer.

Next we need to fill in the blanks in the setState function

const setState = newStateOrCb => {
  const newStateFn =
    typeof newStateOrCb === 'function'
      ? newStateOrCb
      : () => newStateOrCb;
  // 1
  const ownState = currentHook.state;
  // 2
  const previousState = ownState[0];
  // 3
  const newState = newStateFn(previousState);
  // 4
  ownState[0] = newState;
  // 5
  onUpdate();
}

1. We will store the current state of our component’s instance under the state key, and we will make it reflect the return value of use state, meaning each state will be a tuple where the first item is the current state value, and the second item the state setter function
2. We retrieve our previous state by accessing the first item of our state
3. We compute the new state based on the previous state
4. We update the state value with the new state
5. We call the onUpdate function to re-render the UI

Now we only need to fill in the last blanks and update the return statements, here is the completed createMakeUseState function

const createMakeUseState =
  (onUpdate, hooksMap) => (VDOMPointer, isFirstRender) => {

   let currentHook = hooksMap[VDOMPointer];

    return initialState => {

      if (isFirstRender) {
        const computedInitialState =
          typeof initialState === 'function' ? initialState() : initialState;

        const setState = newStateOrCb => {
          const newStateFn =
            typeof newStateOrCb === 'function'
              ? newStateOrCb
              : () => newStateOrCb;
          const ownState = currentHook.state;
          const previousState = ownState[0];
          const newState = newStateFn(previousState);
          ownState[0] = newState;
          onUpdate();
        }
        // 1
        currentHook.state = [computedInitialState, setState];
      }
      // 2
      return currentHook.state;
    };
  };

We just set our currentHook.state value to be a tuple with the first item, the state value, as the computed initial state, and the second value as the setState function.

Then we always return the currentHook.state from the function.

An important note here, while it would theoretically be possible to recreate the setState function on every call to useState, React’s documentation states:
Note

React guarantees that setState function identity is stable and won’t change on re-renders. This is why it’s safe to omit from the useEffect or useCallback dependency list.

So we must keep the same function on each call in order not to break other APIs.

You should now see our todo app starting to work! 

HOORAY! 🎉

You can start creating a todo and the UI will update…

But then you might notice something weird happens: the counter updates to the list of TODO 😲

The reason for that is some of our components have two calls to useState, but for now our useState only handles one state.

Let’s fix that!

Handling several states

If following along, it is now time to go to the first branch git checkout chapter-2/step-2

We had already mentioned the fact that developers can call useState several times within one component to keep track of several states, but to reduce the complexity of our first task we had ignored that so far.

In terms of code updates, you will see most of the new code in the branch is the code we covered above, with one very small addition, before updating the state and re-rendering, we added a condition to verify that the state has indeed changed.

As a reminder, the way React tracks each state is by order of invocation, so the “unique id” of each state would be their invocation index.
For example, given the following component

const ComponentWithTwoStates = () => {
  const [counter, setCounter] = useState(0);
  const [userHasClicked, setUserHasClicked] = useState(false);
}

We could give the counter state the ID 0 and the userHasClicked the ID 1.

If you are coding along, you can look at the instructions left in the code to get you started.

To be able to keep several states per component, we will need to rethink our data structure a bit.

So far, we have held a unique state in our hooks in the hooks map as its value.

Instead, we will now need an array of states, where each item in the array will be a different state the component’s instance tracks.

So for the example above, we could imagine the data structure to be:

Or for the currentHook specifically:

const currentHook = {
  state: [[0, setState1], [false, setState2]]
};

We will then need a mechanism to allow each useState to know which state it is dealing with.

For this, we can use a Ref-like variable as an index (technically a let variable could also work in this instance).

Its initial value will be 0, and we will increment it on each call to useState.

const createMakeUseState =
  (onUpdate, hooksMap) => (VDOMPointer, isFirstRender) => {
    const stateIndexRef = { current: 0 };
    let currentHook = hooksMap[VDOMPointer];

    return initialState => {
      const stateIndex = stateIndexRef.current;
      stateIndexRef.current += 1;
    }
}

Now we have a way to know which precise state each useState is dealing with.

Then we only need to update our code to use the data structure we decided upon earlier, so the final code looks like this:

const createMakeUseState =
  (onUpdate, hooksMap) => (VDOMPointer, isFirstRender) => {
    const stateIndexRef = { current: 0 };
    let currentHook = hooksMap[VDOMPointer];
    if (isFirstRender) {
      // 1
      currentHook.state = [];
    }

    return initialState => {
      const stateIndex = stateIndexRef.current;
      stateIndexRef.current += 1;
      if (isFirstRender) {
        const computedInitialState =
          typeof initialState === 'function' ? initialState() : initialState;

        const setState = newStateOrCb => {
          const newStateFn =
            typeof newStateOrCb === 'function'
              ? newStateOrCb
              : () => newStateOrCb;

          // 2
          const ownState = currentHook.state[stateIndex];
          const previousState = ownState[0];
          const newState = newStateFn(previousState);
          const shouldUpdateState = isStatesDiffer(previousState, newState);

          if (shouldUpdateState) {
            ownState[0] = newState;
            onUpdate();
          }
        };
        // 3
        currentHook.state[stateIndex] = [computedInitialState, setState];
      }
      // 4
      return currentHook.state[stateIndex];
    }
}

1. We initialize the state array to empty on the first render so we are ready to keep track of states
2. The own state of each useState is now dependent on the index at which it was called
3. Similarly to 2., we set the state based on the stateIndex now
4. Similarly to 2 and 3, we return the state for useState based on it’s call order

Conclusion

That’s a wrap 🎁 for this third article!

Congratulations on following along! 👏

We are really getting somewhere now!

Our react-like library renders all our components, and is even able to track state so it re-renders after a state changes!

That’s really powerful! ⚛️

Take the chance to play a bit with the TODO app, see it re-render and think about all that was involved to get here. Pin a little medal ⭐ on you for taking on the journey!

After playing with it for a bit, you will probably notice that the input field is behaving a bit erratically, every time you type in a character, it loses the focus 😠 

That’s because at the moment, we fully re-render every single component on the page for each state update.

This means the input that had the focus before the state update has been destroyed after it, and has been replaced by a new one 😱

That’s one of the reasons we would want to only update the things in the DOM that need to be updated (in this case, only update the value of the input field, instead of completely removing/reading it). Another reason is of course performance, as it is not very efficient to always re-render the whole DOM while only a very small portion of it changed. 

Hope to see you in the next article to see how we can improve that with the diffing algorithm and targeted updates!

---

This article was written as a trio in collaboration with Ömer Gülen and Sean Blundell.