DEV Community πŸ‘©β€πŸ’»πŸ‘¨β€πŸ’»

Lev Izraelit
Lev Izraelit

Posted on

Implementing the useState Hook

Introduction

I felt uneasy the first time I read about hooks in React. Their inner workings seemed too magical. I remember looking at a simple example and trying to make sense of how it worked under the hood:

function Counter() {
  const [count, setCount] = useState(0)
  return (
    <div>
        The count is: {count}
        <button onClick={() => setCount(count + 1)}>+</button>
    </div>
  )
}
Enter fullscreen mode Exit fullscreen mode

It was clear what the example was doing. You click the + button, and the count is incremented. But where was the value of count being stored, and how were we getting the correct value, even though 0 was passed every time? Even as I started incorporating hooks into my apps, I had few clear answers. So I started searching for sources that described how hooks work under the hood. Finally, I decided to try and reimplement some of the core hooks myself.

This post details my process of reimplementing the useState hook. For me, the goal was never to exactly match the real implementation. The goal was to gain some insight into how some like useState can be implemented.

Classes and state

Generally speaking, state includes any value that changes over time, when that value needs to be remembered by the program. For React class components, the concept of state is translated directly into the state object. The idea is to encapsulate all (or at least most) of the changing values in one place. We initialized the state object with some default values when the class is created, and then modify these values indirectly by calling the setState method:

class Counter extends React.Component {
    constructor() {
      this.state = {
        count: 0
      }
    }

    increment = () => {
      this.setState({
        count: this.state.count + 1
      }) 
    }

    render() {
      return (
        <>
          <div>count: {this.state.count}</div>
          <button onClick={this.increment}>+</button>
        </>
      )
    }
}
Enter fullscreen mode Exit fullscreen mode

The setState method recreates the component's state by merging the existing state with the new object that was passed as an argument. If we were to implement the base setState, it would look something like this:

  setState(newPartialState) {
    this.state = {
      ...this.state,
      ...newPartialState
    }
    // rerender the component
  }
Enter fullscreen mode Exit fullscreen mode

Functions and State

Unlike an object or class, a function cannot internally maintain state. This is the reason, in React, that a functional component is also called a stateless functional component. So I'd come to expect a functional component to work the same way as a simple add function - given the same input, I would expect to always get the same output. If I needed state, I would have to create a parent class component, and have that component pass down the state:

// The Counter functional component will receive 
// the count and a setCount function 
// from a parent class component
const Counter = ({ count, setCount }) => (
  <>
    <div>count: {count}</div>
    <button onClick={() => setCount(count + 1)}>+</button>
  </>
)

class CounterContainer extends React.Component {
  // shorthand for using a constructor
  state = {
    count: 0
  }

  setCount = (newCount) => {
    this.setState({
      count: newCount
    }) 
  }

  render() {
    return (
      <Counter count={this.state.count} setCount={this.setCount}>
    )
  }
}
Enter fullscreen mode Exit fullscreen mode

In a sense, the useState hook gives us a way to tell React that we need something like that parent class component, without having to create it ourselves. We simply tell React that we want to use state, and React will create that state for us.

Functions that use state

As a first attempt around creating a parent class component, we could try and have a function component directly modify a global variable:

let count = 0;

const Counter = () => (
  <>
    <div>{count}</div>
    <button onClick={() => count++}>+</button>
  </>
)
Enter fullscreen mode Exit fullscreen mode

This, however, doesn't quite work. Even though value of count is changing, the Counter component does not re-render to show the new value. We stil need something similar to a setState call, which would rerender the component when the value of count changes. We can make a setCount function that does just that:

let count = 0

function setCount(newCount) {
  count = newCount
  ReactDOM.render(<Counter />)
}

const Counter = () => (
  <>
    <div>{count}</div>
    <button onClick={() => setCount(count++)}>+</button>
  </>
)
Enter fullscreen mode Exit fullscreen mode

This works! To ensure count and setCount are always used together, we can put them inside an object. Let's call this object MyReact:

const MyReact = {
  count: 0,
  setCount(newCount) {
    this.count = newCount;
    ReactDOM.render(<Counter />)
  }
}
Enter fullscreen mode Exit fullscreen mode

To make things even clearer, let's create a useCount function that returns an object with count and setCount:

  useCount() {
    return { 
      count: this.count,
      setCount: this.setCount
    }
  }
Enter fullscreen mode Exit fullscreen mode

Next, we would want to allow the caller of useCount to pass an initial value. This presented us with a problem. we only need to set the initial value on the very first time that useCount is called. On any subsequent call, we would want use the existing value of useCount. One solution is adding a stateInitialized variable. We'll initially set it to false, and set it to true on the first time that useCount is called:

  stateInitialized: false,
  useCount(initialValue) {
    if (!this.stateInitialized) {
      this.count = initialValue;
      this.stateInitialized = true;
    }
    // ...
  }
Enter fullscreen mode Exit fullscreen mode

Now that we got the basics working, we can make MyReact more general by renaming the count variable to state, and the method names to useState and setState. Also, we'll return state and setState in an array, to allow for easy renaming:

const MyReact = {
  state: null,
  stateInitialized: false,
  setState(newState) {
    this.state = newState;
    ReactDOM.render(<Counter/>, rootElement);
  },
  useState(initialValue) {
    if (!this.stateInitialized) {
      this.stateInitialized = true;
      this.state = initialValue;
    }
    return [this.state, this.setState];
  }
};

const Counter = () => {
  const [count, setCount] = MyReact.useState(0)
  // ...
}
Enter fullscreen mode Exit fullscreen mode

We can also added a render method to MyReact, and call this method instead of calling ReactDOM.render. This will allow us to save the Counter component as part of MyReact:

  // ...
  setState(newState) {
    this.state = newState;
    ReactDOM.render(<this.component/>, this.rootElement);
  },
  // ...
  render(component, rootElement) {
    this.component = component;
    this.rootElement = rootElement;
    ReactDOM.render(<this.component/>, this.rootElement);
  }
  // ..

// later 
MyReact.render(Counter)
Enter fullscreen mode Exit fullscreen mode

Multiple state variables

The next step is to enable MyReact to manage multiple variables. The first step is to replace the single state variable with an array of state variables. Now we would need some way to know, each time setState was being called, which state variable is the one that needs to change. We can achieve this by relying on the call order to useState. Take, for example, the two subsequent calls below:

const MyCounter = () => {
  const [count, setCount] = MyReact.useState(0);
  const [name, setName] = MyReact.useState("");
}
Enter fullscreen mode Exit fullscreen mode

The MyReact.useState methods would always be executed in the same order, first returning the values of count1, setCount1, and then returning the values of name, setName. This will be the case as as long as MyReact.useState is not called inside conditional block, where the condition isn't always true or false.

Now, since we have two or more state variables, each state variable will need to have a corresponding setState method. We can achieve this by using an array of objects, where object stores the state value and the corresponding setState method. We can call each of the objects a statePair and the arrays that holds them stateArray.

[{ value: count, setCount }, { value: name, setName }, ...]
Enter fullscreen mode Exit fullscreen mode

We now need a way to track which element of the array is being used at any given time. For example, having the two calls to MyReact.useState above, the first call should return the [count, setCount] and the second call should return [name, setName]. We can use a variable to track this value. Let's call this variable currentStateIndex.

The currentStateIndex will be reset to 0 whenever any setState is called. When the value of currentStateIndex becomes equal to the length of the array, we will create a new pair of state an setState.

const MyReact = {
  stateArr: [],
  currentStateIndex: 0,
  component: null,
  useState(initialValue) {
    // if we reached beyond the last element of the array
    // We will need create a new state
    if (this.currentStateIndex === this.stateArr.length) {
      const statePair = {
        value: initialValue,
        setState(newValue) {
          statePair.value = newValue;
          MyReact.currentStateIndex = 0;
          ReactDOM.render(<MyReact.component />, rootElement);
        }
      };

      this.stateArr.push(statePair);
    }
    // get the current state and setState before incrementing the index
    const currentStatePair = this.stateArr[this.currentStateIndex];
    this.currentStateIndex += 1;
    return [currentStatePair.value, currentStatePair.setState];
  },
  render(component, rootElement) {
    this.component = component;
    this.rootElement = rootElement;
    ReactDOM.render(<this.component />, this.rootElement);
  }
};
Enter fullscreen mode Exit fullscreen mode

Example

Given the above implementation, let's try and follow an example of a component that uses two state variables:

const Counter = () => {
  const [count1, setCount1] = MyReact.useState(0);
  const [count2, setCount2] = MyReact.useState(0);
  return (
    <>
      <div>
        The first count is: {count1}
        <button onClick={() => setCount1(count1 + 1)}>+</button>
      </div>
      <div>
        The second count is: {count2}
        <button onClick={() => setCount2(count2 + 1)}>+</button>
      </div>
    </>
  )
}

MyReact.render(Counter)
Enter fullscreen mode Exit fullscreen mode

Below is a sandbox with MyReact and the Counter component:

Following the example, these would be the inital values of MyReact:

MyReact {  
  stateArr: [],
  currentStateIndex: 0,
  component: null,
}
Enter fullscreen mode Exit fullscreen mode

After the first call to useState:

const Counter = () => {
  const [count1, setCount1] = MyReact.useState(0); // <--
Enter fullscreen mode Exit fullscreen mode

The values of MyReact will be:

MyReact {  
  stateArr: [{ value: 0, setState: fn() }],
  currentStateIndex: 1,
  component: Counter,
}
Enter fullscreen mode Exit fullscreen mode

After the second call to useState:

const Counter = () => {
  const [count1, setCount1] = MyReact.useState(0); 
  const [count2, setCount2] = MyReact.useState(0); // <--
Enter fullscreen mode Exit fullscreen mode

The values of MyReact will be:

MyReact {  
  stateArr: [{ value: 0, setState: fn() }, { value: 0, setState: fn() }],
  currentStateIndex: 2,
  component: Counter,
}
Enter fullscreen mode Exit fullscreen mode

Now, if the first + button is pressed, the values of MyReact would become:

MyReact {  
  stateArr: [{ value: 1, setState: fn() }, { value: 0, setState: fn() }],
  currentStateIndex: 0,
  component: Counter,
}
Enter fullscreen mode Exit fullscreen mode

Which would lead to Counter being rendered again. On the subsequent calls to useState, only the currentStateIndex will be incremented, while the existing elements of stateArr will be returned.

Conclusion

So, we've arrived at something pretty similar to React's useState hook. I cannot say if understanding the internal workings of hooks would make someone a better React developer. But I do feel that is is worthwhile to try and understand how abstractions can be created - this can help us better understand the ones that have already been made, and to make new abstractions of our own.

Top comments (4)

Collapse
fijiwebdesign profile image
Gabirieli Lalasava • Edited on

I'm trying to wrap my head around useState() working in multiple contexts. In a single context this works since the execution order is guaranteed.

useState() <-- state, setter at index 0
useState() <-- state, setter at index 1

But if you have not guaranteed execution order (many components and random renders) you can't count indexes.

How do you get useState() to keep a reference for your specific scope/context without supplying it with your context or a reference to save to.

Update:

Just read: medium.com/the-guild/under-the-hoo...

It explains useState hook pretty well. The component context is derived from the current execution context because react controls the render.
So calling useState() inside that render knows the context because react knows it's rendering that context at that time. Which is why useState() can only be used inside a function render.

Your version of using an index is correct, the implementation just uses different semantics akin to middleware:

{
  memoizedState: 'foo',
  next: {
    memoizedState: 'bar',
    next: {
      memoizedState: 'bar',
      next: null
    }
  }
}
Enter fullscreen mode Exit fullscreen mode

instead of incrementing the index it resolves the next state though the next property chain.
(I'm guessing because it's faster to traverse a single next property than a long array while updating an index - might be just choice of semantics)

Thanks for the information, hope that helps as well.

Collapse
lizraeli profile image
Lev Izraelit Author

Sorry for the delayed response. That was a good question and a great follow-up explanation. This is exactly right - each useState hook is called while React is rendering a particular component. React stores the memoized state from these calls as part of its in-memory representation of each component. Hopefully your explanation will be helpful to others who read this article in the future.

Collapse
lucastrvsn profile image
Lucas Trevisan

Thank you for this. This post need to be in the react hooks page in docs!

Collapse
shylockness profile image
shylockness

Thank you for this well written article !
It answered completely my doubts about useState hook.

🌚 Browsing with dark mode makes you a better developer.

It's a scientific fact.