feat: support slices pattern type (#16)

src/computed.test.ts

@@ -1,4 +1,4 @@
-import { create } from "zustand"
+import { create, StateCreator } from "zustand"
 import { computed, ComputedStateOpts } from "./computed"
 
 type Store = {
@@ -132,3 +132,48 @@ describe("custom config", () => {
     expect(computeStateMock).toHaveBeenCalledTimes(4)
   })
 })
+
+describe("slices pattern", () => {
+  const computeStateMock = jest.fn(computeState)
+  const makeStore = () => {
+    const createXySlice = computed<Store, ComputedStore, [], [], Pick<Store, "count" | "dec">>(
+      (set) => ({
+        count: 1,
+        dec: () => set((state) => ({ count: state.count - 1 })),
+      }),
+      computeStateMock
+    )
+    const createCountSlice: StateCreator<Store & ComputedStore, [], [], Pick<Store, "x" | "y" | "inc">> = (set) => ({
+      x: 1,
+      y: 1,
+      // this should not trigger compute function
+      inc: () => set((state) => ({ count: state.count + 2 })),
+    })
+    return create<Store & ComputedStore>()((...a) => ({
+      ...createXySlice(...a),
+      ...createCountSlice(...a),
+    }))
+  }
+
+  beforeEach(() => {
+    computeStateMock.mockClear()
+  })
+
+  test("computed works on slices pattern example", () => {
+    const useStore = makeStore()
+    expect(computeStateMock).toHaveBeenCalledTimes(1)
+    expect(useStore.getState().count).toEqual(1)
+    expect(useStore.getState().countSq).toEqual(1)
+    useStore.getState().inc()
+    expect(useStore.getState().count).toEqual(3)
+    expect(useStore.getState().countSq).toEqual(1)
+    expect(computeStateMock).toHaveBeenCalledTimes(1)
+    useStore.getState().dec()
+    expect(useStore.getState().count).toEqual(2)
+    expect(useStore.getState().countSq).toEqual(4)
+    expect(computeStateMock).toHaveBeenCalledTimes(2)
+    useStore.setState({ count: 4 })
+    expect(useStore.getState().countSq).toEqual(16)
+    expect(computeStateMock).toHaveBeenCalledTimes(3)
+  })
+})

src/computed.ts

@@ -12,11 +12,12 @@ export type ComputedStateCreator = <
   A extends object,
   Mps extends [StoreMutatorIdentifier, unknown][] = [],
   Mcs extends [StoreMutatorIdentifier, unknown][] = [],
+  U = T,
 >(
-  f: StateCreator<T, [...Mps, ["chrisvander/zustand-computed", A]], Mcs>,
+  f: StateCreator<T, [...Mps, ["chrisvander/zustand-computed", A]], Mcs, U>,
   compute: (state: T) => A,
   opts?: ComputedStateOpts<T>,
-) => StateCreator<T, Mps, [["chrisvander/zustand-computed", A], ...Mcs], T & A>
+) => StateCreator<T, Mps, [["chrisvander/zustand-computed", A], ...Mcs], U & A>
 
 type Cast<T, U> = T extends U ? T : U
 type Write<T, U> = Omit<T, keyof U> & U