useSelectiveState / localizing state access to child components to avoid unnecessary re-renders

childState.jsx

import React, { useState } from 'react';

/*
It's common to be in a situation where your react component renders a bunch of children,
and only one of those children rely on some state your component keeps. You don't want
to re-render ALL of your children just because you need to keep some state relevant to
only one of them, so the natural refactoring is to localize that state to the one child.
But... what if your parent (or one of the other children) is actually responsible for
*setting* that state, or maybe you just have a couple of disparate components that use
that state, while most of them don't touch it?

I ran into this situation with some "saving status" state. I have a progress indicator
which spins when I perform operations in other components, and then it stops spinning.
The spinner is not a child component of those other components; it's in a fixed location
on the page.

e.g.,

function Parent() {
  const [saving, setSaving] = useState("inactive");

  async function onTitleChange(newTitle) {
    setSaving("saving");
    await mutation.commit(newTitle);
    setSaving("inactive");
  }

  return <div>
    <Title onTitleChange={onTitleChange} />
    <Spinner saving={saving} />
    <Many /> <Other /> <Expensive /> <Components />
   </div>;
}

function Spinner({saving}) {
  if (saving === "saving") return <div>SPIN SPIN SPIN</div>;
  else return null;
}

I don't want to re-render "Many other expensive components", because there are many of
them and they take some time to render. The only idiomatic way I've heard to deal with
this issue is to wrap them in React.memo wrappers so that they don't get re-rendered
when their props don't change. I didn't like this for the following reasons:
- it required me to update many components to add those React.memo wrappers
- memoizing many components can lead to other performance issues (it takes resources to
  manage those memos)
- those components aren't even related to the problem I'm trying to solve! The "state
  problem" is between Parent and Spinner; I just want a way to be explicit about which
  bits of state apply to which child components, and only re-render those children when
  the state important to them changes.

What I wanted was a way to localize my state to `Child`, but still have a way to set it
from the parent. So I implemented a hacky "Provide a callback for registering a callback
that allows the parent to invoke a function in the child" pattern, then I thought about
how I could make it a bit more clean & clear for this particular situation.

So what I came up with is this:

function Parent() {
  const [SavingConsumer, setSaving] = useSelectiveState("inactive");

  async function onTitleChange(newTitle) {
    setSaving("saving");
    await mutation.commit(newTitle);
    setSaving("inactive");
  }

  return <div>
    <Title onTitleChange={onTitleChange} />
    <SavingConsumer>{saving => <Spinner saving={saving} />}</SavingConsumer>
    <Many /> <Other /> <Expensive /> <Components />
   </div>;
}

function Spinner({saving}) {
  if (saving === "saving") return <div>SPIN SPIN SPIN</div>;
  else return null;
}

So, now:
- The entirety of parent does not get rendered when the "saving" state is changed
- only the Spinner gets re-rendered, since we wrapped it in <SavingConsumer>
- the implementation of Spinner didn't even change, only Parent. The refactoring was
  non-invasive.

The implementation follows.
*/


class ChildState {
  constructor(initial) {
    this.childStateCallback = undefined;
    this.initial = initial;
  }

  registerCallback(cb) {
    this.childStateCallback = cb;
  }

  set(value) {
    if (!this.childStateCallback) {
      this.initial = value;
    } else {
      this.initial = undefined;
      this.childStateCallback(value);
    }
  }

  Consumer(props) {
    return <SelectiveStateConsumer state={this} {...props} />;
  }
}

/** Declare some state that will only be used explicitly in certain subtrees of your
 * component.
 *
 * You generally use this like `useState`, but this is *not* a hook.  It will not cause
 * the calling component to re-render. Instead of returning [currentState, setter] like
 * useState does, it returns [StateConsumer, setter]. You use it like this:
 *
 * const [StateConsumer, setState] = useSelectiveState(null);
 *
 * return <div>
 *   <div onClick={() => setState(Math.random())} />
 *   <StateConsumer>{state => <div>{state}</div>}</state.Consumer>
 * </div>;
 */
export function useSelectiveState(initial) {
  const cs = new ChildState(initial);
  const consumer = cs.Consumer.bind(cs);
  const setter = cs.set.bind(cs);
  return [consumer, setter];
}

export function SelectiveStateConsumer({state, ...props}) {
  const [s, setS] = useState(state.initial);
  state.registerCallback(setS);
  return props.children(s);
}