The future of React.use and React.useMemo - a powerful alternative to context selectors

In this article, I'll review how I optimized an application's problematic React context using Daishi Kato's use-context-selector library, how I developed a bias towards using context selectors to prevent unnecessary re-renders, and how Dan Abramov convinced me that an abstraction I built to make context selectors composable was the wrong approach.

Why I preferred context selectors

Not long ago, I was hired by a company that was building an interactive airbnb-style map search. Their map code relied heavily on a single React context to manage its state, handlers, and effects.

The code looked something like this:

const MapSearchContext = React.createContext();
const useMapSearchContext = React.useContext(MapSearchContext);

const MapSearchProvider = ({ children }) => {
  // 👀 we stuffed most of the map logic into `useMapSearchProvider`
  const mapSearch = useMapSearchProvider();

  return (
    <MapSearchContext.Provider value={mapSearch}>
      {children}
    </MapSearchContext.Provider>
  );
};
const MapSearch = () => {
  return (
    <MapSearchProvider>
      <Container>{/* lots of questionable components here */}</Container>
    </MapSearchProvider>
  );
};

useMapSearchProvider composed several smaller custom hooks that were easy to test and, therefore, easy to change. But an experienced React developer will recognize a potential problem with this code:

Unnecessary re-renders in components that call React.useContext's wrapper, useMapSearchContext.

Why?

Consider the following code that uses a similar strategy of providing the return value of a hook to a context provider:

const Context = React.createContext();

const useFavoritePet = () => {
  const [lastNaughtyAction, setLastNaughtAction] =
    React.useState("Frank - stole food");
  const [favoritePet, setFavoritePet] = React.useState("Frank");

  const onNaughtyAction = React.useCallback(
    (petName, action) => setLastNaughtAction(`${petName} - ${action}`),
    []
  );

  useEffect(() => {
    if (lastNaughtyAction.startsWith("Frank")) {
      setFavoritePet("Coconut");
    } else {
      setFavoritePet("Frank");
    }
  }, [lastNaughtyAction]);

  return {
    favoritePet,
    lastNaughtyAction,
    onNaughtyAction,
  };
};

const MyFavoritePetProvider = ({ children }) => {
  const value = useFavoritePet();
  return <Context.Provider value={value}>{children}</Context.Provider>;
};

const MyFavoritePet = () => {
  return (
    <MyFavoritePetProvider>
      <>
        <FavoritePetPhoto />
        <PetControls />
      </>
    </MyFavoritePetProvider>
  );
};

const FavoritePetPhoto = () => {
  const { favoritePet } = React.useContext(Context);
  return <img src={`/${favoritePet}.png`} />;
};

const PetControls = () => {
  // 👀 calling `useContext` here means that this component will
  // re-render everytime `onNaughtyAction` is called since `onNaughtyAction`
  // calls `React.setState`. But this re-render is unnecessary since
  // `onNaughtyAction` never changes.
  const { onNaughtyAction } = React.useContext(Context);
  return (
    <>
      <button
        onClick={() => onNaughtyAction("Coconut", "scratches girlfriend")}
      >
        Coconut scratches girlfriend
      </button>
      <button onClick={() => onNaughtyAction("Frank", "chases coconut")}>
        Frank chases coconut
      </button>
    </>
  );
};

In the above code, <PetControls /> is going to re-render anytime onNaughtyAction is called since onNaughtyAction causes a state change in useFavoritePet. Any state change in the hook supplied to Context.Provider will cause useFavoritePet, and thereby React.useContext(Context), to return a new reference, and deconstructing a changed reference inside of a React component will trigger a re-render.

So circling back to the map code, if we write lots of state changing code inside of useMapSearchProvider, we can expect lots of unnecessary re-renders in components that only need a small slice of useMapSearchContext's return value, since useMapSearchContext's returns a new reference with each state change.

Over time this context-based design did, in fact, lead to performance problems.

Optimizing React.useContext

I was tasked with optimizing their code, but also instructed to do so in the smallest way possible, since this was their most important feature. I considered breaking up their single context into several smaller contexts, but I struggled to break up useMapSearchProvider into separate hooks. Instead, in the spirit of keeping my changes concise, I researched ways to make a single context more performant.

My research led me to RFC 119, a React proposal created in 2019, which outlined a potential solution to prevent unnecessary re-renders when calling React.useContext. The idea was to add support for passing a selector function as the second parameter likeso:

const thingICareAbout = React.useContext(Ctx, (ctx) => ctx.thingICareAbout);

React's internals would take care to ensure that if anything not "selected" from the ctx changed, re-renders wouldn't happen.

But the RFC was still in an "open" state and the React core team had no official plans to implement it. Nonetheless, the Github discussion led me to a talented developer named Daishi Kato, who was nice enough to open source a similar solution in the user space called use-context-selector.

Using Daishi's library plus some of my own code, I was able to implement a selector as the first argument to useMapSearchContext, which would ensure that our map's components would only re-render when the slice of the context they were interested in changed likeso:

// This will only trigger a re-render if `state` or `city` change.
const state = useMapSearchContext((ctx) => ctx.state);
const city = useMapSearchContext((ctx) => ctx.city);

Several small PRs later, my optimization work was complete. Components throughout the map's render tree were no longer re-rendering unless the thing they needed from useMapSearchContext's return value had actually changed.

And that's the origin story of how I came to prefer the context selector pattern. I adopted it in personal projects, and even wrote an experimental library that extended Daishi's use-context-selector. I loved that I could use just React contexts and hooks for most, if not all, of my global state management needs.

That is, until a recent interaction with Dan Abramov on Twitter.

A better alternative to context selectors

It all started with this tweet by react-query's maintainer, @TkDoko (Dominik).

source tweet

Dominik was highlighting a future feature coming to React that would make it possible to call React.use(Context) (currently only available in canary) inside of a useMemo function. And interestingly, despite nothing in his example's useMemo dependency array, useMemo would somehow know to trigger a rerender only if foo changed. The return value in this scenario acts as an implicit dependency in useMemo's dependency array 🤯. To make things even more interesting, whenever all of this is realized in a future version of React, this will work for third party hooks too.

But I still had an icky feeling about further relying on useMemo to keep applications performant. I'd witnessed many developers (myself included) shoot themselves in the foot with memoization APIs.

A stubborn loyalty to context selectors

I happened to have twitter open when Dominik shared the above screenshot and decided to reply with a question:

source tweet

Dominik quickly pointed out that the added abstraction was probably not a good idea and that my example left the selector out of the useMemo dependency array (oops!).

His exact response:

"no, the selector would need to go into the dependency array, because it could close over things. At that point, you'd need to memoize that one as well. I don't think it makes too much sense to abstract this further, I'd just inline usages" - @TkDodo

Before I had time to fully process what he'd written, Dan Abramov chimed in with an even juicier response:

source tweet

And a few seconds later he followed up with another post to provide a bit of reasoning:

source tweet

Dan's tweet here is referring to the fact that a classic useContextSelector implementation will take a selector function with the following type signature as its second parameter:

type SelectorFunction = (Context) => any;

And that any selector must know about the entire structure of Context to make its selection. Let's demonstrate Dan's point by implementing his previous useColor/useTheme example with useContextSelector:

useTheme() {
  return useContextSelector(Context, ctx => ctx.theme)
}

useColor() {
  return useContextSelector(Context, ctx => ctx.theme.color)
}

This is what Dan meant when he wrote "you have to know precisely what you’re selecting in advance". In the implementation above, the author of useColor must know that the color value exists at ctx.theme.color. And if the author tries to compose useColor from the return value of useTheme likeso:

useTheme() {
  return useContextSelector(Context, ctx => ctx.theme)
}

// ❌ completely defeats the purpose of using a selector function to prevent
// re-renders since `useColor` will now re-render when any portion of `ctx.theme`
// changes!!
useColor() {
  return useTheme().color
}

useColor will now trigger a re-render when unrelated portions of the ctx.theme change, despite only needing access to ctx.theme.color.

Dan was right. A lack of composability was a major shortcoming of the selector approach outlined in RFC 119.

But I couldn't shake the feeling that if composability was, in fact, the primary concern with using context selectors, then surely an additional abstraction could fix the problem. I wasn't ready to concede that direct usage of React.useMemo was the answer.

Making a composable useContextSelector

I had some free time and decided to spend it writing some code to address Dan's critique of useContextSelector's lack of composability. The result: https://github.com/InterBolt/selectable.

Let's use it to recreate the useColor/useTheme example from above:

import selectable from "@interbolt/selectable";

const useContextSelector = selectable(Context);
const useThemeSelector = useContextSelector.narrow((ctx) => ctx.theme);
// COMPOSABLE!!!
const useColorSelector = useThemeSelector.narrow((theme) => theme.color);

// and we can use them likeso:
const fonts = useThemeSelector((theme) => theme.fonts);
const color = useColorSelector();

The selectable API above creates a base context selector hook and ensures that any hook created via the attached narrow method returns a new hook whose selector param only needs to know about the "narrowed" portion of the context.

The selectable abstraction also handles composing any number of hooks and contexts together when further "narrowing" a context:

import selectable from "@interbolt/selectable";

const useContextSelector = selectable(Context);
const useTheme = useContextSelector.narrow((ctx) => ctx.theme);

// Can narrow with `useUserHook` and `SomeOtherContext` too
const useDerivedSelector = useTheme.narrow(
  useUserHook,
  SomeOtherContext,
  (userHookReturnValue, someOtherContext, theme) =>
    someSelectionFunction(theme, someOtherContext, userHookReturnValue)
);

For the curious, here's the source code for selectable:

Expand code
import React from "react";
import useFutureMemoShim from "./useFutureMemoShim.jsx";
import useFutureShim from "./useFutureShim.jsx";

const getIsCtx = (val) => {
  return (
    typeof val === "object" && val["$$typeof"] === Symbol.for("react.context")
  );
};

const buildHook = (ctxOrHook) => {
  if (getIsCtx(ctxOrHook)) {
    return () => useFutureShim(ctxOrHook);
  }
  return ctxOrHook;
};

const selectable = (rootHookOrCtx = null) => {
  const rootHook = buildHook(rootHookOrCtx);

  const nextNarrow =
    (accumNarrowHooks = [], accumSelectors = []) =>
    (...args) => {
      const nextSelector = args.at(-1);
      const nextNarrowHooks = args.slice(0, -1).map((a) => buildHook(a));
      const nextAccumSelectors = accumSelectors.concat([nextSelector]);
      const nextAccumNarrowHooks = accumNarrowHooks.concat([nextNarrowHooks]);

      const useSelector = (hookSelector = (a) => a) =>
        useFutureMemoShim(() => {
          const rootVal = rootHook();
          const hookOutputs = [];
          let selected = rootVal;

          nextAccumSelectors.forEach((selector, i) => {
            nextAccumNarrowHooks[i].forEach((hook) => {
              hookOutputs.push(hook());
            });
            selected = selector(...hookOutputs, selected, rootVal);
            hookOutputs.length = 0;
          });

          return hookSelector(selected, rootVal);
        }, []);

      useSelector.narrow = nextNarrow(
        nextAccumNarrowHooks,
        nextAccumSelectors
      );
      return useSelector;
    };

  return nextNarrow()((a) => a);
};

export default selectable;

Excited that I had "solved" the composability problem, I went back and shared my work on Twitter. I must say, Dan really doesn't miss a beat. My post was only a few minutes old when I got the notification that he had replied.

tweet source

I was as flattered as I was dissapointed. Here was my reply:

tweet source

My reply was an admission that I basically did all of this to avoid direct use of React.useMemo. But it wasn't until his next two comments that I fully came around to his point of view:

First:

tweet source

And then (there's a little back and forth in-between these two):

tweet source

His first tweet about the compiler is referring to the fact that writing abstractions purely to avoid useMemo is a pretty bad idea if React Forget, the compiler that Meta is testing internally, will soon inject those automatically.

Imagine for a moment that React Forget already exists and my selectable API is no longer necessary.

Writing a codemod to change this:

import selectable from "@interbolt/selectable";

const useContextSelector = selectable(Context);
const useTheme = useContextSelector.narrow((ctx) => ctx.theme);

const selectedValue = useTheme(someSelectionFunction);

to this:

const { theme } = use(Context);

const selectedValue = someSelectionFunction(theme);

is significantly more difficult than writing a codemod to change this:

React.useMemo(() => {
  const { theme } = use(Context);
  return someSelectionFunction(theme);
}, [someSelectionFunction]);

to this:

const { theme } = use(Context);
const selectedValue = someSelectionFunction(theme);

The later codemod can simply remove the useMemo and call it a day. But a codemod that gets rid of the selectable and narrow API must understand exactly how my library works.

Despite my personal grudge against useMemo, Dan Abramov is, once again, on the money.

Conclusion

All of this is to say, React is on the verge of some exciting updates that will remove the need for RFC 119 and use-context-selector-style context selectors. React.useMemo + React.use will provide intelligent re-renders based on the value returned within a useMemo function and, eventually, React Forget will remove the need for useMemos and useCallbacks altogether. If the compiler ships with a codemod, developers can eliminate tons of memoization cruft by simply running a command.

In a matter of days, I went from feeling frustrated that the React core team never merged RFC 119 to thankful that they did not. React's future is bright and I feel more confident by the day that building on top of it is the right move.

That's all folks. I do React/NextJS consulting/contracting work so feel free to email me a [email protected] with any questions or inquiries.