Merge pull request #465 from InfiniteXyy/master

perf: avoid using selector each render

Author: Miloas

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
 }

src/hooks/use-ayanami.ts

@@ -23,7 +23,7 @@ interface Config<S, U> extends Partial<ScopeConfig> {
   selector?: (state: S) => U
 }
 
-export function useAyanami<M extends Ayanami<any>, U = M extends Ayanami<infer S> ? S : never>(
+export function useAyanami<M extends Ayanami<S>, S, U = M extends Ayanami<infer SS> ? SS : never>(
   A: ConstructorOf<M>,
   config?: M extends Ayanami<infer S> ? Config<S, U> : never,
 ): M extends Ayanami<infer S>
@@ -38,7 +38,7 @@ export function useAyanami<M extends Ayanami<any>, U = M extends Ayanami<infer S
   const ayanami = React.useMemo(() => getInstanceWithScope(A, reqScope), [reqScope])
   ayanami.scopeName = scope || DEFAULT_SCOPE_NAME
 
-  const useAyanamiInstanceConfig = React.useMemo<UseAyanamiInstanceConfig>(() => {
+  const useAyanamiInstanceConfig = React.useMemo<UseAyanamiInstanceConfig<S>>(() => {
     return { destroyWhenUnmount: scope === TransientScope, selector }
   }, [reqScope])
 

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<M extends Ayanami<S>, S, U>(
+export function useSubscribeAyanamiState<M extends Ayanami<S>, S, U = S>(
   ayanami: M,
   selector: (state: S) => U = identity,
 ): unknown {
-  const state = ayanami.getState()
+  const [state, setState] = React.useState<U>(() => selector(ayanami.getState()))
 
   const ayanamiRef = React.useRef<Ayanami<S> | null>(null)
   const subscriptionRef = React.useRef<Subscription | null>(null)
-  const stateRef = React.useRef<S>(state)
-
-  const [, forceUpdate] = React.useState({})
+  const stateRef = React.useRef<U>(state)
+  const isFirstRenderRef = React.useRef(true)
 
   if (ayanamiRef.current !== ayanami) {
     ayanamiRef.current = ayanami
@@ -25,11 +24,17 @@ export function useSubscribeAyanamiState<M extends Ayanami<S>, 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<M extends Ayanami<S>, S, U>(
     [subscriptionRef],
   )
 
-  return selector(state)
+  isFirstRenderRef.current = false
+  return state
 }