Merge pull request #66 from igloo-4002/feat/undo-redo-stack

feat: undo redo stack

RESOURCES.md

@@ -6,6 +6,10 @@
 
 - [Zustand](https://docs.pmnd.rs/zustand/getting-started/introduction)
 
+### Undo/Redo
+
+- [Command pattern](https://en.wikipedia.org/wiki/Command_pattern)
+
 ### SUMO - (Simulation of Urban MObility)
 
 [SUMO user documentation](https://sumo.dlr.de/docs/index.html)

src/components/Canvas/Canvas.tsx

@@ -15,6 +15,7 @@ import {
   useStageState,
 } from '~/zustand/useStage';
 import { useToolbarStore } from '~/zustand/useToolbar';
+import { useUndoStore } from '~/zustand/useUndoStore';
 
 import { CarLayer } from './Layers/CarLayer';
 import { ConnectionsLayer } from './Layers/ConnectionsLayer';
@@ -26,6 +27,7 @@ export function Canvas() {
   const network = useNetworkStore();
   const nodes = Object.values(network.nodes);
   const toolbarState = useToolbarStore();
+  const undoStore = useUndoStore();
 
   // Keyboard shortcuts
   useEffect(() => {
@@ -45,6 +47,18 @@ export function Canvas() {
         }
       }
 
+      const isUndoCommand =
+        (e.metaKey || e.ctrlKey) && !e.shiftKey && e.key === 'z';
+      const isRedoCommand =
+        (e.metaKey || e.ctrlKey) && e.shiftKey && e.key === 'z';
+
+      if (isUndoCommand) {
+        undoStore.undo();
+      }
+      if (isRedoCommand) {
+        undoStore.redo();
+      }
+
       if (isEsc && selector.selected) {
         selector.deselect();
       }

src/components/ProjectUploadButton.tsx

@@ -3,10 +3,12 @@ import { ArrowUpTrayIcon } from '@heroicons/react/20/solid';
 import { getNetworkFromUploadedFile } from '~/logic/urbanflo-file-logic';
 import { useNetworkStore } from '~/zustand/useNetworkStore';
 import { useSimulationHistory } from '~/zustand/useSimulationHistory';
+import { useUndoStore } from '~/zustand/useUndoStore';
 
 export function ProjectUploadButton() {
   const networkStore = useNetworkStore();
   const simulationHistoryStore = useSimulationHistory();
+  const undoStore = useUndoStore();
 
   function handleFileUpload(event: React.ChangeEvent<HTMLInputElement>) {
     const file = event.target.files?.[0];
@@ -40,6 +42,10 @@ export function ProjectUploadButton() {
               simulationHistoryStore.updateHistory(historyItem);
             },
           );
+
+          // Since the deleteNode, addNode, and drawEdge methods add to the undo stack,
+          // we want to clear after building the network.
+          undoStore.clearStacks();
         } catch (error) {
           console.error(
             'An error occurred while parsing the JSON file OR the file is invalid',

src/helpers/commands/AddEdgeCommand.ts

@@ -0,0 +1,25 @@
+import { Edge } from '~/types/Network';
+import { Network } from '~/zustand/useNetworkStore';
+import { Command } from '~/zustand/useUndoStore';
+
+export class AddEdgeCommand implements Command {
+  constructor(
+    private network: Network,
+    private edge: Edge,
+  ) {}
+
+  execute() {
+    const fromNode = this.network.nodes[this.edge.from];
+    const toNode = this.network.nodes[this.edge.to];
+
+    if (!fromNode && !toNode) {
+      return;
+    }
+
+    this.network.drawEdge(fromNode, toNode);
+  }
+
+  unexecute() {
+    this.network.deleteEdge(this.edge.id);
+  }
+}

src/helpers/commands/AddNodeCommand.ts

@@ -0,0 +1,18 @@
+import { Node } from '~/types/Network';
+import { Network } from '~/zustand/useNetworkStore';
+import { Command } from '~/zustand/useUndoStore';
+
+export class AddNodeCommand implements Command {
+  constructor(
+    private network: Network,
+    private node: Node,
+  ) {}
+
+  execute() {
+    this.network.addNode(this.node);
+  }
+
+  unexecute() {
+    this.network.deleteNode(this.node.id);
+  }
+}

src/helpers/commands/RemoveEdgeCommand.ts

@@ -0,0 +1,25 @@
+import { Edge } from '~/types/Network';
+import { Network } from '~/zustand/useNetworkStore';
+import { Command } from '~/zustand/useUndoStore';
+
+export class RemoveEdgeCommand implements Command {
+  constructor(
+    private network: Network,
+    private edge: Edge,
+  ) {}
+
+  execute() {
+    this.network.deleteEdge(this.edge.id);
+  }
+
+  unexecute() {
+    const fromNode = this.network.nodes[this.edge.from];
+    const toNode = this.network.nodes[this.edge.to];
+
+    if (!fromNode && !toNode) {
+      return;
+    }
+
+    this.network.drawEdge(fromNode, toNode);
+  }
+}

src/helpers/commands/RemoveNodeCommand.ts

@@ -0,0 +1,18 @@
+import { Node } from '~/types/Network';
+import { Network } from '~/zustand/useNetworkStore';
+import { Command } from '~/zustand/useUndoStore';
+
+export class RemoveNodeCommand implements Command {
+  constructor(
+    private network: Network,
+    private node: Node,
+  ) {}
+
+  execute() {
+    this.network.deleteNode(this.node.id);
+  }
+
+  unexecute() {
+    this.network.addNode(this.node);
+  }
+}

src/zustand/useNetworkStore.ts

@@ -1,6 +1,10 @@
 import { create } from 'zustand';
 
 import { laneWidth } from '~/components/Canvas/Road';
+import { AddEdgeCommand } from '~/helpers/commands/AddEdgeCommand';
+import { AddNodeCommand } from '~/helpers/commands/AddNodeCommand';
+import { RemoveEdgeCommand } from '~/helpers/commands/RemoveEdgeCommand';
+import { RemoveNodeCommand } from '~/helpers/commands/RemoveNodeCommand';
 import {
   edgeDoesIntersect,
   removeItems,
@@ -9,6 +13,8 @@ import {
 } from '~/helpers/zustand/NetworkStoreHelpers';
 import { Connection, Edge, Flow, Node, Route, VType } from '~/types/Network';
 
+import { useUndoStore } from './useUndoStore';
+
 export interface NetworkData {
   documentName: string;
   nodes: Record<string, Node>;
@@ -31,7 +37,7 @@ export interface Network extends NetworkData {
   updateFlow: (flowId: string, flow: Flow) => void;
 }
 
-export const useNetworkStore = create<Network>(set => ({
+export const useNetworkStore = create<Network>((set, get) => ({
   documentName: 'Untitled Document',
   nodes: {},
   edges: {},
@@ -43,8 +49,11 @@ export const useNetworkStore = create<Network>(set => ({
   setDocumentName: name => {
     set({ documentName: name });
   },
-  addNode: (node: Node) =>
-    set(state => ({ nodes: { ...state.nodes, [node.id]: node } })),
+  addNode: (node: Node) => {
+    const undoStore = useUndoStore.getState();
+    undoStore.pushCommand(new AddNodeCommand(get(), node));
+    set(state => ({ nodes: { ...state.nodes, [node.id]: node } }));
+  },
   updateNode: (nodeId, node) => {
     set(state => {
       return {
@@ -55,7 +64,9 @@ export const useNetworkStore = create<Network>(set => ({
       };
     });
   },
-  drawEdge: (from, to) =>
+  drawEdge: (from, to) => {
+    const undoStore = useUndoStore.getState();
+
     set(state => {
       const newEdgeId = `${from.id}_${to.id}`;
       const newEdge: Edge = {
@@ -76,6 +87,8 @@ export const useNetworkStore = create<Network>(set => ({
       if (edgeDoesIntersect(state, pointA, pointB)) {
         return state;
       } else {
+        undoStore.pushCommand(new AddEdgeCommand(get(), newEdge));
+
         // update connections, routes, and flows when an edge is being drawn
         const { newConnections, newRoutes, newFlows } =
           updateAssociatesOnNewEdge(
@@ -93,7 +106,8 @@ export const useNetworkStore = create<Network>(set => ({
           route: newRoutes,
         };
       }
-    }),
+    });
+  },
   updateEdge: (edgeId, edge) => {
     set(state => {
       const updatedEdges = {
@@ -194,6 +208,8 @@ export const useNetworkStore = create<Network>(set => ({
     });
   },
   deleteNode: (id: string) => {
+    const undoStore = useUndoStore.getState();
+    undoStore.pushCommand(new RemoveNodeCommand(get(), get().nodes[id]));
     set(state => {
       const newNodes = { ...state.nodes };
       delete newNodes[id];
@@ -233,6 +249,9 @@ export const useNetworkStore = create<Network>(set => ({
     });
   },
   deleteEdge: (id: string) => {
+    const undoStore = useUndoStore.getState();
+    undoStore.pushCommand(new RemoveEdgeCommand(get(), get().edges[id]));
+
     set(state => {
       const newEdges = { ...state.edges };
       delete newEdges[id];

src/zustand/useUndoStore.ts

@@ -0,0 +1,55 @@
+import { create } from 'zustand';
+
+export interface Command {
+  execute(): void;
+  unexecute(): void;
+}
+
+interface UndoRedoStore {
+  undoStack: Command[];
+  redoStack: Command[];
+  pushCommand: (command: Command) => void;
+  undo: () => void;
+  redo: () => void;
+  clearStacks: () => void;
+  setStacks: (undoStack: Command[], redoStack: Command[]) => void;
+}
+
+export const useUndoStore = create<UndoRedoStore>((set, get) => ({
+  undoStack: [],
+  redoStack: [],
+  pushCommand: command =>
+    set(state => ({ undoStack: [...state.undoStack, command] })),
+  undo: () => {
+    const { undoStack, redoStack } = get();
+    const lastCommand = undoStack.pop();
+
+    if (!lastCommand) {
+      return;
+    }
+
+    lastCommand.unexecute();
+
+    set({
+      redoStack: [...redoStack, lastCommand],
+      undoStack: [...undoStack],
+    });
+  },
+  redo: () => {
+    const { undoStack, redoStack } = get();
+    const lastCommand = redoStack.pop();
+
+    if (!lastCommand) {
+      return;
+    }
+
+    lastCommand.execute();
+
+    set({
+      redoStack: [...redoStack],
+      undoStack: [...undoStack, lastCommand],
+    });
+  },
+  clearStacks: () => set({ undoStack: [], redoStack: [] }),
+  setStacks: (undoStack, redoStack) => set({ undoStack, redoStack }),
+}));