context_with_selector.tsx

import React, { createContext, useRef, useEffect, useContext, useState } from 'react'
type Selector<T, U = any> = (value: T) => U
function shallowEqual(a: unknown, b: unknown): boolean {
  if (a === b) return true
  if (typeof a === 'object' && typeof b === 'object') {
    if (a == null || b == null) return a === b
    if (Array.isArray(a) || Array.isArray(b)) {
      if (!Array.isArray(a) || !Array.isArray(b) || a.length !== b.length) return false
      return a.every((v, i) => v == b[i])
    }
    return Object.keys(a).every(k => a[k] === b[k]) && Object.keys(b).every(k => a[k] === b[k])
  } else {
    return a == b
  }
}
class SelectableContextState<T> {
  subscribers = new Map<Selector<T>, { value: any; callbacks: Set<() => void> }>()
  constructor(public value: T) {}
  update(value: T) {
    this.value = value
    this.subscribers.forEach((subscriber, selector) => {
      const newValue = selector(value)
      if (!shallowEqual(subscriber.value, value)) {
        subscriber.value = newValue
        subscriber.callbacks.forEach(cb => cb())
      }
    })
  }
  subscribe(selector: Selector<T>, callback: () => void) {
    let s = this.subscribers.get(selector)
    if (!s) {
      const value = selector(this.value)
      s = { value, callbacks: new Set() }
      this.subscribers.set(selector, s)
    }
    const subscriber = s
    subscriber.callbacks.add(callback)
    return () => {
      subscriber.callbacks.delete(callback)
      if (subscriber.callbacks.size === 0) this.subscribers.delete(selector)
    }
  }
  select<U>(selector: Selector<T, U>): U {
    return this.subscribers.get(selector)!.value
  }
}
type SelectableContext<T> = {
  Context: React.Context<SelectableContextState<T> | null>
  Provider: React.FC<{ value: T}>
}
function createSelectableContext<T>(): SelectableContext<T> {
  const Context = createContext<SelectableContextState<T> | null>(null)
  const Provider: React.FC<{ value: T }> = ({ value, children }) => {
    const stateRef = useRef<SelectableContextState<T> | null>(null)
    if (!stateRef.current) {
      stateRef.current = new SelectableContextState<T>(value)
    } else {
      stateRef.current.update(value)
    }
    return <Context.Provider value={stateRef.current} children={children} />
  }
  return {
    Context,
    Provider
  }
}

function useSelectableConext<T, U>(context: SelectableContext<T>, selector: Selector<T, U>): U {
  const [, forceUpdate] = useState([])
  const state = useContext(context.Context)
  if (!state) throw 'No Provider Given'
  const unsubscribeRef = useRef<(() => void) | null>(null)
  if (!unsubscribeRef.current) unsubscribeRef.current = state.subscribe(selector, () => forceUpdate([]))
  useEffect(() => () => unsubscribeRef.current!(), [])
  return state.select(selector)
}

type SampleType = { id: number; name: string; items: number[] }
const SampleContext = createSelectableContext<SampleType>()

const nameSelector = (v: SampleType) => v.name
const Foo: React.FC = () => {
  const name = useSelectableConext(SampleContext, v => v.name)
  return <div>name: {name}; {Math.random()}</div>
}
const Bar: React.FC = () => {
  const { id, total } = useSelectableConext(SampleContext, v => ({ id: v.id, total: v.items.reduce((a, b) => a + b) }))
  return <div>id: {id}; total: {total}; {Math.random()}</div>
}

const Sample: React.FC = () => {
  const [value, setValue] = useState({ id: 1, name: 'a', items: [1, 2, 3] })
  return <SampleContext.Provider value={value}>
    <input value={JSON.stringify(value)} onChange={e => { setValue(JSON.parse(e.target.value)) }} />
    <Foo />
    <Bar />
  </SampleContext.Provider>
}