Simpler Control Props Implementation?

current-implementation.tsx

import { useReducer, useRef } from 'react'
import { Switch } from '#shared/switch.tsx'

function callAll<Args extends Array<unknown>>(
	...fns: Array<((...args: Args) => unknown) | undefined>
) {
	return (...args: Args) => fns.forEach(fn => fn?.(...args))
}

export type ToggleState = { on: boolean }
export type ToggleAction =
	| { type: 'toggle' }
	| { type: 'reset'; initialState: ToggleState }

export function toggleReducer(state: ToggleState, action: ToggleAction) {
	switch (action.type) {
		case 'toggle': {
			return { on: !state.on }
		}
		case 'reset': {
			return action.initialState
		}
	}
}

export function useToggle({
	initialOn = false,
	reducer = toggleReducer,
	onChange,
	on: controlledOn,
}: {
	initialOn?: boolean
	reducer?: typeof toggleReducer
	onChange?: (state: ToggleState, action: ToggleAction) => void
	on?: boolean
} = {}) {
	const { current: initialState } = useRef<ToggleState>({ on: initialOn })
	const [state, dispatch] = useReducer(reducer, initialState)
	const onIsControlled = controlledOn != null
	const on = onIsControlled ? controlledOn : state.on

	function dispatchWithOnChange(action: ToggleAction) {
		if (!onIsControlled) {
			dispatch(action)
		}
		onChange?.(reducer({ ...state, on }, action), action)
	}

	const toggle = () => dispatchWithOnChange({ type: 'toggle' })
	const reset = () => dispatchWithOnChange({ type: 'reset', initialState })

	function getTogglerProps<Props>({
		onClick,
		...props
	}: { onClick?: React.ComponentProps<'button'>['onClick'] } & Props) {
		return {
			'aria-checked': on,
			onClick: callAll(onClick, toggle),
			...props,
		}
	}

	function getResetterProps<Props>({
		onClick,
		...props
	}: { onClick?: React.ComponentProps<'button'>['onClick'] } & Props) {
		return {
			onClick: callAll(onClick, reset),
			...props,
		}
	}

	return {
		on,
		reset,
		toggle,
		getTogglerProps,
		getResetterProps,
	}
}

export function Toggle({
	on: controlledOn,
	onChange,
}: {
	on?: boolean
	onChange?: (state: ToggleState, action: ToggleAction) => void
}) {
	const { on, getTogglerProps } = useToggle({ on: controlledOn, onChange })
	const props = getTogglerProps({ on })
	return <Switch {...props} />
}

simpler-implementation.tsx

import { useReducer, useRef } from 'react'
import { Switch } from '#shared/switch.tsx'

function callAll<Args extends Array<unknown>>(
	...fns: Array<((...args: Args) => unknown) | undefined>
) {
	return (...args: Args) => fns.forEach(fn => fn?.(...args))
}

export type ToggleState = { on: boolean }
export type ToggleAction =
	| { type: 'toggle' }
	| { type: 'reset'; initialState: ToggleState }

export function toggleReducer(state: ToggleState, action: ToggleAction) {
	switch (action.type) {
		case 'toggle': {
			return { on: !state.on }
		}
		case 'reset': {
			return action.initialState
		}
	}
}

function useUncontrolledToggle({
	initialOn = false,
	reducer = toggleReducer,
	onChange,
}: {
	initialOn?: boolean
	reducer?: typeof toggleReducer
	onChange?: (state: ToggleState, action: ToggleAction) => void
} = {}) {
	const { current: initialState } = useRef<ToggleState>({ on: initialOn })
	const [state, dispatch] = useReducer(reducer, initialState)
	return useToggle({
		initialState,
		state,
		dispatch(action) {
			dispatch(action)
			onChange?.(state, action)
		},
	})
}

export function useToggle({
	initialState,
	state,
	dispatch,
}: {
	initialState: ToggleState
	state: ToggleState
	dispatch: (action: ToggleAction) => void
}) {
	const { on } = state
	const toggle = () => dispatch({ type: 'toggle' })
	const reset = () => dispatch({ type: 'reset', initialState })

	function getTogglerProps<Props>({
		onClick,
		...props
	}: { onClick?: React.ComponentProps<'button'>['onClick'] } & Props) {
		return {
			'aria-checked': on,
			onClick: callAll(onClick, toggle),
			...props,
		}
	}

	function getResetterProps<Props>({
		onClick,
		...props
	}: { onClick?: React.ComponentProps<'button'>['onClick'] } & Props) {
		return {
			onClick: callAll(onClick, reset),
			...props,
		}
	}

	return {
		on,
		reset,
		toggle,
		getTogglerProps,
		getResetterProps,
	}
}

export function Toggle({
	on: controlledOn,
	initialOn = controlledOn ?? false,
	onChange,
	reducer = toggleReducer,
}: {
	initialOn?: boolean
	on?: boolean
	onChange?: (state: ToggleState, action: ToggleAction) => void
	reducer?: typeof toggleReducer
}) {
	const onIsControlled = controlledOn != null
	if (onIsControlled) {
		return (
			<ControlledToggle
				on={controlledOn}
				onChange={onChange}
				reducer={reducer}
			/>
		)
	} else {
		return (
			<UncontrolledToggle
				initialOn={initialOn}
				onChange={onChange}
				reducer={reducer}
			/>
		)
	}
}

export function ControlledToggle({
	on: controlledOn,
	onChange,
	reducer = toggleReducer,
}: {
	on: boolean
	onChange?: (state: ToggleState, action: ToggleAction) => void
	reducer?: typeof toggleReducer
}) {
	const { current: initialState } = useRef<ToggleState>({ on: controlledOn })
	const state = { on: controlledOn }
	const { on, getTogglerProps } = useToggle({
		initialState,
		state,
		dispatch(action) {
			onChange?.(reducer({ ...state, on }, action), action)
		},
	})
	const props = getTogglerProps({ on })
	return <Switch {...props} />
}

export function UncontrolledToggle({
	initialOn = false,
	onChange,
	reducer = toggleReducer,
}: {
	initialOn?: boolean
	onChange?: (state: ToggleState, action: ToggleAction) => void
	reducer?: typeof toggleReducer
}) {
	const { on, getTogglerProps } = useUncontrolledToggle({
		initialOn,
		onChange,
		reducer,
	})
	const props = getTogglerProps({ on })
	return <Switch {...props} />
}