context state wrapper

atom.tsx

import React, { Dispatch, SetStateAction } from 'react'

type Provider = (props: React.PropsWithChildren<{}>) => JSX.Element

interface Atom<T> {
    Provider: Provider
    SetterContext: React.Context<Dispatch<SetStateAction<T>>>
    ValueContext: React.Context<T>
}

interface AtomProviderProps {
    children?: React.ReactNode
    atoms: Array<Atom<any>>
}

export function AtomProvider(props: AtomProviderProps) {
    return (
        <React.Fragment>
            {props.atoms.reduce((combined, atom, i) => {
                return <atom.Provider key={i}>{combined}</atom.Provider>
            }, props.children)}
        </React.Fragment>
    )
}

export function atom<T>(initial: T): Atom<T> {
    const ValueContext = React.createContext<T>(initial)
    const SetterContext = React.createContext<Dispatch<SetStateAction<T>>>(
        () => {}
    )

    function Provider(props: React.PropsWithChildren<{}>) {
        const [state, setState] = React.useState<T>(initial)

        return (
            <SetterContext.Provider value={setState}>
                <ValueContext.Provider value={state}>
                    {props.children}
                </ValueContext.Provider>
            </SetterContext.Provider>
        )
    }

    return {
        Provider,
        SetterContext,
        ValueContext,
    }
}

export function useAtomValue<T>(atom: Atom<T>) {
    const context = React.useContext(atom.ValueContext)

    if (context === undefined) {
        throw new Error('useAtomValue is missing its atom provider')
    }

    return context
}

export function useAtomSetter<T>(atom: Atom<T>) {
    const context = React.useContext(atom.SetterContext)

    if (context === undefined) {
        throw new Error('useAtomSetter is missing its atom provider')
    }

    return context
}

export function useAtomState<T>(atom: Atom<T>): [T, (v: T) => void] {
    const valueContext = React.useContext(atom.ValueContext)
    const setterContext = React.useContext(atom.SetterContext)

    if (valueContext === undefined) {
        throw new Error('useAtomState is missing its atom provider')
    }

    return [valueContext, setterContext]
}

usage.tsx

const MyStringState = atom<string>('')
  
const atomsToProvide = [MyStringState]

export function MyBaseComponent() {
    return (
        <AtomProvider atoms={atomsToProvide}>
            ...
        </AtomProvider>
    )
}