feat: support nested objects and slice pattern (#7)

* refactor * applyChanges from jotai-valtio * patch function recursively * examples/04 * update CHANGELOG & README * refactor
zustandjs · Jul 14, 2024 · edef291 · edef291
1 parent 653d52d
commit edef291
Show file tree

Hide file tree

Showing 9 changed files with 191 additions and 26 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+## Added
+
+- feat: support nested objects and slice pattern #7
+
 ## [0.3.0] - 2024-07-13
 
 ## Added

diff --git a/README.md b/README.md
@@ -64,6 +64,7 @@ You can also try them directly:
 [01](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/01_counter)
 [02](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/02_methods)
 [03](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/03_middleware)
+[04](https://stackblitz.com/github/zustandjs/zustand-valtio/tree/main/examples/04_slices)
 
 ## Tweets
 

diff --git a/examples/04_slices/index.html b/examples/04_slices/index.html
@@ -0,0 +1,9 @@
+<html>
+  <head>
+    <title>example</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.tsx"></script>
+  </body>
+</html>
diff --git a/examples/04_slices/package.json b/examples/04_slices/package.json
@@ -0,0 +1,22 @@
+{
+  "name": "example",
+  "version": "0.0.0",
+  "private": true,
+  "type": "module",
+  "dependencies": {
+    "react": "latest",
+    "react-dom": "latest",
+    "valtio": "latest",
+    "zustand": "latest",
+    "zustand-valtio": "latest"
+  },
+  "devDependencies": {
+    "@types/react": "latest",
+    "@types/react-dom": "latest",
+    "typescript": "latest",
+    "vite": "latest"
+  },
+  "scripts": {
+    "dev": "vite"
+  }
+}
diff --git a/examples/04_slices/src/app.tsx b/examples/04_slices/src/app.tsx
@@ -0,0 +1,67 @@
+import { create } from 'zustand';
+import { withProxy } from 'zustand-valtio';
+
+const countSlice = {
+  value: 0,
+  inc() {
+    this.value++;
+  },
+  reset() {
+    this.value = 0;
+  },
+};
+
+const textSlice = {
+  value: 'Hello',
+  updateText(newText: string) {
+    this.value = newText;
+  },
+  reset() {
+    this.value = 'Hello';
+  },
+};
+
+const useCounterState = create(
+  withProxy({
+    count: countSlice,
+    text: textSlice,
+    reset() {
+      this.count.reset();
+      this.text.reset();
+    },
+  }),
+);
+
+const Counter = () => {
+  const count = useCounterState((state) => state.count.value);
+  const inc = useCounterState((state) => state.count.inc);
+  const text = useCounterState((state) => state.text.value);
+  const updateText = useCounterState((state) => state.text.updateText);
+  const reset = useCounterState((state) => state.reset);
+  return (
+    <>
+      <p>
+        Count: {count}
+        <button type="button" onClick={inc}>
+          +1
+        </button>
+      </p>
+      <p>
+        <input value={text} onChange={(e) => updateText(e.target.value)} />
+      </p>
+      <p>
+        <button type="button" onClick={reset}>
+          Reset
+        </button>
+      </p>
+    </>
+  );
+};
+
+const App = () => (
+  <div>
+    <Counter />
+  </div>
+);
+
+export default App;
diff --git a/examples/04_slices/src/main.tsx b/examples/04_slices/src/main.tsx
@@ -0,0 +1,10 @@
+import { StrictMode } from 'react';
+import { createRoot } from 'react-dom/client';
+
+import App from './app';
+
+createRoot(document.getElementById('root')!).render(
+  <StrictMode>
+    <App />
+  </StrictMode>,
+);
diff --git a/examples/04_slices/tsconfig.json b/examples/04_slices/tsconfig.json
@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "strict": true,
+    "target": "es2018",
+    "esModuleInterop": true,
+    "module": "esnext",
+    "moduleResolution": "bundler",
+    "skipLibCheck": true,
+    "allowJs": true,
+    "noUncheckedIndexedAccess": true,
+    "exactOptionalPropertyTypes": true,
+    "jsx": "react-jsx"
+  }
+}
diff --git a/package.json b/package.json
@@ -42,7 +42,8 @@
     "test:spec": "vitest run",
     "examples:01_counter": "DIR=01_counter vite",
     "examples:02_methods": "DIR=02_methods vite",
-    "examples:03_middleware": "DIR=03_middleware vite"
+    "examples:03_middleware": "DIR=03_middleware vite",
+    "examples:04_slices": "DIR=04_slices vite"
   },
   "keywords": [
     "react",

diff --git a/src/index.ts b/src/index.ts
@@ -1,15 +1,20 @@
 import type { StateCreator, StoreApi, StoreMutatorIdentifier } from 'zustand';
 import { proxy, snapshot, subscribe } from 'valtio/vanilla';
 
+type AnyFunction = (...args: never[]) => unknown;
+
+type DeepWritable<T> = T extends object
+  ? { -readonly [K in keyof T]: DeepWritable<T[K]> }
+  : T;
+
 // TODO replace with Valtio's Snapshot type in v2
-type Snapshot<T> = T extends (...args: never[]) => unknown
+type Snapshot<T> = T extends AnyFunction
   ? T
   : T extends object
     ? { readonly [K in keyof T]: Snapshot<T[K]> }
     : T;
 
 type StoreWithProxy<T> = {
-  setState: never;
   getProxyState: () => T;
 };
 
@@ -19,10 +24,6 @@ type WithWithProxy<S, _A> = S extends { getState: () => Snapshot<infer T> }
   ? Write<S, StoreWithProxy<DeepWritable<T>>>
   : never;
 
-type DeepWritable<T> = T extends object
-  ? { -readonly [K in keyof T]: DeepWritable<T[K]> }
-  : T;
-
 declare module 'zustand/vanilla' {
   interface StoreMutators<S, A> {
     'zustand-valtio': WithWithProxy<S, A>;
@@ -41,8 +42,51 @@ type WithProxyImpl = <T extends object>(
   initialObject: T,
 ) => StateCreator<Snapshot<T>, [], []>;
 
+const isObject = (x: unknown): x is object =>
+  typeof x === 'object' && x !== null && !(x instanceof Promise);
+
+const patchObjectMethods = <T extends object>(
+  proxyObject: T,
+  fn: (obj: object, key: string, fn: AnyFunction) => void,
+) => {
+  (Object.getOwnPropertyNames(proxyObject) as (keyof T)[]).forEach((key) => {
+    const value = proxyObject[key];
+    if (typeof value === 'function') {
+      fn(proxyObject, key as string, value as AnyFunction);
+    } else if (isObject(value)) {
+      patchObjectMethods(value, fn);
+    }
+  });
+};
+
+const applyChanges = <T extends object>(proxyObject: T, prev: T, next: T) => {
+  (Object.getOwnPropertyNames(prev) as (keyof T)[]).forEach((key) => {
+    if (!(key in next)) {
+      delete proxyObject[key];
+    } else if (Object.is(prev[key], next[key])) {
+      // unchanged
+    } else if (
+      isObject(proxyObject[key]) &&
+      isObject(prev[key]) &&
+      isObject(next[key])
+    ) {
+      applyChanges(
+        proxyObject[key] as unknown as object,
+        prev[key] as unknown as object,
+        next[key] as unknown as object,
+      );
+    } else {
+      proxyObject[key] = next[key];
+    }
+  });
+  (Object.keys(next) as (keyof T)[]).forEach((key) => {
+    if (!(key in prev)) {
+      proxyObject[key] = next[key];
+    }
+  });
+};
+
 const withProxyImpl: WithProxyImpl = (initialObject) => (set, get, api) => {
-  type AnyObject = Record<string, unknown>;
   const proxyState = proxy(initialObject);
   let mutating = 0;
   let updating = 0;
@@ -53,37 +97,30 @@ const withProxyImpl: WithProxyImpl = (initialObject) => (set, get, api) => {
       --updating;
     }
   };
-  Object.keys(proxyState).forEach((key) => {
-    const fn = (proxyState as AnyObject)[key];
-    // TODO this doesn't handle nested objects
-    if (typeof fn === 'function') {
-      (proxyState as AnyObject)[key] = (...args: never[]) => {
+  patchObjectMethods(proxyState, (obj, key, fn) => {
+    Object.defineProperty(obj, key, {
+      value: (...args: never[]) => {
         try {
           ++mutating;
-          return fn.apply(proxyState, args);
+          return fn.apply(obj, args);
         } finally {
           --mutating;
           updateState();
         }
-      };
-    }
+      },
+    });
   });
   type Api = StoreApi<unknown> & StoreWithProxy<typeof initialObject>;
-  delete (api as Api).setState;
   (api as Api).getProxyState = () => proxyState;
   subscribe(proxyState, updateState, true);
   api.subscribe(() => {
     if (!updating) {
       // HACK for persist hydration
-      const state = get() as AnyObject;
-      Object.keys(state).forEach((key) => {
-        const val = state[key];
-        // TODO this doesn't handle nested objects
-        if (typeof val !== 'function') {
-          // XXX this will throw if val is a snapshot
-          (proxyState as AnyObject)[key] = val;
-        }
-      });
+      applyChanges(
+        proxyState,
+        snapshot(proxyState) as typeof proxyState,
+        get() as typeof proxyState,
+      );
     }
   });
   return snapshot(proxyState) as Snapshot<typeof proxyState>;