Zustand Create Context

createContextStore.tsx

import {
	type PropsWithChildren,
	type ReactNode,
	createContext,
	forwardRef,
	useContext,
	useImperativeHandle,
	useRef,
} from 'react'
import { type StoreApi, createStore, useStore } from 'zustand'

type StoreDispatch<State> = StoreApi<State>['setState']

type SelectorFunc<State> = {
	(): State
	<S>(selector?: (state: State) => S): S
}

type ControllerProps<State, S> = {
	selector?: (state: State) => S
	render: (props: {
		state: S
		dispatch: StoreDispatch<State>
	}) => ReactNode
}

export function createContextStore<State>(storeCreator: () => State) {
	const Context = createContext(createStore<State>(storeCreator))

	const useSelector = ((selector: unknown) =>
		useStore(useContext(Context), selector as (state: State) => unknown)) as SelectorFunc<State>

	const Provider = forwardRef<StoreApi<State>, PropsWithChildren>(function Provider(
		{ children },
		ref,
	) {
		const storeRef = useRef(createStore<State>(storeCreator))
		useImperativeHandle(ref, () => storeRef.current)

		return <Context.Provider value={storeRef.current}>{children}</Context.Provider>
	})

	function useDispatch(): StoreDispatch<State> {
		return useContext(Context).setState
	}

	function Controller<S = State>(props: ControllerProps<State, S>) {
		const state = useSelector(props.selector)
		const dispatch = useDispatch()
		const Render = props.render

		return <Render state={state} dispatch={dispatch} />
	}

	return [Provider, useSelector, useDispatch, Controller] as const
}

const [InstanceProvider, useInstance, useInstanceDispatch, InstanceController] = createContextStore(
	() => ({
		foo: 'bar',
		baz: false,
	}),
)

function Child() {
	const baz = useInstance((state) => state.baz)
	const dispatch = useInstanceDispatch()

	if (!baz)
		return (
			<button type='button' onClick={() => dispatch((state) => ({ baz: !state.baz }))}>
				Baz
			</button>
		)

	return (
		<InstanceController
			selector={(state) => state.foo}
			render={({ state: foo }) => <div>{foo}</div>}
		/>
	)
}

function Main() {
	return (
		<InstanceProvider>
			<Child />
		</InstanceProvider>
	)
}