Deriving the implementation of React Context

Usual posts about this subject just dive deep into a particular way of solving a problem. I take a different approach in this post in starting with the problem and trying to come up with a solution of my own.

So in this post, I will try to derive the implementation of React.Context with a simple example of a Counter. We will start with the problem, solve it step by step and abstract the solution and we will have derived React.Context in the process.

Let’s create a component and implement the solution in the traditional way, explaining the problems and improving them as we go along. The final solution will be what React.Context is. Note that I have not looked at the official implementation. This is just a simple attempt grasp the idea of Context.

1. Lets look at the traditional way of creating a Counter

const Result = ({ count }) => <h1>{count}</h1>;

const IncrementCounter = ({ onClick }) => (
  <button onClick={onClick}>Increment</button>
);

const DecrementCounter = ({ onClick }) => (
  <button onClick={onClick}>Decrement</button>
);

const Counter = () => {
  const [count, setCount] = useState(0);
  const onIncrementClick = () => setCount((count) => count + 1);
  const onDecrementClick = () => setCount((count) => count - 1);

  return (
    <Fragment>
      <Result count={count} />
      <IncrementCounter onClick={onIncrementClick} />
      <DecrementCounter onClick={onDecrementClick} />
    </Fragment>
  );
};

This is how we traditionally wrote a Counter component. We have a Counter component (lets call it the parent component from here on) which does the managing of the state and also the rendering the child components, passing the required props into them, all in the same place. Nothing fancy, nothing wrong. Except the business logic is baked into the parent component. Not exactly undesirable but maybe there is a better approach.

The problems with this approach

1. The view layer and the business logic are coupled together.
2. Lack of flexibility with the component API. As a consumer of the Counter, you can only use the <Counter /> component as a whole. You can’t extend its functionality.
3. Prop drilling. As a maintainer of this component, you need to pass the necessary props explicitly to the child components.

And these are the problems we are going to solve and hopefully, we will have derived React.Context by the end of it.

The component API we would like to achieve

  <Counter>
    <Counter.Result />
    <Counter.Increment />
    <Counter.Decrement />
  </Counter>

Here is the component API at a high level we ideally would like to achieve. As you can see, it is different from the traditional variant (<Counter />). It displays the value of the counter, an increment button and a decrement button and its all wrapped inside a parent component. When any of the buttons are clicked, the value of the counter changes. We want to:

1. Expose such an API so as to be able to manage the state separately (in the parent component) and keep the child components as stateless as possible. This solves the problem of decoupling the view layer and the business logic.
2. Be able to decide where to put (or rearrange) the child components or even add more stuff making them flexible.
3. Be able to not have to drill props into the child components. They just get it somehow from the parent component. Let’s see how.

Lets look at the child components

With the Component API I mentioned before in mind, let’s look at the Result, Increment and Decrement Components.

const IncrementCounter = ({ increment }) => {
  return <button onClick={increment}> Increment</button>;
};

const DecrementCounter = ({ decrement }) => {
  return <button onClick={decrement}> Decrement</button>;
};

const ShowResult = ({ count }) => {
  return <h1>{count}</h1>;
};

We want them as dumb as possible meaning they shouldn’t have any business logic of their own. They just get the count, increment and decrement as props, although note that we don’t pass them explicitly in the new API. This is important. We don’t pass these props from the parent component but they are still available as props. How do we achieve this? We can achieve this if we can manage the state in the parent component and somehow magically pass the methods to the child components. And thats exactly what we’ll do next, except its no magic.

2. The parent component

const Counter = ({ children }) => {
  const [count, setCount] = React.useState(0);

  const increment = () => setCount((count) => count + 1);
  const decrement = () => setCount((count) => count - 1);

  return React.Children.map(children, (child) => {
    return React.cloneElement(child, { increment, decrement, count });
  });
};

What are we doing here? We manage the state and the methods here, iterate through the child components, create a clone of them and pass the state values and methods as props. And that’ll do the job for us. We found a way to manage state in the parent component and wrap the state values and methods as props to the child components without having to pass them explicitly.

3. The Provider component

Now, let’s go one step ahead and call this component Provider and pass it a value prop with the state as the value. This value prop will be passed down the child components as regular props.

const Counter = ({ children }) => {
  const [count, setCount] = React.useState(0);

  const increment = () => setCount((count) => count + 1);
  const decrement = () => setCount((count) => count - 1);

  const Provider = ({ value, children }) =>
    React.Children.map(children, (child) => React.cloneElement(child, value));

  return (
    <Provider value={{ increment, decrement, count }}>{children}</Provider>
  );
};

As a side note, I would have named this component an Iterator or just Component or something because Context and Provider being used in multiple places caused a great deal of confusion to me.

4. Abstracting React.createContext

Now, let’s move one more step ahead and move the Provider component outside, to a method called createContext inside a fictitious library called MyReact and use it inside our parent component.

const MyReact = {
    createContext: function () {
      return {
        Provider: ({ children, value }) => {
          return React.Children.map(children, (child) =>
            React.cloneElement(child, value)
          );
        }
      };
    }
  };
};

const Counter = ({ children }) => {
  const { createContext } = MyReact();
  const CounterContext = createContext();
  const [count, setCount] = React.useState(0);
  const increment = () => setCount((counter) => counter + 1);
  const decrement = () => setCount((counter) => counter - 1);

  const value = {
    count,
    increment,
    decrement
  };

  return (
    <CounterContext.Provider value={value}>{children}</CounterContext.Provider>
  );
};

Okay, so far, we have seen how the parent component passes the state down all the child components as regular props. We abstracted that logic and moved it inside our dummy library, MyReact into a method called createContext. The child components then use the state values from their props.

But there is a slight twist to it. In the actual implementation, the child components don’t take the state values from the props. Instead, they get it using a method called useContext. So what might they be doing?

5. Abstracting React.useContext()

Instead of drilling the state as props, let’s store the state in some kind of store. And when the child components request it, lets just return it from the store. Let’s try to implement that logic.

const render = (children) => {
  return React.Children.map(children, (child) => {
    if (child.props.children) {
      child = React.cloneElement(child, {
        children: render(child.props.children)
      });
    }
    return React.cloneElement(child);
  });
};

const MyReact = () => {
  const state = new Map();

  return {
    createContext: function () {
      let Context = Symbol();
      Context = {
        Provider: ({ children, value }) => {
          state.set(Context, value);
          return render(children);
        }
      };
      return Context;
    },
    useContext: (context) => {
      return state.get(context);
    }
  };
};

So we made a bit of change to the createContext method. Every time it is called, it stores the state in a store and just returns a component called Provider which just renders the child components. We also added a new method called useContext which when passed the context, returns the current state.

And thats it. We have abstracted both the createContext and the useContext methods and achieved the same desired result.

End Result

With this, we can write something like this.

const Counter = (() => {
  const { createContext, useContext } = MyReact();
  const Context = createContext();

  return {
    State: ({ children }) => {
      const [count, setCount] = React.useState(0);
      const increment = () => setCount((counter) => counter + 1);
      const decrement = () => setCount((counter) => counter - 1);

      const value = {
        count,
        increment,
        decrement
      };
      return <Context.Provider value={value}>{children}</Context.Provider>;
    },
    Increment: () => {
      const { increment } = useContext(Context);
      return <button onClick={increment}>Increment</button>;
    },
    Decrement: () => {
      const { decrement } = useContext(Context);
      return <button onClick={decrement}>Decrement</button>;
    },
    Result: () => {
      const { count } = useContext(Context);
      return <h1>{count}</h1>;
    }
  };
})();

export default () => (
  <Counter.State>
    <Counter.Result />
    <Counter.Increment />
    <Counter.Decrement />
  </Counter.State>
);

Summary

1. We wrote the component the traditional way
2. We created a parent component which returns a new component with the new child component receiving state values as prop.
3. We abstracted this component and called it Provider.
4. We moved this abstracted component into createContext()
5. We created a map to store the state and created useContext() to return the state.

Conclusion

To me, it’s amazing we can write like this. It was hard to wrap my head around Context for some reason. I find the namings a bit counter intuitive. But with this, I have a decent picture of the internal workings of it. Here is the codesandbox url to explore this case study.