diff --git a/striker-ui/hooks/useProtectedState.ts b/striker-ui/hooks/useProtectedState.ts index 019fc9c0..f486c201 100644 --- a/striker-ui/hooks/useProtectedState.ts +++ b/striker-ui/hooks/useProtectedState.ts @@ -1,4 +1,6 @@ -import { Dispatch, SetStateAction, useState } from 'react'; +import { Dispatch, SetStateAction, useMemo, useState } from 'react'; + +import useProtect from './useProtect'; type SetStateFunction = Dispatch>; @@ -8,17 +10,24 @@ type SetStateReturnType = ReturnType>; const useProtectedState = ( initialState: S | (() => S), - protect: ( + protect?: ( fn: SetStateFunction, ...args: SetStateParameters ) => SetStateReturnType, ): [S, SetStateFunction] => { + const { protect: defaultProtect } = useProtect(); + const [state, setState] = useState(initialState); + const pfn = useMemo( + () => protect ?? defaultProtect, + [defaultProtect, protect], + ); + return [ state, (...args: SetStateParameters): SetStateReturnType => - protect(setState, ...args), + pfn(setState, ...args), ]; };