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.

Comments

[federazionekoinonie]

Thanks for detailed post, I found it very useful.
I tried to use the "under the hood" implementation of context-selector.
1) Is it possible that on first run (of react-native app) there always are 2 renders per consumer caused by wrong initial (default) value returned by useContextSelector? I found that first time it returns a "ref" object (with "current" property), then the expected value.
I think in function createProvider it should be:

. . .
const contextValue = useRef({
      value: valueRef.current, // <<=== ADDED .current HERE

[Romain Trotard]

First of all thanks for the read and your interest in the subject :)

1) I don't code with React native but I guess it should works the same than with in JS. Do you have a repository where I can see the code ?

[federazionekoinonie]

2) when I use a context with 2 props: theme and lang, whenever I change several times the same prop (ie. lang) only the "lang consumer" rerender as expected.
When I update one prop (ie. lang) then the other prop (ie. theme) then the app rerender both props consumers, opposite of what I was expecting: lang consumer shouldn't rerender on theme change...
If I "follow" updating theme value, only theme consumer rerender, but each time I modify "the other prop", both (all) consumers rerender...
Below code sample:

import React, { useState, useCallback } from 'react';
import { StyleSheet, View, Text, Button } from 'react-native';
import { createContext, useContextSelector } from './services/useContextSelector';

export default App;

const PreferencesContext = createContext();

const PreferencesProvider = ({children}) => {
  const [preferences, setPreferences] = useState({theme: 'light', lang: 'it'});
  const contextValue = {preferences, setPreferences};
  return (
    <PreferencesContext.Provider value={contextValue}>
      {children}
    </PreferencesContext.Provider>
  );
};

const LangConsumer = () => {
  const langSelector = useCallback(contextValue => contextValue?.preferences?.lang, []);
  const lang = useContextSelector(PreferencesContext, langSelector);
  return (
    <Text>Language: {lang}</Text>
  );
}

const ThemeConsumer = () => {
  const themeSelector = useCallback(contextValue => contextValue?.preferences?.theme, []);
  const theme = useContextSelector(PreferencesContext, themeSelector);
  return (
    <Text>Theme: {theme}</Text>
  );
}

const StubScreen = () => {
  const setPreferences = useContextSelector(PreferencesContext, contextValue => (contextValue?.current ? contextValue.current?.setPreferences : contextValue.setPreferences));
  const toggleTheme = useCallback(
    () => setPreferences(state => 
      ({ ...state, theme: (state.theme === 'light' ? 'dark' : 'light') })
    ),
    [setPreferences]
  );
  const toggleLang = useCallback(
    () => setPreferences(state => 
      ({ ...state, lang: (state.lang === 'it' ? 'en' : 'it') })
    ),
    [setPreferences]
  );

  return (
    setPreferences ?
    <View style={styles.stubView}>
      <ThemeConsumer />
      <Button onPress={toggleTheme} title="Change Theme" />
      <LangConsumer />
      <Button onPress={toggleLang} title="Change Language" />
    </View>
    :
    null
  );
}

function App() {
  return (
    <PreferencesProvider>
      <StubScreen />
    </PreferencesProvider>
  );
}

const styles = StyleSheet.create({
  stubView: { flex: 1, justifyContent: 'center', alignItems: 'center', }
});

[Romain Trotard]

Indeed, it's really strange. I have tried on snack.expo.dev/ and it has a weird behavior. If I chose the Web device it works fine but not on Android and iOS.

It seems to be due to the fact of always calling setSelectedValue in useContextSelector. Because if I condition it it's working.

I will try to deep dive more and will come back to you ;)

[Saravanan Ramupillai]

@romaintrotard Thanks for a detailed Post. Here is the point i observed, deleting Consumer component is not doing any trick. The trick is the value passed into a value of OriginalProvider. It is a ref whose value never change so the consumer below are not notified automatically. we notify them manually via own listener.

[Romain Trotard]

Glad you liked the Post :)
Yep deleting the Consumer has no effect on the implementation. Everything is based on the the Observer pattern that I named badly subscription / notification pattern in the article. This pattern is used in numerous libraries: react-redux, jotai, ...

[Marco Dinh]

Thanks for your precious works.

I made a package here if anyone is interested.
npmjs.com/package/@fishbot/context...

It's the optimal version of the React Context with Selector, which only re-renders the components that observe the changed value.
This works on both Web and Mobile.