From 1c61e1fca179a6116bb07b17dd5be72f6e193387 Mon Sep 17 00:00:00 2001
From: Charles Kornoelje <33156025+charkour@users.noreply.github.com>
Date: Sat, 10 Apr 2021 21:01:43 -0400
Subject: [PATCH] add redo feature

---
 src/index.ts              | 22 ++++++++++++++++++----
 stories/bears.stories.tsx | 10 ++++++++--
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 952109c..35d8122 100644
--- a/src/index.ts
+++ b/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);
     },
diff --git a/stories/bears.stories.tsx b/stories/bears.stories.tsx
index a781d67..f78ac0a 100644
--- a/stories/bears.stories.tsx
+++ b/stories/bears.stories.tsx
@@ -35,14 +35,19 @@ 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 />
@@ -50,6 +55,7 @@ const App = () => {
       <button onClick={removeAllBears}>remove</button>
       <br />
       <button onClick={undo}>undo</button>
+      <button onClick={redo}>redo</button>
     </div>
   );
 };