Vasyl Putra

July 11, 2024 12:24 PM

Implementing a Undo/Redo Hook in React

Undo/Redo is a commonly used feature in many applications, enhancing user experience by allowing changes to be reverted. React Hooks provide a powerful way to manage state and handle changes in React components.

useUndoableState.tsx
1
import { SetStateAction, useCallback, useState } from "react";
2
import { unstable_batchedUpdates } from "react-dom";
3

4
interface IOptions {
5
  historyMaxLength?: number;
6
}
7

8
/**
9
 * useUndoableState is a custom hook that allows you to manage state with undo and redo functionality.
10
 * It takes an initial state and an options object as parameters.
11
 * The options object can have a historyMaxLength property which specifies the maximum number of states to store in the history.
12
 * The hook returns an object with the current state, a setState function to update the state,
13
 * a resetState function to reset the state to a new value,
14
 * undo and redo functions to move back and forth in the history,
15
 * and the undoHistory and redoHistory arrays that store the history of states.
16
 */
17
export const useUndoableState = <T>(
18
  initialState: T | (() => T),
19
  { historyMaxLength }: IOptions = {},
20
) => {
21
  // Define the state and its setter function
22
  const [state, setState] = useState<T>(initialState);
23

24
  // Define the undo and redo history arrays
25
  const [undoHistory, setUndoHistory] = useState<T[]>([]);
26
  const [redoHistory, setRedoHistory] = useState<T[]>([]);
27

28
  // Define the undo and redo functions
29
  const undo = useCallback(() => {
30
    // If there is anything in the undo history, restore the previous state
31
    if (undoHistory.length) {
32
      const stateToRestore = undoHistory.pop()!;
33

34
      // Update the undo and redo history arrays
35
      setUndoHistory([...undoHistory]);
36
      setRedoHistory([...redoHistory, state]);
37

38
      // Update the state
39
      setState(stateToRestore);
40
    }
41
  }, [undoHistory, redoHistory, state]);
42

43
  const redo = useCallback(() => {
44
    // If there is anything in the redo history, restore the next state
45
    if (redoHistory.length) {
46
      const stateToRestore = redoHistory.pop()!;
47

48
      // Update the undo and redo history arrays
49
      setUndoHistory([...undoHistory, state]);
50
      setRedoHistory([...redoHistory]);
51

52
      // Update the state
53
      setState(stateToRestore);
54
    }
55
  }, [state, redoHistory, undoHistory]);
56

57
  // Define the setState function that updates the state and adds the previous state to the undo history
58
  const handleSetState = useCallback(
59
    (stateOrStateSetter: SetStateAction<T>) => {
60
      // Use batchedUpdates to group multiple state updates into a single update
61
      unstable_batchedUpdates(() => {
62
        // Update the undo history by adding the current state to it
63
        setUndoHistory(
64
          [...undoHistory, state].slice(-Number(historyMaxLength)),
65
        );
66

67
        // Reset the redo history
68
        setRedoHistory([]);
69

70
        // Update the state
71
        setState(stateOrStateSetter);
72
      });
73
    },
74
    [historyMaxLength, state, undoHistory],
75
  );
76

77
  // Define the resetState function that resets the state to a new value
78
  const resetState = useCallback((stateOrStateSetter: SetStateAction<T>) => {
79
    // Use batchedUpdates to group multiple state updates into a single update
80
    unstable_batchedUpdates(() => {
81
      // Reset the undo and redo history
82
      setUndoHistory([]);
83
      setRedoHistory([]);
84

85
      // Update the state
86
      setState(stateOrStateSetter);
87
    });
88
  }, []);
89

90
  // Return the state, setState, resetState, undo, redo, undoHistory, and redoHistory objects
91
  return {
92
    state,
93
    setState: handleSetState,
94
    resetState,
95
    undo,
96
    undoHistory,
97
    redo,
98
    redoHistory,
99
  };
100
};