add redo feature

src/index.ts

@@ -11,17 +11,29 @@ import create from 'zustand';
 
 interface UndoStoreState extends State {
  prevStates: any[];
+ futureStates: any[];
  undo: () => void;
- handle: Function;
+ redo: () => void;
+ setStore: Function;
+ getStore: Function;
 }
 
 // Stores previous actions
 const undoStore = createVanilla<UndoStoreState>((_, get) => ({
  prevStates: [],
  undo: () => {
- get().handle(get().prevStates.pop());
+ const prevState = get().prevStates.pop();
+ get().futureStates.push(get().getStore());
+ get().setStore(prevState);
+ },
+ setStore: () => {},
+ getStore: () => {},
+ futureStates: [],
+ redo: () => {
+ const futureState = get().futureStates.pop();
+ get().prevStates.push(get().getStore());
+ get().setStore(futureState);
  },
- handle: () => {},
 }));
 const { getState, setState } = undoStore;
 export const useUndo = create(undoStore);
@@ -36,7 +48,9 @@ export const undo = <TState extends State>(config: StateCreator<TState>) => (
  args => {
  setState({
  prevStates: [...getState().prevStates, { ...get() }],
- handle: set,
+ setStore: set,
+ futureStates: [],
+ getStore: get,
  });
  set(args);
  },

stories/bears.stories.tsx

@@ -35,21 +35,27 @@ const useStore = create<StoreState>(
 );
 
 const App = () => {
- const { prevStates, undo } = useUndo();
- const { bears, increasePopulation, removeAllBears } = useStore();
+ const { prevStates, undo, futureStates, redo } = useUndo();
+ const store = useStore();
+ const { bears, increasePopulation, removeAllBears } = store;
 
  return (
  <div>
  <h1>🐻 ♻️ Zundo!</h1>
  previous states: {JSON.stringify(prevStates)}
  <br />
+ future states: {JSON.stringify(futureStates)}
+ <br />
+ current state: {JSON.stringify(store)}
+ <br />
  <br />
  bears: {bears}
  <br />
  <button onClick={increasePopulation}>increase</button>
  <button onClick={removeAllBears}>remove</button>
  <br />
  <button onClick={undo}>undo</button>
+ <button onClick={redo}>redo</button>
  </div>
  );
 };