fix: should add subscribe after hot load

package.json

@@ -20,6 +20,7 @@
   ],
   "scripts": {
     "clean": "rm -rf dist",
+    "dev": "rollup -w --config rollup.config.js",
     "bundle": "rollup --config rollup.config.js",
     "build": "run-s clean bundle",
     "test": "jest",

src/hooks.spec.tsx

@@ -80,12 +80,17 @@ describe('useSelector', () => {
     const defaultSelector = selector(defaultFn);
     const strictSelector = selector(strictFn, (select, times) => [select(testBox).count, times]);
     const inlineFn = fn((select: Select) => select(testBox).count);
-    const expectCalled = (isDefault: boolean, isStrict: boolean, isInline: boolean) => {
+    const expectCalled = (
+      isDefault: boolean,
+      isStrict: boolean,
+      isInline: boolean,
+      first?: boolean,
+    ) => {
       expect(defaultFn).toBeCalledTimes(isDefault ? 1 : 0);
       defaultFn.mockClear();
-      expect(strictFn).toBeCalledTimes(isStrict ? 1 : 0);
+      expect(strictFn).toBeCalledTimes(isStrict ? (first ? 2 : 1) : 0);
       strictFn.mockClear();
-      expect(inlineFn).toBeCalledTimes(isInline ? 1 : 0);
+      expect(inlineFn).toBeCalledTimes(isInline ? (first ? 2 : 1) : 0);
       inlineFn.mockClear();
     };
     const { result, dispatch, rerender, waitForNextUpdate } = renderUseSelector(
@@ -99,7 +104,7 @@ describe('useSelector', () => {
       { count: 1, test: { count: 1, greets: ['PRELOAD'] } },
       { multiply: 1 },
     );
-    expectCalled(true, true, true);
+    expectCalled(true, true, true, true);
     expect(result.current).toEqual([1, 1, 1, 1]);
     dispatch(addGreet('HELLO'));
     await waitForNextUpdate();

src/hooks.ts

@@ -3,12 +3,19 @@
  * @author acrazing <[email protected]>
  */
 
-import { useContext, useDebugValue, useEffect, useReducer, useRef } from 'react';
+import { useContext, useDebugValue, useLayoutEffect, useEffect, useReducer, useRef } from 'react';
 import { __Context } from './context';
 import { Selector } from './selector';
 import { Dispatch, Selectable, Snapshot, Store } from './store';
 import { arrayEqual, strictEqual } from './utils';
 
+export const useIsomorphicLayoutEffect =
+  typeof window !== 'undefined' &&
+  typeof window.document !== 'undefined' &&
+  typeof window.document.createElement !== 'undefined'
+    ? useLayoutEffect
+    : useEffect;
+
 /**
  * use context's store
  *
@@ -158,9 +165,58 @@ export function useSelector<Rs extends Selectable[]>(...selectors: Rs): MapSelec
   const store = useStore();
   const lastSelector = useRef<SelectorRef>(defaultSelectorRef);
   const lastStore = useRef<StoreRef>();
+  const lastState = useRef<unknown[]>([]);
+
   if (lastStore.current?.store !== store) {
     lastSelector.current = defaultSelectorRef;
-    lastStore.current?.disposer();
+  }
+
+  if (lastStore.current?.error) {
+    const error = lastStore.current.error;
+    lastStore.current.error = void 0;
+    throw error;
+  }
+
+  const resolveState = () => {
+    if (lastStore.current?.updated) {
+      lastStore.current.updated = false;
+      return lastSelector.current.results;
+    } else {
+      if (lastSelector.current === defaultSelectorRef) {
+        lastSelector.current = { selectors: [], deps: [], snapshots: [], results: [] };
+      }
+      // updates from outside
+      const { selectors: oldSelectors, deps, snapshots, results } = lastSelector.current;
+      for (let i = 0; i < selectors.length; i++) {
+        const old = oldSelectors[i];
+        const newly = selectors[i];
+        if (typeof newly === 'object') {
+          results[i] = store.select(newly);
+          oldSelectors[i] = newly;
+        } else {
+          const newDeps = selectorChanged(old, newly, snapshots[i], store, deps[i]);
+          if (newDeps) {
+            snapshots[i] = void 0;
+            const newSnapshot: Snapshot = {};
+            results[i] = store.select(newly, newSnapshot);
+            deps[i] = newDeps === true ? void 0 : newDeps;
+            snapshots[i] = newSnapshot;
+            oldSelectors[i] = newly;
+          }
+        }
+      }
+      results.length = selectors.length;
+      return results;
+    }
+  };
+
+  let selectedState: any = resolveState();
+
+  useIsomorphicLayoutEffect(() => {
+    lastState.current = [...selectedState];
+  });
+
+  useIsomorphicLayoutEffect(() => {
     lastStore.current = {
       store,
       updated: false,
@@ -200,44 +256,15 @@ export function useSelector<Rs extends Selectable[]>(...selectors: Rs): MapSelec
         }
       }),
     };
-  }
-  useEffect(() => () => lastStore.current?.disposer(), []);
-  if (lastStore.current.error) {
-    const error = lastStore.current.error;
-    lastStore.current.error = void 0;
-    throw error;
-  }
-  let selectedState: any;
-  if (lastStore.current.updated) {
-    lastStore.current.updated = false;
-    selectedState = lastSelector.current.results;
-  } else {
-    if (lastSelector.current === defaultSelectorRef) {
-      lastSelector.current = { selectors: [], deps: [], snapshots: [], results: [] };
-    }
-    // updates from outside
-    const { selectors: oldSelectors, deps, snapshots, results } = lastSelector.current;
-    for (let i = 0; i < selectors.length; i++) {
-      const old = oldSelectors[i];
-      const newly = selectors[i];
-      if (typeof newly === 'object') {
-        results[i] = store.select(newly);
-        oldSelectors[i] = newly;
-      } else {
-        const newDeps = selectorChanged(old, newly, snapshots[i], store, deps[i]);
-        if (newDeps) {
-          snapshots[i] = void 0;
-          const newSnapshot: Snapshot = {};
-          results[i] = store.select(newly, newSnapshot);
-          deps[i] = newDeps === true ? void 0 : newDeps;
-          snapshots[i] = newSnapshot;
-          oldSelectors[i] = newly;
-        }
-      }
+
+    // if something change between render and the effect. eg. dispatch when render
+    if (!arrayEqual(lastState.current, resolveState())) {
+      update();
     }
-    results.length = selectors.length;
-    selectedState = results;
-  }
+
+    return () => lastStore.current?.disposer();
+  }, [store]);
+
   // TODO: print friendly with selector names
   useDebugValue(selectedState, (value: any[]) => {
     return value.reduce((map, value, index) => {