diff --git a/src/hooks/shallow-equal.ts b/src/hooks/shallow-equal.ts index 462712a5..3eb0ee0b 100644 --- a/src/hooks/shallow-equal.ts +++ b/src/hooks/shallow-equal.ts @@ -1,10 +1,9 @@ export const shallowEqual = (a: any, b: any): boolean => { if (a === b) return true if (typeof a === 'object' && typeof b === 'object' && a !== null && b !== null) { - return ( - Object.keys(a).length === Object.keys(b).length && - Object.keys(a).every((key) => a[key] === b[key]) - ) + const keys = Object.keys(a) + if (keys.length !== Object.keys(b).length) return false + return keys.every((key) => key in b && a[key] === b[key]) } return false } diff --git a/src/hooks/use-ayanami.ts b/src/hooks/use-ayanami.ts index c53a0a2c..f5244283 100644 --- a/src/hooks/use-ayanami.ts +++ b/src/hooks/use-ayanami.ts @@ -23,7 +23,7 @@ interface Config extends Partial { selector?: (state: S) => U } -export function useAyanami, U = M extends Ayanami ? S : never>( +export function useAyanami, S, U = M extends Ayanami ? SS : never>( A: ConstructorOf, config?: M extends Ayanami ? Config : never, ): M extends Ayanami @@ -38,7 +38,7 @@ export function useAyanami, U = M extends Ayanami getInstanceWithScope(A, reqScope), [reqScope]) ayanami.scopeName = scope || DEFAULT_SCOPE_NAME - const useAyanamiInstanceConfig = React.useMemo(() => { + const useAyanamiInstanceConfig = React.useMemo>(() => { return { destroyWhenUnmount: scope === TransientScope, selector } }, [reqScope]) diff --git a/src/hooks/use-subscribe-ayanami-state.ts b/src/hooks/use-subscribe-ayanami-state.ts index 2d978011..40e9789c 100644 --- a/src/hooks/use-subscribe-ayanami-state.ts +++ b/src/hooks/use-subscribe-ayanami-state.ts @@ -4,17 +4,16 @@ import identity from 'lodash/identity' import { shallowEqual } from './shallow-equal' import { Ayanami } from '../core' -export function useSubscribeAyanamiState, S, U>( +export function useSubscribeAyanamiState, S, U = S>( ayanami: M, selector: (state: S) => U = identity, ): unknown { - const state = ayanami.getState() + const [state, setState] = React.useState(() => selector(ayanami.getState())) const ayanamiRef = React.useRef | null>(null) const subscriptionRef = React.useRef(null) - const stateRef = React.useRef(state) - - const [, forceUpdate] = React.useState({}) + const stateRef = React.useRef(state) + const isFirstRenderRef = React.useRef(true) if (ayanamiRef.current !== ayanami) { ayanamiRef.current = ayanami @@ -25,11 +24,17 @@ export function useSubscribeAyanamiState, S, U>( } if (ayanami) { - subscriptionRef.current = ayanami.getState$().subscribe((state) => { - const before = selector(stateRef.current) - const after = selector(state) - if (!shallowEqual(before, after)) forceUpdate({}) - stateRef.current = state + subscriptionRef.current = ayanami.getState$().subscribe((moduleState) => { + if (isFirstRenderRef.current) return + if (selector === identity) { + setState(selector(moduleState)) + stateRef.current = selector(moduleState) + } else { + const before = stateRef.current + const after = selector(moduleState) + if (!shallowEqual(before, after)) setState(after) + stateRef.current = after + } }) } } @@ -43,5 +48,6 @@ export function useSubscribeAyanamiState, S, U>( [subscriptionRef], ) - return selector(state) + isFirstRenderRef.current = false + return state }