Let’s build a React from scratch: Part 2— State Management and React Hooks

Disclaimer: The content below is just for learning and have some key insights while using framework like React. While the inspiration is React, the idea is to deliver the core concepts of what it takes to build a library like React at its most basic level.

Welcome, on our journey to building React of our own, it's just getting more and more exciting. In this part, we will try to see and explore the purpose of a library like React which not only helps in building UI Components but also makes them useful and interactive by managing dynamic states within it and updating the UI accordingly. We will also learn about React Hooks along the way and why some of the rules of Hooks actually apply (like no hooks in conditionals) in reality.

• Part 1 — VirtualDOM and React renderer
• Part 2 — State Management and React Hooks
• Part 3 — React Suspense and Concurrent Mode
• Part 4 — Server Side Rendering and its Challenges

🛥 Our Setup

We will continue with our setup and code which we developed in Part-1 of this series. The easiest way to get set up is to fork this code to a brand new app in Stackblitz. Or you can use your own Typescript based setup following the previous instructions.

🤹 Introduction to React Hooks (useState)🎪

Let’s start with an example here, let's say we want to manage a state which holds say the name of the current user (any string for that matter), and we want to be able to modify it and see it reflect on the UI immediately. Well, that’s a mouthful, but how would we go about achieving it? Let's see how to build it from scratch. So, as we have done so far, let’s try to first use useState hook which of course doesn’t exist yet. But we will define it.

// ---- Application --- //
const App = () => {
  const [name, setName] = useState('Arindam');
  return (
    <div draggable>
      <h2>Hello {name}!</h2>
      <p>I am a pargraph</p>
      <input
        type="text"
        value={name}
        onchange={(e) => setName(e.target.value)}
      />
    </div>
  );
};

And as usual, we immediately get Typescript complaining about it.

Let’s define our useState here, please note how it's structured, it takes an initial state useState(initialState) and we are returning two things from here, we are returning the state itself name and a setter for our state setName which is used to modify this state.

// ---- Library --- //
const useState = (initialState) => {
  let state = initialState;
  const setState = (newState) => (state = newState);
  return [state, setState];
};

Let’s also add a few console.log() to see what’s going on how they are used within a ReactComponent.

// ---- Library --- //
const useState = (initialState) => {
  console.log('useState is initialized with value:', initialState);
  let state = initialState;
  const setState = (newState) => {
    console.log('setState is called with newState value:', newState);
    state = newState;
  };
  return [state, setState];
};

Now, if we start changing the input value we have a onchange handler that should call setName() and it should be logged in our console with the new and updated values accordingly. And sure enough, our console shows the same.

setState is in action

But, one big question here is how does our UI know that the state has changed and we should update it. This is where the next big concept of React unfolds, which is to re-render the UI whenever the state changes. And we know for sure, if we makesetState calls that changes the state, will always require a re-render.

🌗 Introducing Re-render for our App 🌝

For the purposes of this post, we will keep the re-render implementation very simple. In real-world, React doesn’t re-render the entire App, it does an intelligent diff on what has changed and what minimal change on the UI will correctly reflect the state and only change that much.

So, our dumb version of the re-render might look like this. which simply calls the render fresh for the App.

// ---- Library --- //
const reRender = () => {
  console.log('reRender-ing :)');
  render(<App />, document.getElementById('myapp'));
};

Let’s use this reRender after any state change in our useState hook and see if that helps in any way to update the UI.

// ---- Library --- //
const useState = (initialState) => {
  console.log('useState is initialized with value:', initialState);
  let state = initialState;
  const setState = (newState) => {
    console.log('setState is called with newState value:', newState);
    state = newState;
    // Render the UI fresh given state has changed.
    reRender();
  };
  return [state, setState];
};

Well, the behaviour, we will see with this is really not what we expected when we change the state.

reRender resets the state and add more UI elements every time state changes

let’s see what’s happening here in detail.

• Firstly, every time the input value or name state, in this case, changes, reRender is getting called. But immediately, during reRender which is glorified render , resets state with useState again with the previous value(initial value) and the new value is lost during the render.
• Secondly, because render function simply appends the DOM Tree to the root node, it's not updating the existing UI, it's adding new UI for <App/> every time reRender is called and that’s what we see here, the same UI being repeated and getting added again on every reRender .

So, how do we go about tackling them? Both are very genuine and serious problems. In this section, let’s handle the second problem first as it relates to the reRender in the next section, we will solve for resetting the state.

Fixing the reRender should be easy. See if you really think about it, when we render for the first time what we have is a blank DOM without any existing UI on our root DOM node myapp . So, during reRender if we simply clean up the UI to make it blank again before the actual render part of reRender , we should effectively be replacing the UI and not adding more and more duplicate elements. Let’s see what I mean here,

// ---- Library --- //
const reRender = () => {
  console.log('reRender-ing :)');
  const rootNode = document.getElementById('myapp');
  // reset/clean whatever is rendered already
  rootNode.innerHTML = '';
  // then render Fresh
  render(<App />, rootNode);
};

With this simple change, we have fixed the duplication of UI altogether. Here is how it looks now when changing state now.

Well, it indeed solves the second problem and reRender() function is now all good. But, our objective of changing the state and getting that reflected on the UI is still not happening because every time render is called it resets all-state back to initialValue . Interesting problem, isn’t it, let’s solve for this in next section.

🍔 State-fulness and Global State Management 🍟

Fundamentally, every time useState is called, we are indeed resetting the state without considering it is already in a modified state or not. This is the root cause of why in reRender cycle we simply overwrite what already exists.

// in useState() function
let state = initialState;

So, let’s try to keep this state outside of useState and add a check if the state has been modified which needs to be preserved and not overwritten with initialState on reRender cycles. We will start off simple, just have one variable outside of useState keeping track of this. (Why outside, because whatever is inside useState will be reset again when its called during reRender() )

// ---- Library --- //
let myAppState;
const useState = (initialState) => {
  // Check before setting AppState to initialState (reRender)
  myAppState = myAppState || initialState;
  console.log('useState is initialized with value:', myAppState);
  const setState = (newState) => {
    console.log('setState is called with newState value:', newState);
    myAppState = newState;
    // Render the UI fresh given state has changed.
    reRender();
  };
  return [myAppState, setState];
};

With this change, let’s see how our App behaves on input change.

And Voila! We have our updated state reflected everywhere on the UI. It starts off with an initialState and then as things change or state changes, it persists that state changes during reRender cycles and updates the UI every time that happens. Congratulations!! 🍨 🍨 now you have a React Library of your own which does state management and provides two-way binding on the UI. Yay!

But we are not done yet! we wanna go further, this is all good when you just have one state to manage. But, we as library author don’t know how many such states will be needed by the actual application user. And we should be able to handle arbitrary number of states in an App independently of each other having their own update and render cycles. So, how do we solve for that. First let’s try to see if we try to cramp more states in our current code what’s going to happen.

Say, in our App we now want to have two pieces of state, one name and other may be a counter which increments or decrements on a button click. Let’s quickly write out the Component code which might look like this,

// ---- Application --- //
const App = () => {
  const [name, setName] = useState('Arindam');
  const [count, setCount] = useState(0);
  return (
    <div draggable>
      <h2>Hello {name}!</h2>
      <p>I am a pargraph</p>
      <input
        type="text"
        value={name}
        onchange={(e) => setName(e.target.value)}
      />
      <h2> Counter value: {count}</h2>
      <button onclick={() => setCount(count + 1)}>+1</button>
      <button onclick={() => setCount(count - 1)}>-1</button>
    </div>
  );
};

If we simply try run our App in current state let’s see what happens,

Sure enough, useState is called twice for two states. But, both ultimately operate on the same global state myAppState and they way we have set things up, second time with userState is trying to initialize it sees there is already a value for myAppState and doesn’t reset it because of everything on reRender we learned so far. So, at this point pretty much both the states are pointing to the same global state and it gets really weird when we start interacting like if I click+1 on the counter now, this is what the result looks like.

So, the +1 button tried to do setCount(count+1) where the count is not a number but a string Arindam which means 1 will be concatenated by the + operation and as everyone is sharing the same state effectively, it gets reflected everywhere! ⚠️ ⚠️

Now let’s see how we can tackle this problem.

👨🏻‍🎤 Managing Multiple States with useState() 😵‍💫

First of all, we have to acknowledge some reality about the problem at hand for multiple states,

• We as library Authors don’t know how many states will be needed and where
• We definitely can not overwrite someone’s state with some other state while operating on them.

For the first concern, we need to have a way where every time, useState() is used in our component, we need to vend out a completely different global state and keep track of its reference within useState() which will be used to update the same. If we can manage this, the second problem will automatically be solved as there is no overwrite happening anymore on each other toes.

One simple Data Structure we can think of for the above requirement can be a simple global state array where each element of the array is managing a unique state for the component. And a state cursor or a pointer that points to the correct index in the state array for the current state variable. Visually, this might look like this.

Alright, let’s start working on setting up our state management with this context in mind.

// ---- Library --- //
const myAppState = [];
let myAppStateCursor = 0;

Now, every time, useState() is called we will create a state on the myAppState array at index myAppStateCursor and also keep the cursor value locally so that we can refer to the same cursor value for the same state.

Lastly, we will increment the myAppStateCursor to prepare it for the next useState() call which will point to the next location on the global myAppState array for the next piece of state.

Let’s see how our useState() looks with these changes,

// ---- Library --- //
const useState = (initialState) => {
  // get the cursor for this useState
  const stateCursor = myAppStateCursor;
  // Check before setting AppState to initialState (reRender)
  myAppState[stateCursor] = myAppState[stateCursor] || initialState;
  console.log(
    `useState is initialized at cursor ${stateCursor} with value:`,
    myAppState
  );
  const setState = (newState) => {
    console.log(
      `setState is called at cursor ${stateCursor} with newState value:`,
      newState
    );
    myAppState[stateCursor] = newState;
    // Render the UI fresh given state has changed.
    reRender();
  };
  // prepare the cursor for the next state.
  myAppStateCursor++;
  console.log(`stateDump`, myAppState);
  return [myAppState[stateCursor], setState];
};

Let’s see if this makes things work, with different states for the App which are independently updating the UI. Sure enough, we see an independent initialState for both of our states.

I am also printing the cursor value here to show you something important (let’s also dump the global state in useState for debugging). Let’s see what happens when we try to increment the count or maybe try to change the name state variable.

1. Increment the counter (clicked on the +1 button)
2. Change the name (Arindam -> Arindam Paul)

global state is drifting away with every re-render cycle

The Observation here is that after a setState is called when the reRender is happening, it’s basically adding new states on the next locations of the array and not using the previous cursor position. So, here we reRender two times and we see six state variables in the global state array. And in reality, our setState() calls have modified the values but useState() is not using them after the reRender as it points to a different location now on the array for the same state.

This is happening, because every time we call useState we are incrementing the global cursor myAppStateCursor to make room for the next piece of state for our App. But, no one is resetting them back during areRender() for them to start using the same references again, after any state update. So, we will do this resetting ourselves, by simply setting myAppStateCursor to its initial value 0 in our reRender() function before it actually starts rendering again. Let’s see how our reRender function is modified to reset the cursor.

// ---- Library --- //
const reRender = () => {
  console.log('reRender-ing :)');
  const rootNode = document.getElementById('myapp');
  // reset/clean whatever is rendered already
  rootNode.innerHTML = '';
  // Reset the global state cursor
  myAppStateCursor = 0;
  // then render Fresh
  render(<App />, rootNode);
};

And the MAGIC happens! our App refers to the same state references between renders and every useState() vends out a completely separate piece of state with its own setter and reRender cycle. Woo! hoo! see both states being updated in action here. (You can also see the global state after every state change for reference)

✌🏽 Why the rules of React 🤞🏽

With the above context, you can appreciate the inner workings of hooks and why some of the critical rules actually apply, see this from React documentation

Rules of Hooks

Because we have this global Array of states for a component and every useState() call specifically refers to a particular index in this state Array, you can now imagine, if we start using hooks like useState inside, loops and conditionals which invariably branches out without predictability so all of a sudden all the cursor pointers pointing to their corresponding state will be messed up if a condition skips a useState call between renders or a loop introduces more useState call than originally intended.

It becomes so obvious to see and really appreciate the design once you have an understanding ground up like this.

🚵🏼‍♀️ Conclusion 🏊🏼‍♀️

We have come a long way in building a library like React which does real useful things for us, in the first part we show how it abstracts the creation of actual UI elements to a VirtualDOM, processes it through its own render function to render it out on the UI. In this part, we kind of took off from there, and start adding dynamic states to our app which handles the UI updates when the state changes automatically on state change. For this, we introduced our own famous reRender function which does a clean up on UI and on the global state before it renders the App again. Lastly, we saw how to manage multiple states using Hooks.

If you have made it this far, I really wanna congratulate you 👌🏽👌🏽, give yourself some kudos! you have done an amazing job! 👏🏽

But, our journey continues to be more and more exciting things on the road ahead of us. In our next part, we will learn the current hot topic of React Suspense and Concurrent Mode which allows the apps to have async nature and can load components, resources dynamically when needed making our Apps much faster and slick. Stay tuned! Happy Coding!

Code for this section of our discussion can be found here.