User-land implementation of useContextSelector (no dependencies) to avoid Context re-renders.

Pablo Garcia
5 min readJul 16, 2022

--

When working with the React context API, you might have noticed that it is difficult to have state updates without re-rendering every single component that is listening to the context. To fix this, we want to expose functions (methods/getters) only, so that the value of the context never changes. That is what I tried to achieve with this implementation.

I try to keep my posts short, but this is the type of hook that requires a bit more explanation.

TL;DR;

There is a Code Sandbox at the bottom if you want to copy paste the solution or see it working.

Disclosure

  1. You might be able to fix this by simply moving the state into its own provider(s). Having an individual provider for the state will allow us to re-render only the components that are interested in state changes. Note, if you have many listening components (i.e. a very wide branch in the rendering tree), then splitting is not ideal. For example, an infinite scroll app where every component in the list requires a different property from the state.
  2. There are already other implementations, such as: https://www.npmjs.com/package/@fluentui/react-context-selector, but it is not completely “user-land” because it uses calculateChangedBits, an undocumented feature to prevent React from re-rendering if a value is changed. Also, there is https://www.npmjs.com/package/use-context-selector, which I think is a better alternative to the former package, but it uses scheduler internally and I just wanted to try without using any dependencies.
  3. I do not like the amount of code I used to get this working, but it was necessary.
  4. Lastly, but most important, there is work being done around this: https://github.com/reactwg/react-18/discussions/73.

Code

First, the helper hooks:

  1. useRefVariable it simply keeps a forever-updated reference by assigning the value on every render. You can read more here.
  2. useStateRef works like useState the difference being it doesn’t trigger a re-render and it takes in a callback to notify the component of updates to the state. Internally, it uses Object.is() to skip updates if the state didn’t change.
  3. useListenersRef uses useStateRef to create a list of listeners and it notifies the listeners when the state changes.
/**
* Keeps a reference and it is always updated on every render.
*/
function useRefVariable(value) {
const ref = useRef();
ref.current = value;
return ref;
}
/**
* Same as useState, but it doesn't trigger a re-render when the state changes,
* instead it notifies the `onChange` callback.
*
* This is specially useful when you are only interested on parts of the state.
* For example, when working with the Context API, to prevent re-renders of all
* the components listening to state changes.
*
* @param {*} init
* @param {Function} onChange callback executed when the state changes
*/
function useStateRef(init, onChange) {
const onChangeRef = useRefVariable(onChange);
const stateRef = useRef(init);
const setState = useCallback(
(value) => {
// Allow passing functions when needed
if (value instanceof Function) {
value = value(stateRef.current);
}
// Skip updates if it didn't change
if (Object.is(value, stateRef.current)) {
return;
}
// Update reference
stateRef.current = value;
// Notify of changes
if (onChangeRef.current instanceof Function) {
onChangeRef.current(stateRef.current);
}
},
[stateRef, onChangeRef]
);
return [stateRef, setState];
}
/**
* Creates a list of listeners that we can use to subscribe, unsubscribe,
* and notify.
*
* It uses useStateRef internally to prevent re-renders when adding/removing
* listeners.
*/
function useListenersRef() {
const state = useStateRef([]);
const [listenersRef, setListeners] = state; const removeListener = useCallback(
(listener) => {
setListeners((listeners) =>
listeners.filter((l) => !Object.is(l, listener))
);
},
[setListeners]
);
const addListener = useCallback(
(listener) => {
setListeners((listeners) => [...listeners, listener]);
// Allow removing listeners
return function cleanup() {
removeListener(listener);
};
},
[setListeners, removeListener]
);
const notifyListeners = useCallback(
(message) => {
listenersRef.current.forEach((listener) => listener(message));
},
[listenersRef]
);
return {
listenersRef,
addListener,
removeListener,
notifyListeners,
};
}

With these helpers, we can now build the useContextSelector API (useContextSelector & useContextStateRef):

import { useContext, useEffect, useState } from "react";
import { useRefVariable } from "./useRefVariable";
import { useListenersRef } from "./useListenersRef";
import { useStateRef } from "./useStateRef";
/**
* Used in conjunction with useContextStateRef, this function subscribes to
* state changes based on what selector returns and notifies the components
* when changes occur.
* @param {Context} Context React Context (i.e. React.createContext())
* @param {Function} selector used to select parts of the context.
*/
function useContextSelector(Context, selector) {
// Context value should never change, ever!
const context = useContext(Context);
// Keep track of updates to the selected state
const [selectedState, setSelectedState] = useState(() => selector(context));
// Selector Helpers
const selectorRef = useRefVariable(selector);
const prevStateRef = useRefVariable(selectedState);
// Add listener for context state updates
useEffect(() => {
const cleanup = context.onStateChange(() => {
const value = selectorRef.current(context);
if (!Object.is(prevStateRef.current, value)) {
setSelectedState(() => value);
}
});
// Remove listener
return () => cleanup();
}, [context, selectorRef, prevStateRef, setSelectedState]);
return selectedState;
}
/**
* Hook used in conjunction with useContextSelector. It creates a state for the
* context API that is consumed by useContextSelector (via subscribers/listeners)
* without triggering a state update.
* @param {*} init
*/
function useContextStateRef(init) {
// Allow adding state-change listeners
const { addListener, notifyListeners } = useListenersRef();
// State doesn't trigger a re-render.
const [stateRef, setState] = useStateRef(init, (state) => {
// Trigger listeners when state changes
notifyListeners(state);
});
return [stateRef, setState, addListener];
}

useContextStateRef works as follows:

  1. uses useStateRef to hold the state of the context and uses useListenersRef to allow useContextSelector to subscribe to state changes.
  2. We must pass the onStateChange for useContextSelector to subscribe to changes, and we must pass stateRef for consumers to obtain the state.
  3. Note: we could refactor this a bit because useContextSelector receives the state updates from useStateRef as parameters for onStateChange.

useContextSelector works as follows:

  1. It takes in the React Context as the first parameter and a selector as the second parameter. The selector is used to select the part of the context that the component is interested in.
  2. It is the hook that will maintain the state for each of the components, so it uses useState (not useStateRef) to be able to re-render the components listening.
  3. It keeps a reference (useRefVariable) to the selector and the selectedState so that we don’t have to run the effect if the selector is an arrow function or if the state changes. Read more here.
  4. It has an effect where it listens to changes from the context. This effect should run very few times because all the dependencies are static.

Usage

import { createContext, useCallback, useMemo } from "react";
import { useContextSelector } from "./useContextSelector";
import { useContextStateRef } from "./useContextStateRef";
const CounterContext = createContext();export const useCounterContextSelector = (selector) =>
useContextSelector(CounterContext, selector);
export function CounterProvider({ children }) {
const [stateRef, setState, onStateChange] = useContextStateRef({
foo: 'bar'
});
const someMethod = useCallback(
() => setState((s) => ({ foo: '...' })),
[setState]
);
// Value must never change so the consumers do not get re-rendered.
const value = useMemo(
() => ({ stateRef, onStateChange, someMethod }),
[stateRef, onStateChange, someMethod]
);
return (
<CounterContext.Provider value={value}>
{children}
</CounterContext.Provider>
);
}

And a component would look as follows:

import { useContextSelector } from './CounterContext'function Component() {
const foo = useContextSelector(({ stateRef }) => stateRef.current.foo)
return (
<>{foo}</>
)
}

Playground

--

--

Pablo Garcia

Senior Engineer at Netflix, ex-Staff Architect 2 at PayPal. M.S. in Computer Science w/specialization in Computing Systems. B.Eng. in Computer Software.