use-context-selector demystified

In a previous article about React context performance, I mentionned the library use-context-selector that allows you to avoid useless re-render.

Today, I will refresh your memory by putting an example how to use the library. Then, I will explain how it works under the hood, you will see that it's amazing :)

A quick example

use-context-selector exposes:

createContext: a function to create a React context (yep like the React one). You can pass an optional initial value.
useContextSelector: a hook to get data from the context. It takes as first parameter the created context, and as second parameter a selector, if an identity function is passed (i.e. v => v), you will watch all changes of the context.
useContext: a hook to be notified of all changes made in the context (like the React one).

Note: In reality, the lib exposes also: useContextUpdate, BridgeProvider and useBridgeValue that I don't gonna talk about in this article.

Then you used it:

import {
  createContext,
  useContextSelector,
} from "use-context-selector";

const MyContext = createContext();

function MyProvider({ children }) {
  const [value, setValue] = useState("Initial value");

  return (
    <MyContext.Provider value={{ value, setValue }}>
      {children}
    </MyContext.Provider>
  );
}

function ComponentUsingOnlySetter() {
  const setValue = useContextSelector(
    MyContext,
    (state) => state.setValue
  );

  return (
    <button
      type="button"
      onClick={() => setValue("Another value")}
    >
      Change value
    </button>
  );
}

function ComponentUsingOnlyValue() {
  const value = useContextSelector(
    MyContext,
    (state) => state.value
  );

  return <p>The value is: {value}</p>;
}

function App() {
  return (
    <MyProvider>
      <ComponentUsingOnlySetter />
      <ComponentUsingOnlyValue />
    </MyProvider>
  );
}

As you can see it's as simple than using context with the React API.

But unlike the previous example, I would advise you to make a custom hook to select from the context not to make leak the context in all your application and to have an easy API without having to always pass the context:

import {
  createContext,
  useContextSelector,
} from "use-context-selector";

const MyContext = createContext();

const useMyContext = (selector) =>
  useContextSelector(MyContext, selector);

// I just rewrite this component,
// but it will be the same for the other one
function ComponentUsingOnlyValue() {
  const value = useMyContext((state) => state.value);

  return <p>The value is: {value}</p>;
}

Warning: Contrary to the React API, you don't have access to a Consumer component from the context. The Consumer can be useful when you have class components (and not functional component), in this case I recommend you to make an HOC that will use the useContextSelector. Or migrate to functional components :)

Ok, now you've just seen how to use it let's deep dive in the implementation.

Under the hood

We want to override the behavior which trigger a re-render of all Consumers when the data changes in the context.
So we are going to implement our own system of subscription / notify, where:

Consumers register to a custom Provider.
The custom Provider notifies Consumers where there are data changes.
The listener (in each Consumer) will recalculate the selected value and compare it to the previous one and trigger a render if it's not the same (thanks to useState or useReducer).

We are going to use a Provider to be able to register, and to put also the current data.
As you can imagine, you have to put them in an object with a stable reference and mutate this object.

Context creation

Let's implement the function to create the context named createContext. This method will just:

create a React context thanks to the react API.
remove the Consumer component from it.
override the Provider by our own implementation.

import { createContext as createContextOriginal } from "react";

function createContext(defaultValue) {
  // We are going to see next how to store the defaultValue
  const context = createContextOriginal();

  delete context.Consumer;

  // Override the Provider by our own implem
  // We are going next to implement the `createProvider` function
  context.Provider = createProvider(context.Provider);

  return context;
}

Registration system

We are going to implement the following pattern:

Let's get starting by implementing the createProvider function:

import { useRef } from "react";

function createProvider(ProviderOriginal) {
  return ({ value, children }) => {
    // Keep the current value in a ref
    const valueRef = useRef(value);
    // Keep the listeners in a Set
    // For those who doesn't know Set
    // You can compare it to Array
    // But only store unique value/reference
    // And give a nice API: add, delete, ...
    const listenersRef = useRef(new Set());
    // We don't want the context reference to change
    // So let's store it in a ref
    const contextValue = useRef({
      value: valueRef,
      // Callback to register a listener
      registerListener: (listener) => {
        // Add the listener in the Set of listeners
        listenersRef.current.add(listener);
        // Return a callback to unregister/remove the listener
        return () => listenersRef.current.delete(listener);
      },
      listeners: new Set(),
    });

    useEffect(() => {
      // Each time the value change let's:
      // - change the valueRef
      // - notify all listeners of the new value
      valueRef.current = value;
      listenersRef.current.forEach((listener) => {
        listener(value);
      });
    }, [value]);

    return (
      <ProviderOriginal value={contextValue.current}>
        {children}
      </ProviderOriginal>
    );
  };
}

And the useContextSelector and its listener is:

import { useContext, useEffect } from "react";

export default function useContextSelector(
  context,
  selector
) {
  const { value, registerListener } = useContext(context);

  // In the next part we will how to really implement this
  const selectedValue = selector(value);

  useEffect(() => {
    const updateValueIfNeeded = (newValue) => {
      // We are going to implement the logistic in the next part
    };

    const unregisterListener = registerListener(
      updateValueIfNeeded
    );

    return unregisterListener;
  }, [registerListener, value]);

  return selectedValue;
}

Now, we have a subscription / notification working. We can now focus on the implementation of the listener named here updateValueIfNeeded.

Listener implementation

The purpose of the listener is to calculate the new selected value and to return it.
To achieve this, we will use a state. But in the real implementation they use a reducer because they handle many things that I don't in my implementation, for example: version of the state, it manages when the parent renders and there is changes made in the context value that has not been yet notify to consumers.

The useContextSelector becomes:

import {
  useContext,
  useEffect,
  useRef,
  useState,
} from "react";

export default function useContextSelector(
  context,
  selector
) {
  const { value, registerListener } = useContext(context);
  // We use a state to store the selectedValue
  // It will re-render only if the value changes
  // As you may notice, I lazily initialize the value
  const [selectedValue, setSelectedValue] = useState(() =>
    selector(value)
  );
  const selectorRef = useRef(selector);

  useEffect(() => {
    // Store the selector function at each render
    // Because maybe the function has changed
    selectorRef.current = selector;
  });

  useEffect(() => {
    const updateValueIfNeeded = (newValue) => {
      // Calculate the new selectedValue
      const newSelectedValue =
        selectorRef.current(newValue);
      // Always update the value
      // React will only re-render if the reference has changed
      // Use the callback to be able to select callback too
      // Otherwise it will the selected callback
      setSelectedValue(() => newSelectedValue);
    };

    const unregisterListener = registerListener(
      updateValueIfNeeded
    );

    return unregisterListener;
  }, [registerListener, value]);

  return selectedValue;
}

Default value of context

Remember, I don't have handle the default value when creating the context. Now that we know what the format of the object stored in the context, we can do it:

import { createContext as createContextOriginal } from "react";

function createContext(defaultValue) {
  // Just put the defaultValue
  // And put a noop register function
  const context = createContextOriginal({
    value: {
      current: defaultValue,
    },
    register: () => {
      return () => {};
    }
  });

  delete context.Consumer;

  // Override the Provider by our own implem
  // We are going next to implement the `createProvider` function
  context.Provider = createProvider(context.Provider);

  return context;
}

And here we go with a simplified re-implementation of use-context-selector.

Conclusion

Looking to implementation of libraries is really something that I enjoyed because it allows you to discover the magic that is hidden.
In this case it's the implementation of a subscription / notification pattern. This pattern is also present in the react-redux implementation for performance purposes.
The library already handles the concurrent mode thanks to useContextUpdate.
By the way, Daishi Kato (the creator of many libs including this one) made a talk at the React conf 2021 to manages concurrent mode in state libraries that I found great.

Last but not least, here is a little codesandbox with my implementation if you want to play with it:

Do not hesitate to comment and if you want to see more, you can follow me on Twitter or go to my Website.