yorkie-team · LakHyeonKim · Sep 15, 2023
diff --git a/examples/react-tldraw/src/hooks/types.ts b/examples/react-tldraw/src/hooks/types.ts
@@ -13,6 +13,23 @@ export type YorkieDocType = {
   assets: JSONObject<Record<string, JSONObject<TDAsset>>>;
 };
 
+export type TlType = {
+  shapes: Record<string, JSONObject<TDShape>>;
+  bindings: Record<string, JSONObject<TDBinding>>;
+  assets: Record<string, JSONObject<TDAsset>>;
+};
+
+export type HistoryType = {
+  undoStack: Array<CommandType>;
+  redoStack: Array<CommandType>;
+};
+
+export type CommandType = {
+  snapshot: TlType;
+  undo: () => void;
+  redo: () => void;
+};
+
 export type YorkiePresenceType = {
   tdUser: TDUser;
 };
diff --git a/examples/react-tldraw/src/hooks/useMultiplayerState.ts b/examples/react-tldraw/src/hooks/useMultiplayerState.ts
@@ -13,8 +13,14 @@ import * as yorkie from 'yorkie-js-sdk';
 import randomColor from 'randomcolor';
 import { uniqueNamesGenerator, names } from 'unique-names-generator';
 import _ from 'lodash';
+import useUndoRedo from './useUndoRedo';
 
-import type { Options, YorkieDocType, YorkiePresenceType } from './types';
+import type {
+  Options,
+  YorkieDocType,
+  YorkiePresenceType,
+  TlType,
+} from './types';
 
 // Yorkie Client declaration
 let client: yorkie.Client;
@@ -25,6 +31,7 @@ let doc: yorkie.Document<YorkieDocType, YorkiePresenceType>;
 export function useMultiplayerState(roomId: string) {
   const [app, setApp] = useState<TldrawApp>();
   const [loading, setLoading] = useState(true);
+  const { push, undo, redo } = useUndoRedo();
 
   // Callbacks --------------
 
@@ -55,6 +62,37 @@ export function useMultiplayerState(roomId: string) {
     [roomId],
   );
 
+  // undo
+
+  const onUndo = useCallback(() => {
+    undo();
+  }, [roomId]);
+
+  // redo
+
+  const onRedo = useCallback(() => {
+    redo();
+  }, [roomId]);
+
+  // Subscribe to changes
+  function handleChanges() {
+    const root = doc.getRoot();
+
+    // Parse proxy object to record
+    const shapeRecord: Record<string, TDShape> = JSON.parse(
+      root.shapes.toJSON!(),
+    );
+    const bindingRecord: Record<string, TDBinding> = JSON.parse(
+      root.bindings.toJSON!(),
+    );
+    const assetRecord: Record<string, TDAsset> = JSON.parse(
+      root.assets.toJSON!(),
+    );
+
+    // Replace page content with changed(propagated) records
+    app?.replacePageContent(shapeRecord, bindingRecord, assetRecord);
+  }
+
   // Update Yorkie doc when the app's shapes change.
   // Prevent overloading yorkie update api call by throttle
   const onChangePage = useThrottleCallback(
@@ -65,6 +103,13 @@ export function useMultiplayerState(roomId: string) {
     ) => {
       if (!app || client === undefined || doc === undefined) return;
 
+      // Object that stores the latest state value of yorkie doc before the client changes
+      const currentYorkieDocSnapshot: TlType = {
+        shapes: {},
+        bindings: {},
+        assets: {},
+      };
+
       const getUpdatedPropertyList = <T extends object>(
         source: T,
         target: T,
@@ -76,18 +121,26 @@ export function useMultiplayerState(roomId: string) {
 
       Object.entries(shapes).forEach(([id, shape]) => {
         doc.update((root) => {
+          const rootShapesToJS = root.shapes.toJS!();
           if (!shape) {
+            currentYorkieDocSnapshot.shapes[id] = rootShapesToJS[id];
             delete root.shapes[id];
           } else if (!root.shapes[id]) {
+            currentYorkieDocSnapshot.shapes[id] = undefined!;
             root.shapes[id] = shape;
           } else {
             const updatedPropertyList = getUpdatedPropertyList(
               shape,
-              root.shapes[id]!.toJS!(),
+              rootShapesToJS[id],
             );
-
+            currentYorkieDocSnapshot.shapes[id] =
+              {} as yorkie.JSONObject<TDShape>;
             updatedPropertyList.forEach((key) => {
               const newValue = shape[key];
+              const snapshotValue = rootShapesToJS[id][key];
+              (currentYorkieDocSnapshot.shapes[id][
+                key
+              ] as typeof snapshotValue) = snapshotValue;
               (root.shapes[id][key] as typeof newValue) = newValue;
             });
           }
@@ -96,18 +149,26 @@ export function useMultiplayerState(roomId: string) {
 
       Object.entries(bindings).forEach(([id, binding]) => {
         doc.update((root) => {
+          const rootBindingsToJS = root.bindings.toJS!();
           if (!binding) {
+            currentYorkieDocSnapshot.bindings[id] = rootBindingsToJS[id];
             delete root.bindings[id];
           } else if (!root.bindings[id]) {
+            currentYorkieDocSnapshot.bindings[id] = undefined!;
             root.bindings[id] = binding;
           } else {
             const updatedPropertyList = getUpdatedPropertyList(
               binding,
-              root.bindings[id]!.toJS!(),
+              rootBindingsToJS[id],
             );
-
+            currentYorkieDocSnapshot.bindings[id] =
+              {} as yorkie.JSONObject<TDBinding>;
             updatedPropertyList.forEach((key) => {
               const newValue = binding[key];
+              const snapshotValue = rootBindingsToJS[id][key];
+              (currentYorkieDocSnapshot.bindings[id][
+                key
+              ] as typeof snapshotValue) = snapshotValue;
               (root.bindings[id][key] as typeof newValue) = newValue;
             });
           }
@@ -118,14 +179,16 @@ export function useMultiplayerState(roomId: string) {
       // Document key for assets should be asset.id (string), not index
       Object.entries(app.assets).forEach(([, asset]) => {
         doc.update((root) => {
+          const rootAssetsToJS = root.assets.toJS!();
+          currentYorkieDocSnapshot.assets[asset.id] = rootAssetsToJS[asset.id];
           if (!asset.id) {
             delete root.assets[asset.id];
-          } else if (root.assets[asset.id]) {
+          } else if (!root.assets[asset.id]) {
             root.assets[asset.id] = asset;
           } else {
             const updatedPropertyList = getUpdatedPropertyList(
               asset,
-              root.assets[asset.id]!.toJS!(),
+              rootAssetsToJS[asset.id],
             );
 
             updatedPropertyList.forEach((key) => {
@@ -135,8 +198,112 @@ export function useMultiplayerState(roomId: string) {
           }
         });
       });
+
+      // Command object for action
+      // Undo, redo work the same way
+      // undo(): Save yorkie doc's state before returning
+      // redo(): Save yorkie doc's state before moving forward
+      const command = {
+        snapshot: currentYorkieDocSnapshot,
+        undo: () => {
+          const currentYorkieDocSnapshot: TlType = {
+            shapes: {},
+            bindings: {},
+            assets: {},
+          };
+          const snapshot = command.snapshot;
+          Object.entries(snapshot.shapes).forEach(([id, shape]) => {
+            doc.update((root) => {
+              const rootShapesToJS = root.shapes.toJS!();
+              if (!shape) {
+                currentYorkieDocSnapshot.shapes[id] = rootShapesToJS[id];
+                delete root.shapes[id];
+              } else if (!root.shapes.toJS!()[id]) {
+                currentYorkieDocSnapshot.shapes[id] = undefined!;
+                if (shape.id) root.shapes[id] = shape;
+              } else {
+                currentYorkieDocSnapshot.shapes[id] =
+                  {} as yorkie.JSONObject<TDShape>;
+                (
+                  Object.keys(snapshot.shapes[id]) as Array<keyof TDShape>
+                ).forEach((key) => {
+                  const snapshotValue = snapshot.shapes[id][key];
+                  const newSnapshotValue = rootShapesToJS[id][key];
+
+                  (currentYorkieDocSnapshot.shapes[id][
+                    key
+                  ] as typeof newSnapshotValue) = newSnapshotValue;
+                  (root.shapes[id][key] as typeof snapshotValue) =
+                    snapshotValue;
+                });
+              }
+            });
+          });
+
+          Object.entries(snapshot.bindings).forEach(([id, binding]) => {
+            doc.update((root) => {
+              const rootBindingsToJs = root.bindings.toJS!();
+              if (!binding) {
+                currentYorkieDocSnapshot.bindings[id] = rootBindingsToJs[id];
+                delete root.bindings[id];
+              } else if (!root.bindings.toJS!()[id]) {
+                currentYorkieDocSnapshot.bindings[id] = undefined!;
+                if (binding.id) root.bindings[id] = binding;
+              } else {
+                currentYorkieDocSnapshot.bindings[id] =
+                  {} as yorkie.JSONObject<TDBinding>;
+                (
+                  Object.keys(snapshot.bindings[id]) as Array<keyof TDBinding>
+                ).forEach((key) => {
+                  const snapshotValue = snapshot.bindings[id][key];
+                  const newSnapshotValue = rootBindingsToJs[id][key];
+
+                  (currentYorkieDocSnapshot.bindings[id][
+                    key
+                  ] as typeof newSnapshotValue) = newSnapshotValue;
+                  (root.bindings[id][key] as typeof snapshotValue) =
+                    snapshotValue;
+                });
+              }
+            });
+          });
+
+          Object.entries(snapshot.assets).forEach(([, asset]) => {
+            doc.update((root) => {
+              const rootAssetsToJs = root.assets.toJS!();
+              currentYorkieDocSnapshot.assets[asset.id] =
+                rootAssetsToJs[asset.id];
+              if (!asset.id) {
+                delete root.assets[asset.id];
+              } else if (!root.assets.toJS!()[asset.id]) {
+                root.assets[asset.id] = asset;
+              } else {
+                const updatedPropertyList = getUpdatedPropertyList(
+                  asset,
+                  rootAssetsToJs[asset.id],
+                );
+
+                updatedPropertyList.forEach((key) => {
+                  const newValue = asset[key];
+                  (root.assets[asset.id][key] as typeof newValue) = newValue;
+                });
+              }
+            });
+          });
+          command.snapshot = currentYorkieDocSnapshot;
+          // Reflect changes locally
+          handleChanges();
+        },
+        redo: () => {
+          command.undo();
+          handleChanges();
+        },
+      };
+
+      // Create History
+      push(command);
     },
-    60,
+    20,
     false,
   );
 
@@ -168,25 +335,6 @@ export function useMultiplayerState(roomId: string) {
 
     window.addEventListener('beforeunload', handleDisconnect);
 
-    // Subscribe to changes
-    function handleChanges() {
-      const root = doc.getRoot();
-
-      // Parse proxy object to record
-      const shapeRecord: Record<string, TDShape> = JSON.parse(
-        root.shapes.toJSON!(),
-      );
-      const bindingRecord: Record<string, TDBinding> = JSON.parse(
-        root.bindings.toJSON!(),
-      );
-      const assetRecord: Record<string, TDAsset> = JSON.parse(
-        root.assets.toJSON!(),
-      );
-
-      // Replace page content with changed(propagated) records
-      app?.replacePageContent(shapeRecord, bindingRecord, assetRecord);
-    }
-
     let stillAlive = true;
 
     // Setup the document's storage and subscriptions
@@ -294,5 +442,7 @@ export function useMultiplayerState(roomId: string) {
     onChangePage,
     loading,
     onChangePresence,
+    onUndo,
+    onRedo,
   };
 }
diff --git a/examples/react-tldraw/src/hooks/useUndoRedo.ts b/examples/react-tldraw/src/hooks/useUndoRedo.ts
@@ -0,0 +1,37 @@
+import type { CommandType, HistoryType } from './types';
+
+const history: HistoryType = {
+  undoStack: [],
+  redoStack: [],
+};
+
+const useUndoRedo = () => {
+  const { undoStack, redoStack } = history;
+
+  const push = (command: CommandType) => {
+    undoStack.push(command);
+    redoStack.length = 0;
+  };
+
+  const undo = () => {
+    if (undoStack.length === 0) return;
+    const command: CommandType | undefined = undoStack.pop();
+    if (command) {
+      command.undo();
+      redoStack.push(command);
+    }
+  };
+
+  const redo = () => {
+    if (redoStack.length === 0) return;
+    const command: CommandType | undefined = redoStack.pop();
+    if (command) {
+      command.redo();
+      undoStack.push(command);
+    }
+  };
+
+  return { push, undo, redo };
+};
+
+export default useUndoRedo;