switch config for better types and save bytes (#83)

Author: charkour

packages/zundo/__tests__/options.test.ts

@@ -431,9 +431,9 @@ describe('Middleware options', () => {
 
   describe('secret internals', () => {
     it('should have a secret internal state', () => {
-      const { __handleUserSet, __onSave } =
+      const { __handleSet, __onSave } =
         store.temporal.getState() as TemporalStateWithInternals<MyState>;
-      expect(__handleUserSet).toBeInstanceOf(Function);
+      expect(__handleSet).toBeInstanceOf(Function);
       expect(__onSave).toBe(undefined);
     });
     describe('onSave', () => {
@@ -508,29 +508,29 @@ describe('Middleware options', () => {
 
     describe('handleUserSet', () => {
       it('should update the temporal store with the pastState when called', () => {
-        const { __handleUserSet } =
+        const { __handleSet } =
           store.temporal.getState() as TemporalStateWithInternals<MyState>;
         act(() => {
-          __handleUserSet(store.getState());
+          __handleSet(store.getState());
         });
         expect(store.temporal.getState().pastStates.length).toBe(1);
       });
 
       it('should only update if the the status is tracking', () => {
-        const { __handleUserSet } =
+        const { __handleSet } =
           store.temporal.getState() as TemporalStateWithInternals<MyState>;
         act(() => {
-          __handleUserSet(store.getState());
+          __handleSet(store.getState());
         });
         expect(store.temporal.getState().pastStates.length).toBe(1);
         act(() => {
           store.temporal.getState().pause();
-          __handleUserSet(store.getState());
+          __handleSet(store.getState());
         });
         expect(store.temporal.getState().pastStates.length).toBe(1);
         act(() => {
           store.temporal.getState().resume();
-          __handleUserSet(store.getState());
+          __handleSet(store.getState());
         });
       });
 

packages/zundo/src/index.ts

@@ -5,7 +5,12 @@ import type {
   StoreApi,
 } from 'zustand';
 import { createVanillaTemporal } from './temporal';
-import type { TemporalState, Write, ZundoOptions } from './types';
+import type {
+  TemporalState,
+  TemporalStateWithInternals,
+  Write,
+  ZundoOptions,
+} from './types';
 
 type Zundo = <
   TState,
@@ -27,51 +32,61 @@ declare module 'zustand/vanilla' {
   }
 }
 
-const zundoImpl =
-  <TState>(
-    config: StateCreator<TState, [], []>,
-    {
-      partialize = (state: TState) => state,
-      handleSet: userlandSetFactory = (handleSetCb) => handleSetCb,
-      ...restOptions
-    } = {} as ZundoOptions<TState>,
-  ): StateCreator<TState, [], []> =>
-  (set, get, _store) => {
-    type TState = ReturnType<typeof config>;
-    type StoreAddition = StoreApi<TemporalState<TState>>;
-
-    const temporalStore = createVanillaTemporal<TState>(set, get, partialize, restOptions);
-
-    const store = _store as Mutate<
-      StoreApi<TState>,
-      [['temporal', StoreAddition]]
-    >;
-    const setState = store.setState;
-
-    // TODO: should temporal be only temporalStore.getState()?
-    // We can hide the rest of the store in the secret internals.
-    store.temporal = temporalStore;
+const zundoImpl = <TState>(
+  config: StateCreator<TState, [], []>,
+  {
+    partialize = (state: TState) => state,
+    handleSet = (handleSetCb) => handleSetCb,
+    ...restOptions
+  } = {} as ZundoOptions<TState>,
+): StateCreator<TState, [], []> => {
+  type StoreAddition = StoreApi<TemporalState<TState>>;
+  type StoreWithAddition = Mutate<
+    StoreApi<TState>,
+    [['temporal', StoreAddition]]
+  >;
+  const configWithTemporal = (
+    set: StoreApi<TState>['setState'],
+    get: StoreApi<TState>['getState'],
+    store: StoreWithAddition,
+  ) => {
+    store.temporal = createVanillaTemporal<TState>(
+      set,
+      get,
+      partialize,
+      restOptions,
+    );
 
-    const curriedUserLandSet = userlandSetFactory(
-      temporalStore.getState().__handleUserSet,
+    const curriedHandleSet = handleSet(
+      (store.temporal.getState() as TemporalStateWithInternals<TState>)
+        .__handleSet,
     );
 
-    const modifiedSetState: typeof setState = (state, replace) => {
+    const setState = store.setState;
+    // Modify the setState function to call the userlandSet function
+    store.setState = (state, replace) => {
+      // Get most up to date state. The state from the callback might be a partial state.
+      // The order of the get() and set() calls is important here.
       const pastState = partialize(get());
       setState(state, replace);
-      curriedUserLandSet(pastState);
+      curriedHandleSet(pastState);
     };
-    store.setState = modifiedSetState;
 
-    const modifiedSetter: typeof set = (state, replace) => {
-      // Get most up to date state. Should this be the same as the state in the callback?
-      const pastState = partialize(get());
-      set(state, replace);
-      curriedUserLandSet(pastState);
-    };
-
-    return config(modifiedSetter, get, _store);
+    return config(
+      // Modify the set function to call the userlandSet function
+      (state, replace) => {
+        // Get most up to date state. The state from the callback might be a partial state.
+        // The order of the get() and set() calls is important here.
+        const pastState = partialize(get());
+        set(state, replace);
+        curriedHandleSet(pastState);
+      },
+      get,
+      store,
+    );
   };
+  return configWithTemporal as StateCreator<TState, [], []>;
+};
 
 export const temporal = zundoImpl as unknown as Zundo;
 export type { ZundoOptions, Zundo, TemporalState };

packages/zundo/src/temporal.ts

@@ -11,7 +11,8 @@ export const createVanillaTemporal = <TState>(
   partialize: (state: TState) => TState,
   { equality, onSave, limit } = {} as Omit<ZundoOptions<TState>, 'handleSet'>,
 ) => {
-  return createStore<TemporalStateWithInternals<TState>>()((set, get) => {
+
+  return createStore<TemporalStateWithInternals<TState>>((set, get) => {
     return {
       pastStates: [],
       futureStates: [],
@@ -68,7 +69,7 @@ export const createVanillaTemporal = <TState>(
       },
       // Internal properties
       __onSave: onSave,
-      __handleUserSet: (pastState) => {
+      __handleSet: (pastState) => {
         const trackingStatus = get().trackingStatus,
           onSave = get().__onSave,
           pastStates = get().pastStates.slice(),
@@ -77,6 +78,7 @@ export const createVanillaTemporal = <TState>(
           trackingStatus === 'tracking' &&
           !equality?.(currentState, pastState)
         ) {
+          // This naively assumes that only one new state can be added at a time
           if (limit && pastStates.length >= limit) {
             pastStates.shift();
           }

packages/zundo/src/types.ts

@@ -16,7 +16,7 @@ export interface TemporalStateWithInternals<TState> {
 
   setOnSave: (onSave: onSave<TState>) => void;
   __onSave: onSave<TState>;
-  __handleUserSet: (pastState: TState) => void;
+  __handleSet: (pastState: TState) => void;
 }
 
 export interface ZundoOptions<TState, PartialTState = TState> {