React doesn't need state management tool, I said

From time to time someone still tells me that is using REDUX or similar tool in their project. I usually respond, that I wouldn't use it as now with hooks and context API you don't need it.

But context API usually brings performance problems and is also a bit awkward to use it properly, so today I'll try to show how to avoid common problems, and also build your own (micro) state management tool, without any compromises.

Naive solution

Basic idea is to manage state in one component and pass the whole it by context so it's accessible from all child components, so we can avoid props drilling.

export const StateContext = createContext(null);
const Provider = () => {
  return (
    <StateContext.Provider value={state}>
      <ChildComponent />
    </StateContext.Provider>
  )
}

Using dispatch

However you also need some way how to modify the state from children, you could pass individual functions to the context, but I personally don't like that as the state will get complex very fast. I like idea of dispatching events (similarly as in REDUX), so we basically pass one function which you can use to dispatch all different actions that you need. We could pass it through the same context as the state, but I don't like mixing it with the state, so I pass it through a separate context.

const StateContext = createContext(null);
const DispatchContext = createContext(null);

export const Provider = () => {
  const [state, setState] = useState(...)

  const dispatch = (action) => {
    switch (action.type) {
      case 'CHANGE_STATE':
        setState(action.payload)
        break;
      ...
    }
  }

  return (
    <StateContext.Provider value={{state, ...}}>
      <DispatchContext.Provider value={dispatch}>
        <ChildComponent />
      </DispatchContext.Provider>
    </StateContext.Provider>
  )
}

I also like creating hook for getting the dispatch function to make it more explicit:

export const useDispatch = () => {
  return useContext(DispatchContext)
}

Basically we are separating data from actions - provider component provides data to children. Children can dispatch actions to modify the data, but it's controlled by provider component, so it has control over it. Dispatched actions can be understood similarly as e.g. dom events, except we know who will receive it.

Now let's look at the performance side as if we want to use this as a replacement of REDUX, it needs to be able to handle big states with a lot of components subscribed.

Avoiding unnecessary children re-creation

In this configuration we are really inefficient, as all the children will get re-rendered every time we change something in the state. This happens because every time we update state in Provider component, all it's children will get re-created. We could use React.memo on children to avoid this, however nicer solution is to pass children from component above, so when the Provider is updated, children will stay the same. And we only update actual context consumers.

export const Provider = ({ children }) => {

  ...

  return (
    <StateContext.Provider value={{state, ...}}>
      <DispatchContext.Provider value={dispatch}>
        {children}
      </DispatchContext.Provider>
    </StateContext.Provider>
  )
}

In parent we do:

export const Parent = ({ children }) => {
  return (
    <Provider>
      <ChildComponent />
    </Provider>
  )
}

Now the provider component is managing the context, but is not managing children (only passing them). It took me a while to understand this subtle difference as it is quite small change in the code, with big consequences.

The trick is to understand, that when we put <ChildComponent >, we are basically creating new React.Node every time, so all the children are re-rendered, unless we wrap them in React.memo.

So with this change, we update only components which are using the context.

Avoiding dispatch causing re-renders

Currently dispatch function is re-created every time the state is changed, which mean that all components using it will get re-rended, even though they are not using StateContext. Usually if we want to have stable function react documentation advices to use useCallback, but in this case it will help us only partially, because, that will basically cause "caching" of dispatch function and we wouldn't be able to use outer scope variables without including them into dependencies - and then the dispatch function would still get recreated when dependencies change. We will need to use ref to help us with this.

...

export const Provider = ({ children }) => {
  const [state, setState] = useState(...)

  const dispatchRef = useRef()

  // new function with every render
  const dispatchRef.current = (action) => {
    switch (action.type) {
      case 'CHANGE_STATE':
        // we can use outer scope without restrictions
        setState({...action.payload, ...state})
        break;
      ...
    }
  }

  // stable dispatch function
  const dispatch = useCallback(
    (action: ActionType) => dispatchRef.current(action),
    [dispatchRef]
  );

  return (
    <StateContext.Provider value={{state, ...}}>
      <DispatchContext.Provider value={dispatch}>
        {children}
      </DispatchContext.Provider>
    </StateContext.Provider>
  )
}

This way stable dispatch function is passed to the DispatchContext and we can use outer scope without limitations.

Subscribable context

Last optimization we'll need is ability of the component subscribe only to part of the state. Now components can only use whole state and even when they need just small piece (e.g. one boolean value), they'll get notified every we change the state. This is not the best practice as we would still get unnecessary re-renders. The way to solve this is through use-context-selector.

This library is quite simple and it allows to use selector function, to "pick" what we want from the state.

import { createContext } from 'use-context-selector';

const StateContext = createContext(null);

export const Provider = ({ children }) => {
  return (
    <StateContext.Provider value={{state, ...}}>
      <DispatchContext.Provider value={dispatch}>
        {children}
      </DispatchContext.Provider>
    </StateContext.Provider>
  )
}

import { useContextSelector } from 'use-context-selector';

export const Subscriber = () => {
  const somePart = useContextSelector(StateContext, context => context.somePart)
}

Oh, wait that is cheating! You said you will only use Context API!

This library is quite simple wrapper of React.Context api. It uses ref to wrap passed value, so that components are not re-rendering automatically and then it keeps list of subscribers. When value changes it runs all the subscribed functions and if the value from the selector is different than before it forces the subscribed Component to re-render. Similar concept is used e.g. in redux useSelector hook. So I say, it's quite standard solution and why build a new one, when it already exists?

PS: there is even a open RFC to add something like this directly into react

Final product

We can wrap this whole functionality to be reusable (+ add typescript types)

import React, { useCallback, useRef } from 'react';
import { createContext, useContextSelector } from 'use-context-selector';

type DispatchType<ActionType, DispatchReturn> = (
  action: ActionType
) => DispatchReturn;

type SelectorType<StateType> = (state: StateType) => any;

export const createProvider = <
  StateType,
  ActionType,
  DispatchReturn,
  ProviderProps
>(
  body: (
    props: ProviderProps
  ) => [state: StateType, dispatch: DispatchType<ActionType, DispatchReturn>]
) => {
  const StateContext = createContext<StateType>(null as any);
  const DispatchContext = React.createContext<
    DispatchType<ActionType, DispatchReturn>
  >(null as any);

  const Provider: React.FC<ProviderProps> = ({ children, ...props }) => {
    const [state, _dispatch] = body(props as any);
    const dispatchRef = useRef(_dispatch);

    dispatchRef.current = _dispatch;

    // stable dispatch function
    const dispatch = useCallback(
      (action: ActionType) => dispatchRef.current?.(action),
      [dispatchRef]
    );

    return (
      <StateContext.Provider value={state}>
        <DispatchContext.Provider value={dispatch}>
          {children}
        </DispatchContext.Provider>
      </StateContext.Provider>
    );
  };

  const useDispatch = () => React.useContext(DispatchContext);
  const useStateContext = (selector: SelectorType<StateType>) =>
    useContextSelector(StateContext, selector);

  return [Provider, useDispatch, useStateContext] as const;
};

Usage example

type ActionType =
  | { type: 'CHANGE_STATE'; payload: ... }
  ...

export const [
  TranslationsContextProvider,
  useTranslationsDispatch,
  useTranslationsSelector,
] = createProvider(
  (props /* provider props */) => {
    const [state1, setState1] = useState(...)
    const [state2, setState2] = useState(...)
    const {data, isLoading} = useQuery(...)

    const dispatch = (action: ActionType) => {
      switch (action.type) {
        case 'CHANGE_STATE':
          setState(action.payload)
          break;
        ...
      }
    }

    const state = {
      state1,
      state2,
      data,
      isLoading
    }

    // don't forget to return state and dispatch function
    return [state, dispatch]
  })