From 69b19ced8f5da2802edf8c1dc793300af448b19e Mon Sep 17 00:00:00 2001
From: forehalo <forehalo@gmail.com>
Date: Sat, 18 Sep 2021 16:03:18 +0800
Subject: [PATCH] fix(react): state selector may use old closures

---
 packages/core/src/store.ts                  | 11 +----
 packages/react/src/__tests__/react.spec.tsx |  9 ++--
 packages/react/src/index.browser.ts         | 50 +++++++++++----------
 packages/vue/src/index.ts                   |  1 +
 4 files changed, 36 insertions(+), 35 deletions(-)

diff --git a/packages/core/src/store.ts b/packages/core/src/store.ts
index 1bd01436..f82df9db 100644
--- a/packages/core/src/store.ts
+++ b/packages/core/src/store.ts
@@ -17,7 +17,6 @@ export class Store<S> implements IStore<S> {
   private readonly reducer: Reducer<S, Action>
   private readonly epic$: BehaviorSubject<Epic>
   private actionSub = new Subscription()
-  private readonly stateSub = new Subscription()
   private readonly initAction: Action<null> = {
     type: INIT_ACTION_TYPE_SYMBOL,
     payload: null,
@@ -47,14 +46,8 @@ export class Store<S> implements IStore<S> {
    */
   setup(defaultState: S) {
     this.internalState = defaultState
-
+    this.state$.next(defaultState)
     this.subscribeAction()
-    this.stateSub.add(
-      this.state$.subscribe((state) => {
-        this.internalState = state
-      }),
-    )
-    this.state$.next(this.state)
     this.log(this.initAction)
     this.isReady = true
   }
@@ -101,6 +94,7 @@ export class Store<S> implements IStore<S> {
       if (process.env.NODE_ENV !== 'production' && newState === undefined) {
         console.warn(`${action.type} produced an undefined state, you may forget to return new State in @Reducer`)
       }
+      this.internalState = newState
       this.state$.next(newState)
     }
     this.log(action)
@@ -114,7 +108,6 @@ export class Store<S> implements IStore<S> {
   }
 
   dispose() {
-    this.stateSub.unsubscribe()
     this.actionSub.unsubscribe()
     this.action$.complete()
     this.state$.complete()
diff --git a/packages/react/src/__tests__/react.spec.tsx b/packages/react/src/__tests__/react.spec.tsx
index 054ce9a3..e0ee0ba7 100644
--- a/packages/react/src/__tests__/react.spec.tsx
+++ b/packages/react/src/__tests__/react.spec.tsx
@@ -65,10 +65,12 @@ describe('React components test', () => {
   const TestComponent = () => {
     const [localCount, setLocalCount] = useState(0)
     const state = useModuleState(CountModel, {
-      selector: (state) => state.count + localCount,
+      selector: (state) => {
+        return state.count + localCount
+      },
       dependencies: [localCount],
     })
-    spy()
+    spy(state)
     stub.callsFake(() => {
       setLocalCount(localCount + 1)
     })
@@ -87,7 +89,8 @@ describe('React components test', () => {
   it('should render three times while change local state which in dependencies list', () => {
     const reactNode = create(<TestComponent />)
     act(() => stub())
-    expect(spy.callCount).toBe(3)
+    expect(spy.callCount).toBe(2)
+    expect(spy.args).toEqual([[0], [1]])
     expect(reactNode).toMatchSnapshot()
   })
 })
diff --git a/packages/react/src/index.browser.ts b/packages/react/src/index.browser.ts
index 200798e3..57b22528 100644
--- a/packages/react/src/index.browser.ts
+++ b/packages/react/src/index.browser.ts
@@ -1,7 +1,8 @@
 import { EffectModule, ActionOfEffectModule } from '@sigi/core'
 import { ConstructorOf, IStore } from '@sigi/types'
 import { useEffect, useState, useRef, useMemo, useCallback } from 'react'
-import { distinctUntilChanged, map, skip } from 'rxjs/operators'
+import { identity } from 'rxjs'
+import { skip, tap } from 'rxjs/operators'
 
 import { useInstance } from './injectable-context'
 import { shallowEqual } from './shallow-equal'
@@ -32,38 +33,41 @@ function _useModuleState<S, U = S>(
   }
   // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
   dependencies = dependencies || []
-  const isFirstRendering = useRef(true)
-  const [appState, setState] = useState(() => {
-    const initialState = store.state
-    return selector ? selector(initialState) : initialState
-  })
-
-  const subscribe = useCallback(() => {
-    return store.state$
-      .pipe(
-        map((s) => (selector ? selector(s) : s)),
-        distinctUntilChanged((s1, s2) => equalFn(s1, s2)),
-        // skip initial state emission
-        skip(isFirstRendering.current ? 1 : 0),
-      )
-      .subscribe(setState)
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [store, ...dependencies])
+  const stateRef = useRef<S | U>()
+  const depsRef = useRef<any[]>()
+  const selectorRef = useRef<(s: S) => S | U>()
+
+  if (!equalFn(depsRef.current, dependencies)) {
+    selectorRef.current = selector ?? identity
+    stateRef.current = selectorRef.current(store.state)
+  }
+  depsRef.current = dependencies
 
-  const subscription = useMemo(() => {
-    return subscribe()
+  const [_, _flipSig] = useState(false)
+
+  const tryUpdateState = useCallback((state: S) => {
+    const newState = selectorRef.current!(state)
+    if (!equalFn(stateRef.current, newState)) {
+      stateRef.current = newState
+      _flipSig((v) => !v)
+    }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [subscribe])
+  }, [])
+
+  const subscribe = useCallback(
+    () => store.state$.pipe(skip(1), tap(tryUpdateState)).subscribe(),
+    [store, tryUpdateState],
+  )
+  const subscription = useMemo(() => subscribe(), [subscribe])
 
   useEffect(() => {
-    isFirstRendering.current = false
     const maybeReSubscribeInConcurrencyMode = subscription.closed ? subscribe() : subscription
     return () => {
       maybeReSubscribeInConcurrencyMode.unsubscribe()
     }
   }, [subscription, subscribe])
 
-  return appState
+  return stateRef.current!
 }
 
 export function useModuleState<M extends EffectModule<any>>(
diff --git a/packages/vue/src/index.ts b/packages/vue/src/index.ts
index 543bc2b5..9c189465 100644
--- a/packages/vue/src/index.ts
+++ b/packages/vue/src/index.ts
@@ -103,6 +103,7 @@ export const reactive = <
       }
     }
     if (changed) {
+      store['internalState'] = latestState
       store.state$.next(latestState)
     }
     if (typeof originalBeforeUpdate === 'function') {