From fe8fa811df7e14e3d6311c2970425f7005360000 Mon Sep 17 00:00:00 2001
From: lawvs <18554747+lawvs@users.noreply.github.com>
Date: Fri, 12 Apr 2024 21:41:36 +0800
Subject: [PATCH] feat: can undo and redo changes

---
 src/components/config-panel.tsx | 41 +++++++++++++++++-
 src/state.ts                    | 74 ++++++++++++++++++++++++++++++++-
 src/y-type.ts                   |  4 +-
 3 files changed, 115 insertions(+), 4 deletions(-)

diff --git a/src/components/config-panel.tsx b/src/components/config-panel.tsx
index 0f9b5c9..6443fff 100644
--- a/src/components/config-panel.tsx
+++ b/src/components/config-panel.tsx
@@ -1,9 +1,11 @@
-import { useConfig, useYDoc } from "../state";
+import { Redo, Undo } from "lucide-react";
+import { useConfig, useUndoManager, useYDoc } from "../state";
 import { fileToYDoc } from "../utils";
 import { ConnectButton } from "./connect-button";
 import { ExportButton } from "./export-button";
 import { FullScreenDropZone } from "./full-screen-drop-zone";
 import { LoadButton } from "./load-button";
+import { Button } from "./ui/button";
 import { Label } from "./ui/label";
 import { Switch } from "./ui/switch";
 import { useToast } from "./ui/use-toast";
@@ -12,6 +14,8 @@ export function ConfigPanel() {
   const [yDoc, setYDoc] = useYDoc();
   const { toast } = useToast();
   const [config, setConfig] = useConfig();
+  const { undoManager, canRedo, canUndo, undoStackSize, redoStackSize } =
+    useUndoManager();
 
   return (
     <div className="flex w-64 flex-col gap-4">
@@ -77,6 +81,41 @@ export function ConfigPanel() {
         <Label htmlFor="editable-switch">Editable</Label>
       </div>
 
+      {config.editable && (
+        <div className="flex items-center space-x-2">
+          <Button
+            className="flex-1"
+            variant="outline"
+            disabled={!canUndo}
+            onClick={() => {
+              if (!undoManager.canUndo()) {
+                console.warn("Cannot undo", undoManager);
+                return;
+              }
+              undoManager.undo();
+            }}
+          >
+            <Undo className="mr-2 h-4 w-4" />
+            Undo({undoStackSize})
+          </Button>
+          <Button
+            className="flex-1"
+            variant="outline"
+            disabled={!canRedo}
+            onClick={() => {
+              if (!undoManager.canRedo()) {
+                console.warn("Cannot redo", undoManager);
+                return;
+              }
+              undoManager.redo();
+            }}
+          >
+            <Redo className="mr-2 h-4 w-4" />
+            Redo({redoStackSize})
+          </Button>
+        </div>
+      )}
+
       <ExportButton />
 
       <FullScreenDropZone
diff --git a/src/state.ts b/src/state.ts
index 060cb17..405f70c 100644
--- a/src/state.ts
+++ b/src/state.ts
@@ -1,12 +1,82 @@
-import { atom, useAtom } from "jotai";
+import { atom, useAtom, useAtomValue } from "jotai";
+import { useEffect, useState } from "react";
 import * as Y from "yjs";
 
-const yDocAtom = atom(new Y.Doc());
+const TRACK_ALL_ORIGINS = Symbol();
+
+function createUndoManager(doc: Y.Doc) {
+  const undoManager = new Y.UndoManager([], {
+    doc,
+    trackedOrigins: new Set([TRACK_ALL_ORIGINS]),
+  });
+  doc.on("update", () => {
+    // The UndoManager can only track shared types that are created
+    // See https://discuss.yjs.dev/t/global-document-undo-manager/2555
+    const keys = Array.from(doc.share.keys());
+    if (!keys.length) return;
+    const scope = keys.map((key) => doc.get(key));
+    undoManager.addToScope(scope);
+    // undoManager.addTrackedOrigin(origin);
+  });
+  doc.on("beforeTransaction", (transaction) => {
+    // Try to track all origins
+    // Workaround for https://github.com/yjs/yjs/issues/624
+    transaction.origin = TRACK_ALL_ORIGINS;
+  });
+  return undoManager;
+}
+
+const defaultYDoc = new Y.Doc();
+const defaultUndoManager = createUndoManager(defaultYDoc);
+
+const undoManagerAtom = atom<Y.UndoManager>(defaultUndoManager);
+
+const yDocAtom = atom(defaultYDoc, (get, set, newDoc: Y.Doc) => {
+  get(undoManagerAtom).destroy();
+  const undoManager = createUndoManager(newDoc);
+  set(undoManagerAtom, undoManager);
+  get(yDocAtom).destroy();
+  set(yDocAtom, newDoc);
+});
 
 export const useYDoc = () => {
   return useAtom(yDocAtom);
 };
 
+export const useUndoManager = () => {
+  const undoManager = useAtomValue(undoManagerAtom);
+  const [state, setState] = useState({
+    canUndo: undoManager.canUndo(),
+    canRedo: undoManager.canRedo(),
+    undoStackSize: undoManager.undoStack.length,
+    redoStackSize: undoManager.redoStack.length,
+  });
+
+  useEffect(() => {
+    const callback = () => {
+      setState({
+        canUndo: undoManager.canUndo(),
+        canRedo: undoManager.canRedo(),
+        undoStackSize: undoManager.undoStack.length,
+        redoStackSize: undoManager.redoStack.length,
+      });
+    };
+    callback();
+
+    undoManager.on("stack-item-added", callback);
+    undoManager.on("stack-item-popped", callback);
+    return () => {
+      undoManager.off("stack-item-added", callback);
+      undoManager.off("stack-item-popped", callback);
+    };
+  }, [state, undoManager]);
+
+  return {
+    undoManager,
+    ...state,
+  };
+};
+
 export type Config = {
   parseYDoc: boolean;
   showDelta: boolean;
diff --git a/src/y-type.ts b/src/y-type.ts
index 026217f..fcf1da8 100644
--- a/src/y-type.ts
+++ b/src/y-type.ts
@@ -19,7 +19,9 @@ export function guessType(abstractType: Y.AbstractType<unknown>) {
     return Y.Map;
   }
   if (abstractType._length > 0) {
-    return Y.Array;
+    // TODO distinguish between Y.Text and Y.Array
+    return Y.Text;
+    // return Y.Array;
   }
   return Y.AbstractType;
 }