Minimal react store with selector

store.tsx

/*
 * Usage:
 *
 * const abContext = Store.createContext()
 *
 * function Parent(props: { a: number, b: number }) {
 *   return (
 *     <Store.Provider context={abContext} value={props}>
 *       <ChildrenA>
 *       <ChildrenB>
 *     </Store.Provider>
 *   )
 * }
 *
 * function ChildrenA() {
 *   const a = useStoreSelector(abContext, store => store.a)
 *   return (
 *     <div>{a}</div>
 *   )
 * }
 *
 * function ChildrenB() {
 *   const b = useStoreSelector(abContext, store => store.b)
 *   return (
 *     <div>{b}</div>
 *   )
 * }
 *
 */


import React from 'react';
import useLazyRef from './useLazyRef';

type Store<T> = {
  value: T,
  listeners: Set<(value: T) => void>,
}

export function createContext<T>() {
  return React.createContext<Store<T>>(undefined as any);
}

function createStore<T>(value: T): Store<T> {
  return {
    value,
    listeners: new Set(),
  }
}

export function Provider<T>({ context, value, children }: {
  context: React.Context<Store<T>>,
  value: T,
  children: React.ReactNode
}) {
  const store = useLazyRef(createStore, value);

  React.useEffect(() => {
    if (store.current.value === value) {
      return;
    }
    store.current.value = value;
    store.current.listeners.forEach(listener => {
      listener(value);
    })
  })

  return (
    <context.Provider value={store.current}>
      {children}
    </context.Provider>
  )
}

const NEVER = [] as unknown[];

export function useStoreSelector<T, U>(
  context: React.Context<Store<T>>,
  selector: (value: T) => U,
) {
  const store = React.useContext(context)
  const [state, setState] = React.useState<U>(selector(store.value))

  React.useEffect(() => {
    const handleUpdate = (value: T) => {
      setState(selector(value));
    }

    store.listeners.add(handleUpdate)
    return () => {
      store.listeners.delete(handleUpdate)
    }
  }, NEVER)

  return state;
}

This makes sense! Plus, we already use selectors for the state.

What about updating the state? Should we use another context provider in Provider component for the state update callback and another hook to get it from context?

Yes, we'd have the same thing for the API state. The nice thing is the API calls are indeed already wrapped in selectors and with useGridSelector, which means we can do the change without having to modify or break the API. This also means we can safely React.memo everything, and we'll get small granular updates for everything.

There would be one store for root props, one store for state.

In practice, we may need to use use-sync-external-store / React.useSyncExternalStore, which does the same logic as above but prevents tearing with React 18 concurrent rendering. (see details: reactwg/react-18#70)

My point was that we also need a state setter for updating the state in the context.
We can either put both state and state setter in the same context provider or use another context provider for the state setter.

It can be stored in the same context. The idea is that we only store a ref/handle to the store in the context, and that ref stays constant. What we're building is a reactive-programming system that operates outside React, and only reaches into React when it detects a change.

The store would be something like:

// React.useSyncExternalStore compatible interface
class Store {
  listeners = new Set<Function>()
  value = null

  constructor(value) {
    this.value = value
  }

  subscribe = (fn: Function) => {
    this.listeners.add(fn)
    return () => this.listeners.delete(fn)
  }

  getSnapshot = () => this.value

  setValue = (value) => {
    this.value = value
    this.listeners.forEach(l => l(value))
  }
}

We can add this store to our API ref, and our API's setState would call the store's setValue with itself.

const apiRef = useRef(...)
apiRef.current.store = new Store(apiRef.current)
apiRef.current.setState = (state) => {
  // ...
  apiRef.current.store.setValue(state)
  // ...
}

And useGridSelector would make use of it:

function useGridSelector(apiRef, selector) {
  const [state, setState] = useState(selector(apiRef.store.getSnapshot()))
  useEffect(() => {
    return apiRef.store.subscribe((value) => {
      setState(selector(value)) // setState already has the logic to avoid re-rendering if `newValue === oldValue`
    })
  }, [])
  return state
}

Alright, I missed the fact that the value in the context is stable 👍

The content makes sense to me.

The only question I would have is whether can we fall into a multi-rendering tree case? I mean, if Parent rerender, it's going to try to render ChildrenA. At the same time, when Parent's context value, he's going to try to render ChildrenA's useStoreSelector. How is React able to not render to batch these render signals to not render twice?