fix: hydrate rows subscribed in listeners

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@amos/root",
   "private": true,
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "license": "MIT",
   "workspaces": [
     "packages/*",

packages/amos-babel/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-babel",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": [],
   "peerDependencies": {

packages/amos-boxes/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-boxes",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-core/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-core",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-io/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-io",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-persist/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-persist",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-persist/src/enhancer.spec.ts

@@ -75,7 +75,7 @@ describe('withPersist', () => {
     ]);
     const r2 = store.select(userMapBox.getItem(Morty.id));
     expect(r2).toBe(Morty);
-    await sleep(1);
+    await store.dispatch(hydrate([]));
     const r3 = store.select(userMapBox.getItem(Morty.id));
     expect(r3).toEqual(Morty.set('firstName', 'F2'));
     expectCalledWith(getMulti, [[toKey(userMapBox, Morty.id)]]);
@@ -112,4 +112,36 @@ describe('withPersist', () => {
     const r5 = store.select(darkModeBox);
     expect([r4, r5]).toEqual([false, 3 * 2 + 1]);
   });
+
+  it('should not hydrate non-selected row', async () => {
+    const storage = new MemoryStorage();
+    await storage.setMulti([
+      [toKey(countBox), 1, 1],
+      [toKey(userMapBox, Rick.id), 1, Rick.set('firstName', 'F1')],
+    ]);
+    const store = createStore(void 0, withPersist({ storage, includes: () => true }));
+    store.select(countBox);
+    await store.dispatch(hydrate([]));
+    expect(store.select(userMapBox.getItem(Rick.id)).firstName).toBe(Rick.firstName);
+  });
+
+  it('should hydrate rows selected in listener', async () => {
+    const storage = new MemoryStorage();
+    await storage.setMulti([
+      [toKey(countBox), 1, 1],
+      [toKey(userMapBox, Rick.id), 1, Rick.set('firstName', 'F1')],
+    ]);
+    const store = createStore(void 0, withPersist({ storage, includes: () => true }));
+    store.subscribe(() => {
+      if (store.select(countBox) === 1) {
+        store.select(userMapBox.getItem(Rick.id));
+      }
+    });
+    store.select(countBox);
+    await store.dispatch(hydrate([]));
+    expect(store.select(countBox)).toBe(1);
+    expect(store.select(userMapBox.getIn(Rick.id, 'firstName'))).toBe(Rick.firstName);
+    await store.dispatch(hydrate([]));
+    expect(store.select(userMapBox.getIn(Rick.id, 'firstName'))).toBe('F1');
+  });
 });

packages/amos-persist/src/enhancer.ts

@@ -3,8 +3,8 @@
  * @author junbao <[email protected]>
  */
 
-import { Box, type Selectable, type Selector, StoreEnhancer } from 'amos-core';
-import { append, isAmosObject, once, override, PartialRequired } from 'amos-utils';
+import { Box, type Mutation, type Selectable, type Selector, StoreEnhancer } from 'amos-core';
+import { append, isAmosObject, once, override, PartialRequired, StackObserver } from 'amos-utils';
 import { createHydrate } from './hydrate';
 import { createPersist } from './persist';
 import { persistBox, type PersistState } from './state';
@@ -23,24 +23,19 @@ export function withPersist(options: PartialRequired<PersistOptions, 'storage'>)
     const hydrate = createHydrate(store, finalOptions);
     const persist = createPersist(store, finalOptions);
 
+    const selecting = new StackObserver();
     const initial = new Map<string, any>();
 
     const state: PersistState = {
       ...finalOptions,
-      selecting: false,
       init: once(async () => options.storage.init?.()),
       snapshot: new Map(),
+      select: selecting.observe((s: any) => store.select(s)),
       getInitial: (box) => {
         if (initial.has(box.key)) {
           return initial.get(box.key);
         }
-        const selecting = state.selecting;
-        state.selecting = true;
-        try {
-          store.select(box);
-        } finally {
-          state.selecting = selecting;
-        }
+        state.select(box);
         return state.getInitial(box);
       },
       hydrate,
@@ -52,6 +47,21 @@ export function withPersist(options: PartialRequired<PersistOptions, 'storage'>)
       initial.set(box.key, initialState);
     });
 
+    let dispatchingMutation = 0;
+    override(store, 'dispatch', (dispatch) => {
+      return (dispatchables: any) => {
+        if (!isAmosObject<Mutation>(dispatchables, 'mutation')) {
+          return dispatch(dispatchables);
+        }
+        try {
+          dispatchingMutation++;
+          return dispatch(dispatchables);
+        } finally {
+          dispatchingMutation--;
+        }
+      };
+    });
+
     let selectingRows: PersistRowKey<any> | undefined = void 0;
     override(store, 'select', (select) => {
       return (selectable: any): any => {
@@ -64,7 +74,7 @@ export function withPersist(options: PartialRequired<PersistOptions, 'storage'>)
         }
         try {
           const r = select(selectable as Selectable);
-          if (state.selecting) {
+          if (selecting.count || dispatchingMutation) {
             return r;
           }
           // hydrate rows for multi-row boxes

packages/amos-persist/src/hydrate.ts

@@ -4,7 +4,7 @@
  */
 
 import { Box, type Mutation, type Store } from 'amos-core';
-import { fromJS, type ID, isArray, must, nextSerialTicker, tryFinally } from 'amos-utils';
+import { fromJS, type ID, isArray, must, nextSerialTicker } from 'amos-utils';
 import { persistBox } from './state';
 import type { PersistEntry, PersistKey, PersistOptions } from './types';
 import { fromKey, toKey, toPersistOptions } from './utils';
@@ -70,63 +70,55 @@ export const createHydrate = (store: Store, finalOptions: PersistOptions) => {
       exactKeys.length ? state.storage.getMulti(exactKeys.map(([k, i]) => toKey(k, i))) : [],
     ]);
 
-    tryFinally(
-      () => {
-        state.selecting = true;
-        const tablePrefixMap: Record<string, PersistEntry[]> = {};
-        const exactEntries: Mutation[] = [];
-        exacts.forEach((value, index) => {
-          if (value === null) {
-            return;
-          }
-          const [key, rowId] = exactKeys[index];
-          const box = Box.get(key);
-          if (rowId !== void 0) {
-            if (!tablePrefixMap[key]) {
-              tablePrefixMap[key] = [];
-              prefixBoxes.push(key);
-              prefixes.push(tablePrefixMap[key]);
+    const tablePrefixMap: Record<string, PersistEntry[]> = {};
+    const exactEntries: Mutation[] = [];
+    exacts.forEach((value, index) => {
+      if (value === null) {
+        return;
+      }
+      const [key, rowId] = exactKeys[index];
+      const box = Box.get(key);
+      if (rowId !== void 0) {
+        if (!tablePrefixMap[key]) {
+          tablePrefixMap[key] = [];
+          prefixBoxes.push(key);
+          prefixes.push(tablePrefixMap[key]);
+        }
+        tablePrefixMap[key].push([toKey(key, rowId), ...value]);
+      } else if (state.select(box) === state.getInitial(box)) {
+        const js = migrate(box, value[0], '', value[1]);
+        if (js !== void 0) {
+          exactEntries.push(box.setState(fromJS(state.getInitial(box), js)));
+        }
+      }
+    });
+
+    store.dispatch(
+      prefixes
+        .map((p, i): Mutation[] => {
+          const box = Box.get(prefixBoxes[i]);
+          const curr = state.select(box);
+          const data: Record<string, any> = {};
+          for (let [k, v, d] of p) {
+            const id = fromKey(k);
+            if (id === void 0) {
+              continue;
+            }
+            if (box.table!.getRow(curr, id) !== box.table!.getRow(state.getInitial(box), id)) {
+              continue;
             }
-            tablePrefixMap[key].push([toKey(key, rowId), ...value]);
-          } else if (store.select(box) === state.getInitial(box)) {
-            const js = migrate(box, value[0], '', value[1]);
-            if (js !== void 0) {
-              exactEntries.push(box.setState(fromJS(box.getInitialState(), js)));
+            d = migrate(box, v, id, d);
+            if (d !== void 0) {
+              data[id] = d;
             }
           }
-        });
-
-        store.dispatch(
-          prefixes
-            .map((p, i): Mutation[] => {
-              const box = Box.get(prefixBoxes[i]);
-              const curr = store.select(box);
-              const data: Record<string, any> = {};
-              for (let [k, v, d] of p) {
-                const id = fromKey(k);
-                if (id === void 0) {
-                  continue;
-                }
-                if (box.table!.getRow(curr, id) !== box.table!.getRow(state.getInitial(box), id)) {
-                  continue;
-                }
-                d = migrate(box, v, id, d);
-                if (d !== void 0) {
-                  data[id] = d;
-                }
-              }
-              if (Object.keys(data).length === 0) {
-                return [];
-              }
-              return [box.setState(box.table!.hydrate(curr, data))];
-            })
-            .flat()
-            .concat(exactEntries),
-        );
-      },
-      () => {
-        state.selecting = false;
-      },
+          if (Object.keys(data).length === 0) {
+            return [];
+          }
+          return [box.setState(box.table!.hydrate(curr, data))];
+        })
+        .flat()
+        .concat(exactEntries),
     );
   }, finalOptions.onError);
 };

packages/amos-persist/src/state.ts

@@ -10,7 +10,7 @@ import { toKey } from './utils';
 
 export interface PersistState extends PersistOptions {
   init: () => Promise<void>;
-  selecting: boolean;
+  select: Select;
   getInitial: (box: Box) => any;
   snapshot: Map<string, any>;
   hydrate: NextTicker<PersistKey<any>, void>;

packages/amos-react/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-react",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": [],
   "peerDependencies": {

packages/amos-shapes/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-shapes",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-testing/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-testing",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-typescript/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-typescript",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": [],
   "peerDependencies": {

packages/amos-utils/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos-utils",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "private": true,
   "files": []
 }

packages/amos-utils/src/misc.ts

@@ -4,7 +4,7 @@
  */
 
 import { isIterable } from './equals';
-import { ValueOrConstructor, ValueOrFunc } from './types';
+import { type AnyFunc, ValueOrConstructor, ValueOrFunc } from './types';
 
 export function must<T>(value: T, message: string): asserts value {
   if (!value) {
@@ -51,6 +51,21 @@ export function toArray<T>(items: T[] | Iterable<T> | T): T[] {
   return [items as T];
 }
 
+export class StackObserver {
+  count = 0;
+
+  observe<T extends AnyFunc>(fn: T) {
+    return ((...args: any[]) => {
+      this.count++;
+      try {
+        return fn(...args);
+      } finally {
+        this.count--;
+      }
+    }) as T;
+  }
+}
+
 export interface NextTicker<T, R> {
   (...items: T[]): void;
   wait(...items: T[]): Promise<R>;

packages/amos/package.json

@@ -1,6 +1,6 @@
 {
   "name": "amos",
-  "version": "0.3.0-beta.32",
+  "version": "0.3.0-beta.33",
   "description": "An out-of-the-box state management library designed for your large-scale projects.",
   "keywords": [
     "amos",