feat(zustand): support slice zustand store

unadlib · Jan 1, 2025 · ab23bb5 · ab23bb5
1 parent 69ffbb1
commit ab23bb5
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 8 deletions.
diff --git a/packages/coaction-zustand/src/index.ts b/packages/coaction-zustand/src/index.ts
@@ -5,18 +5,55 @@ type BindZustand = <T>(
   initializer: StateCreator<T, [], []>
 ) => StateCreator<T, [], []>;
 
+const storeMap = new Map<string, StoreApi<any>>();
+
 /**
  * Bind a store to Zustand
  */
 export const bindZustand = ((initializer: StateCreator<any, [], []>) =>
   (set, get, zustandStore) => {
     let coactionStore: Store<object>;
     const internalBindZustand = createBinder<BindZustand>({
-      handleStore: (store, rawState, state, internal) => {
+      handleStore: (store, rawState, state, internal, key) => {
+        coactionStore = store;
+        if (key) {
+          if (zustandStore.getState() === (internal.rootState as any)[key])
+            return;
+          (internal.rootState as any)[key] = zustandStore.getState();
+          storeMap.set(key, zustandStore);
+          let isCoactionUpdated = false;
+          zustandStore.subscribe(() => {
+            if (!isCoactionUpdated) {
+              (internal.rootState as any)[key] =
+                zustandStore.getState() as object;
+              if (coactionStore.share === 'client') {
+                throw new Error('client zustand store cannot be updated');
+              } else if (coactionStore.share === 'main') {
+                // emit to all clients
+                coactionStore.setState({
+                  [key]: zustandStore.getState()
+                });
+              }
+            }
+            internal.listeners.forEach((listener) => listener());
+          });
+          if (internal.updateImmutable) return;
+          internal.updateImmutable = (state: any) => {
+            isCoactionUpdated = true;
+            try {
+              for (const _key of Object.keys(state)) {
+                const zustandStore = storeMap.get(_key)!;
+                zustandStore.setState(state[_key], true);
+              }
+            } finally {
+              isCoactionUpdated = false;
+            }
+          };
+          return;
+        }
         if (zustandStore.getState() === internal.rootState) return;
         let isCoactionUpdated = false;
         internal.rootState = zustandStore.getState() as object;
-        coactionStore = store;
         // TODO: check Slice Zustand store
         zustandStore.subscribe(() => {
           if (!isCoactionUpdated) {

diff --git a/packages/core/src/binder.ts b/packages/core/src/binder.ts
@@ -32,10 +32,26 @@ export function createBinder<F = (...args: any[]) => any>({
    * handleStore is a function to handle the store object.
    */
   handleStore: (
+    /**
+     * Coaction store
+     */
     store: Store<object>,
+    /**
+     * The raw state object from 3rd party library.
+     */
     rawState: object,
+    /**
+     * 3rd party library state object to Coaction state object.
+     */
     state: object,
-    internal: Internal<object>
+    /**
+     * internal Coaction API.
+     */
+    internal: Internal<object>,
+    /**
+     * the key of the slice state object.
+     */
+    key?: string
   ) => void;
 }) {
   return (<S extends object>(state: S): S => {

diff --git a/packages/core/src/getInitialState.ts b/packages/core/src/getInitialState.ts
@@ -7,12 +7,28 @@ export const getInitialState = <T extends CreateState>(
   createState: any,
   internal: Internal<T>
 ) => {
-  const makeState = (fn: (...args: any[]) => any) => {
+  const makeState = (
+    /**
+     * createState is a function to create the state object.
+     */
+    createState: (
+      setState: (state: any) => void,
+      getState: () => any,
+      store: Store<T>
+    ) => any,
+    /**
+     * the key of the slice state object.
+     */
+    key?: string
+  ) => {
     // make sure createState is a function
-    if (process.env.NODE_ENV !== 'production' && typeof fn !== 'function') {
+    if (
+      process.env.NODE_ENV !== 'production' &&
+      typeof createState !== 'function'
+    ) {
       throw new Error('createState should be a function');
     }
-    let state = fn(store.setState, store.getState, store);
+    let state = createState(store.setState, store.getState, store);
     // support 3rd party library store like zustand, redux
     if (state.getState) {
       state = state.getState();
@@ -23,7 +39,7 @@ export const getInitialState = <T extends CreateState>(
     // support bind store like mobx
     if (state[bindSymbol]) {
       const rawState = state[bindSymbol].bind(state);
-      state[bindSymbol].handleStore(store, rawState, state, internal);
+      state[bindSymbol].handleStore(store, rawState, state, internal, key);
       delete state[bindSymbol];
       return rawState;
     }
@@ -32,7 +48,9 @@ export const getInitialState = <T extends CreateState>(
   return store.isSliceStore
     ? Object.entries(createState).reduce(
         (stateTree, [key, value]) =>
-          Object.assign(stateTree, { [key]: makeState(value as Slice<any>) }),
+          Object.assign(stateTree, {
+            [key]: makeState(value as Slice<any>, key)
+          }),
         {} as ISlices<Slice<any>>
       )
     : makeState(createState as Slice<any>);