From 50ea6e118493069f59cc258615e11ab46a279fa5 Mon Sep 17 00:00:00 2001
From: Logan <logan.zartman@utexas.edu>
Date: Sun, 17 Mar 2024 12:09:37 -0700
Subject: [PATCH] refactor into useUndoable

---
 src/app/reacjin/ReacjinEditor.tsx | 61 +++++++++++---------------
 src/app/reacjin/useUndoable.ts    | 73 +++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 36 deletions(-)
 create mode 100644 src/app/reacjin/useUndoable.ts

diff --git a/src/app/reacjin/ReacjinEditor.tsx b/src/app/reacjin/ReacjinEditor.tsx
index e1d4096..f4acc22 100644
--- a/src/app/reacjin/ReacjinEditor.tsx
+++ b/src/app/reacjin/ReacjinEditor.tsx
@@ -29,6 +29,7 @@ import {Panel} from '@/app/reacjin/Panel';
 import {PanelProvider} from '@/app/reacjin/PanelContext';
 import {pluginByID} from '@/app/reacjin/plugins/registry';
 import {Toolbar} from '@/app/reacjin/Toolbar';
+import {useUndoable} from '@/app/reacjin/useUndoable';
 import {MotionDiv} from '@/lib/framer-motion';
 
 export default function ReacjinEditor() {
@@ -36,59 +37,47 @@ export default function ReacjinEditor() {
   const editorAreaRef = useRef<HTMLDivElement>(null);
   const [imageSize, setImageSize] = useState([256, 256]);
   const [zoom, setZoom] = useState(1);
-  const [layers, _setLayers] = useState<Layers>(() => [
-    createLayer('text', {
-      text: 'Hello, world!',
-      autoFitText: false,
-      fontSize: 32,
-      fontFamily: 'sans-serif',
-      fillStyle: 'white',
-      strokeStyle: 'black',
-      strokeWidth: 2,
-      textAlign: 'center',
-      lineHeight: 1.1,
-    }),
-    createLayer('image', {src: 'https://picsum.photos/256'}),
-  ]);
-  const [undoStack, setUndoStack] = useState<Layers[]>([]);
-  const [redoStack, setRedoStack] = useState<Layers[]>([]);
+  const {
+    state: layers,
+    setState: setLayers,
+    undo,
+    redo,
+  } = useUndoable(
+    useState<Layers>(() => [
+      createLayer('text', {
+        text: 'Hello, world!',
+        autoFitText: false,
+        fontSize: 32,
+        fontFamily: 'sans-serif',
+        fillStyle: 'white',
+        strokeStyle: 'black',
+        strokeWidth: 2,
+        textAlign: 'center',
+        lineHeight: 1.1,
+      }),
+      createLayer('image', {src: 'https://picsum.photos/256'}),
+    ]),
+  );
   const [selectedLayerID, setSelectedLayerID] = useState<string | null>(null);
   const [computedCache] = useState(() => new ComputedCache());
   const [computing, setComputing] = useState(false);
   const [dropping, setDropping] = useState(false);
 
-  const setLayers = useCallback(
-    (action: React.SetStateAction<Layers>) => {
-      setUndoStack((undoStack) => [...undoStack, layers]);
-      setRedoStack([]);
-      _setLayers((layers) =>
-        typeof action === 'function' ? action(layers) : action,
-      );
-    },
-    [layers],
-  );
-
   useEffect(() => {
     function handler(event: KeyboardEvent) {
       if (event.repeat) return;
       if (event.key === 'z' && (event.metaKey || event.ctrlKey)) {
         event.preventDefault();
-        if (undoStack.length < 1) return;
-        _setLayers(undoStack[undoStack.length - 1]);
-        setUndoStack((undoStack) => undoStack.slice(0, -1));
-        setRedoStack((redoStack) => [...redoStack, layers]);
+        undo();
       }
       if (event.key === 'y' && (event.metaKey || event.ctrlKey)) {
         event.preventDefault();
-        if (redoStack.length < 1) return;
-        _setLayers(redoStack[redoStack.length - 1]);
-        setRedoStack((redoStack) => redoStack.slice(0, -1));
-        setUndoStack((undoStack) => [...undoStack, layers]);
+        redo();
       }
     }
     document.addEventListener('keydown', handler, false);
     return () => document.removeEventListener('keydown', handler, false);
-  }, [layers, redoStack, undoStack]);
+  }, [layers, redo, undo]);
 
   const setZoomToSize = useCallback(
     (size: number) => {
diff --git a/src/app/reacjin/useUndoable.ts b/src/app/reacjin/useUndoable.ts
new file mode 100644
index 0000000..553eb85
--- /dev/null
+++ b/src/app/reacjin/useUndoable.ts
@@ -0,0 +1,73 @@
+import React, {useCallback, useMemo, useRef, useState} from 'react';
+
+type Undoable<T> = {
+  state: T;
+  setState: React.Dispatch<React.SetStateAction<T>>;
+  undo: () => void;
+  redo: () => void;
+  canUndo: boolean;
+  canRedo: boolean;
+};
+
+function isStateUpdater<T>(
+  action: React.SetStateAction<T>,
+): action is (prevState: T) => T {
+  return typeof action === 'function';
+}
+
+export function useUndoable<T>(
+  stateTuple: [T, React.Dispatch<React.SetStateAction<T>>],
+): Undoable<T> {
+  const [state, _setState] = stateTuple;
+  const [undoStack, setUndoStack] = useState<T[]>([]);
+  const [redoStack, setRedoStack] = useState<T[]>([]);
+  const stateRef = useRef<T>(state);
+  const undoStackRef = useRef<T[]>([]);
+  const redoStackRef = useRef<T[]>([]);
+  const canUndo = undoStack.length > 0;
+  const canRedo = redoStack.length > 0;
+
+  stateRef.current = state;
+  undoStackRef.current = undoStack;
+  redoStackRef.current = redoStack;
+
+  const setState: React.Dispatch<React.SetStateAction<T>> = useCallback(
+    (action) => {
+      const state = stateRef.current;
+      setUndoStack((undoStack) => [...undoStack, state]);
+      setRedoStack([]);
+      _setState((state) => (isStateUpdater(action) ? action(state) : action));
+    },
+    [_setState],
+  );
+
+  const undo = useCallback(() => {
+    const state = stateRef.current;
+    const undoStack = undoStackRef.current;
+    if (undoStack.length < 1) return;
+    _setState(undoStack[undoStack.length - 1]);
+    setUndoStack((undoStack) => undoStack.slice(0, -1));
+    setRedoStack((redoStack) => [...redoStack, state]);
+  }, [_setState]);
+
+  const redo = useCallback(() => {
+    const state = stateRef.current;
+    const redoStack = redoStackRef.current;
+    if (redoStack.length < 1) return;
+    _setState(redoStack[redoStack.length - 1]);
+    setRedoStack((redoStack) => redoStack.slice(0, -1));
+    setUndoStack((undoStack) => [...undoStack, state]);
+  }, [_setState]);
+
+  return useMemo(
+    () => ({
+      state,
+      setState,
+      undo,
+      redo,
+      canUndo,
+      canRedo,
+    }),
+    [state, setState, undo, redo, canUndo, canRedo],
+  );
+}