Implementing a Undo/Redo Hook in React

Here’s the story (as code) of how I went from typical OOP to React hooks paradigm mind set on this particular problem.

Undo/Redo is a great invention that is not simple to manage and users expect it everywhere they use software. Understanding how to implement the bare bones of it is something you shouldn’t ignore on your next project.

The mechanics of Undo/Redo

It’s a rather simple thing to explain, because we’ve all used it before. The perspective of the implementation presented here is the following.

We begin with a document. The document is represented by a state. We can think of that state as a plain old object.

Document state changes as the user makes progress on their work.

In simple implementation, older states are just disposed/ignored. If you have undo functionality in mind, we can just capture those states so they can just be brought back. We could maybe store them in a stack, ready to be popped.

We can store the history of states in a stack

So what happens when the user clicks undo? We recover the state out of the undo stack. In order to avoid losing the current state we can introduce a new stack: the redo-stack, where we can push it.

First we push the current tate into the redo stack

Once the current state is all safe from oblivion we can pop out a state out of the undo stack and make it the current state.

Then we pop out the state out of the undo stack

Now let’s try this in a functioning scenario. We will create a Line editor where the user can draw lines.

Getting Ready

Things we use here:

• React
• Typescript — If you’re not using it, start right now.
• Next.js

Creating a Line Editor

Lets begin with a simple line editor, one simple use case:

To Draw Line: 

1. Mouse Down: Start line with point A at (x, y)
2. Mouse Move: Set line point B at (x, y)
3. Mouse Up: End line

Implementing this into a modest React component:

type Point = [number, number];
type Line = [Point, Point];
type Document = Line[];

export default function Home() {

  const [document, setDocument] = useState<Document>([]);
  const [isDrawing, setIsDrawing] = useState(false);

  const handleMouseDown = (e: React.MouseEvent) => {
    const point: Point = [e.clientX, e.clientY];
    setDocument([...document, [point, point]]);
    setIsDrawing(true);
  };

  const handleMouseUp = (e: React.MouseEvent) => {
    setIsDrawing(false);
  };

  const handleMouseMove = (e: React.MouseEvent) => {
    if(isDrawing && document.length > 0){
      const lastLine = document[document.length - 1];
      const lastLineUpdated: Line = [lastLine[0], [e.clientX, e.clientY]];
      setDocument([...document.slice(0, document.length - 1), lastLineUpdated]);
    }
  };

  return (
    <main>
      <svg 
        width="100vw" 
        height="100vh"
        onMouseDown={handleMouseDown} 
        onMouseMove={handleMouseMove}
        onMouseUp={handleMouseUp}
      >
        {document.map(([[x1, y1], [x2, y2]]) => 
          <line 
            key={JSON.stringify([x1, y1, x2, y2])}
            {...{x1, y1, x2, y2}}
          />
        )}
      </svg>
    </main>
  )
}

Ignore implementation details, code was kept short for simplicity and demonstration purposes. A few notes about the code:

• Convenience types define our Document.
• Document is just an array of Lines, Lines are a tuple of points, points are a tuple of numbers.
• isDrawing flag to understand if a line drawing is in process
• mouseDown creates a new line
• mouseMove moves the second point of the last line if isDrawing
• Rendering is just passing the coordinates to the <line> elements

Simple component, with a flag for understanding when a line is being drawn. The render method just renders lines with the collected coordinates.

And a few lines of CSS:

body {
  margin: 0;
  padding: 0;
}

line {
  stroke: #000;
  stroke-width: 1px;
}

We have an editor now:

Our Line Editor Live | Source

Implementing Undo

How the state changes? Imagine a scenario where the user draws three lines but regrets the last one, so undo is pressed after drawing the third one.

// State 0 - No lines
[]

// State 1 - One line
[a]

// State 2 - Two lines
[a, b]

// State 3 - Three lines
[a, b, c]

// Oops, regret that last one
undo();

// State 4 - Two lines (recovered from State 2)
[a, b]

So far so good. That last undo just means that we can reverse to the last state before that regretful decision.

From a structural point of view, this means we could just implement a stack of states, and when undo is invoked we just pop out the last state. We can actually repeatedly pop out states until the stack is empty.

We can write a rather simple toy to handle this:

class Editor<D>{

    private document: D;
    private undoStack: D[] = [];

    constructor(initialDocument: D) {
        this.document = initialDocument;
        this.save();
    }

    public getDocument(): D {
        return this.document;
    }

    public setDocument(newDocument: D){
        this.document = newDocument;
    }

    public store(): void {
        this.undoStack.push(this.document);
    }

    public undo(): void {
        const desired = this.undoStack.pop();

        if (desired) {
            this.document = desired;
        }
    }
}

This Editor class manages the state so we can undo changes to the state.

• The T type represents the type of document that this editor manages.
• The document local variable represents the current state in a staging phase.
• The save method pushes the document to the undo stack
• The undo method pops out the stage.

Let’s give our toy a test drive:

// Create editor with an empty state
const editor = new Editor<Document>([]);

// Add a line starting at (1, 1)
editor.setDocument([[1, 1], [1, 1]]);

...

// We wait for the user set the end point. Let's say at (2, 2)
editor.setDocument([[1, 1], [2, 2]]);

// Now we save the state so it can be "undone"
editor.store();

// This prints [[[1, 1], [2, 2]]] (array with one line)
console.log(editor.getDocument())

// We can undo it
editor.undo();

// Voila! This prints []
console.log(editor.getDocument())

Now, this is not how React works. React already provides state management and hooks as means to group, augment and and specialize that management.

We can translate that class to a React Hook:

function useUndo<T>(initialState: T): {
    document: T;
    setDocument(s: T): void;
    undo(): void;
    save(): void;
}{
    const [document, setDocument] = useState(initialState);
    const [undoStack, setUndoStack] = useState<T[]>([initialState]);

    const save = () => {
        setUndoStack([...undoStack, document]);
    };

    const undo = () => {
        if(undoStack.length > 1){
            const desired = undoStack[undoStack.length - 2];
            setUndoStack(undoStack.slice(0, undoStack.length - 1));
            setDocument(desired);
        }
    };

    return { document, setDocument, undo, save };
}

This hook does the same job as the class Editor we wrote earlier, but handles state the React way, exposing the methods as handy properties.

These methods can be harnessed in our editor with a couple of modifications:

type Point = [number, number];
type Line = [Point, Point];
type Document = Line[];

export default function EditorPage() {
  const [isDrawing, setIsDrawing] = useState(false);
  const { document, setDocument, undo, store } = useUndo<Document>([]);

  const handleMouseDown = (e: React.MouseEvent) => {
    const point: Point = [e.clientX, e.clientY];
    setDocument([...document, [point, point]]);
    setIsDrawing(true);
  };

  const handleMouseUp = (e: React.MouseEvent) => {
    setIsDrawing(false);
  };

  const handleMouseMove = (e: React.MouseEvent) => {
    if (isDrawing && document.length > 0) {
      const lastLine = document[document.length - 1];
      const lastLineUpdated: Line = [lastLine[0], [e.clientX, e.clientY]];
      setDocument([...document.slice(0, document.length - 1), lastLineUpdated]);
    }
  };

  useEffect(() => {
    if (!isDrawing) {
      store();
    }
  }, [isDrawing]);

  return (
    <main>
        <button onClick={undo}>Undo</button>
        <svg
            width="100vw"
            height="80vh"
            onMouseDown={handleMouseDown}
            onMouseMove={handleMouseMove}
            onMouseUp={handleMouseUp}
        >
            {document.map(([[x1, y1], [x2, y2]]) =>
                <line
                    key={JSON.stringify([x1, y1, x2, y2])}
                    {...{ x1, y1, x2, y2 }}
                />
            )}
        </svg>
    </main>
  )
}

• Replace useState for useUndo
• Added useEffect to save state when a line is finished
• Added button for undo operation

Undo working Live | Source

Now what about Redo?

Supplementing the hook for Redo functionality is just a matter of juggling the popped states to another redo stack:

function useUndoRedo<T>(initialState: T): {
    document: T;
    setDocument(s: T): void;
    undo(): void;
    redo(): void;
    store(): void;
}{
    const [document, setDocument] = useState(initialState);
    const [undoStack, setUndoStack] = useState<T[]>([initialState]);
    const [redoStack, setRedoStack] = useState<T[]>([]);

    const store = () => {
        setUndoStack([...undoStack, document]);
        setRedoStack([]);
    };

    const undo = () => {
        if(undoStack.length > 1){
            const last = undoStack[undoStack.length - 1];
            const desired = undoStack[undoStack.length - 2];
            setUndoStack(undoStack.slice(0, undoStack.length - 1));
            setRedoStack([...redoStack, last]);
            setDocument(desired);
        }
    };

    const redo = () => {
        if(redoStack.length > 0){
            const last = redoStack[redoStack.length - 1];
            setUndoStack([...undoStack, last]);
            setRedoStack(redoStack.slice(0, redoStack.length - 1));
            setDocument(last);
        }
    };

    return { document, setDocument, undo, redo, store };
}

And we add the button to the component:

type Point = [number, number];
type Line = [Point, Point];
type Document = Line[];

export default function EditorPage() {
  const [isDrawing, setIsDrawing] = useState(false);
  const { document, setDocument, undo, redo, store } = useUndoRedo<Document>([]);

  const handleMouseDown = (e: React.MouseEvent) => {
    const point: Point = [e.clientX, e.clientY];
    setDocument([...document, [point, point]]);
    setIsDrawing(true);
  };

  const handleMouseUp = (e: React.MouseEvent) => {
    setIsDrawing(false);
  };

  const handleMouseMove = (e: React.MouseEvent) => {
    if (isDrawing && document.length > 0) {
      const lastLine = document[document.length - 1];
      const lastLineUpdated: Line = [lastLine[0], [e.clientX, e.clientY]];
      setDocument([...document.slice(0, document.length - 1), lastLineUpdated]);
    }
  };

  useEffect(() => {
    if (!isDrawing) {
      store();
    }
  }, [isDrawing]);

  return (
    <main>
        <button onClick={undo}>Undo</button>
        <button onClick={redo}>Redo</button>
        <svg
            width="100vw"
            height="80vh"
            onMouseDown={handleMouseDown}
            onMouseMove={handleMouseMove}
            onMouseUp={handleMouseUp}
        >
            {document.map(([[x1, y1], [x2, y2]]) =>
                <line
                    key={JSON.stringify([x1, y1, x2, y2])}
                    {...{ x1, y1, x2, y2 }}
                />
            )}
        </svg>
    </main>
  )
}

Undo / Redo working Live | Source

The final version has a few extra touches to make it more obvious and usable.

Hope you have fun understanding and implementing.

• Check out the Live Version
• Don’t forget about the Source Code