perf: Table decorations (#10812)

* perf: CodeFence decos

* perf: Caching for table decos

* PR feedback

shared/editor/nodes/CodeFence.ts

@@ -7,7 +7,13 @@ import {
   Schema,
   Node as ProsemirrorNode,
 } from "prosemirror-model";
-import { Command, Plugin, PluginKey, TextSelection } from "prosemirror-state";
+import {
+  Command,
+  EditorState,
+  Plugin,
+  PluginKey,
+  TextSelection,
+} from "prosemirror-state";
 import { Decoration, DecorationSet } from "prosemirror-view";
 import { toast } from "sonner";
 import { Primitive } from "utility-types";
@@ -185,6 +191,19 @@ export default class CodeFence extends Node {
   }
 
   get plugins() {
+    const createActiveCodeBlockDecoration = (state: EditorState) => {
+      const codeBlock = findParentNode(isCode)(state.selection);
+      if (!codeBlock) {
+        return DecorationSet.empty;
+      }
+      const decoration = Decoration.node(
+        codeBlock.pos,
+        codeBlock.pos + codeBlock.node.nodeSize,
+        { class: "code-active" }
+      );
+      return DecorationSet.create(state.doc, [decoration]);
+    };
+
     return [
       CodeHighlighting({
         name: this.name,
@@ -243,20 +262,21 @@ export default class CodeFence extends Node {
         },
       }),
       new Plugin({
-        props: {
-          decorations(state) {
-            const codeBlock = findParentNode(isCode)(state.selection);
-
-            if (!codeBlock) {
-              return null;
+        key: new PluginKey("code-fence-active"),
+        state: {
+          init: (_, state) => createActiveCodeBlockDecoration(state),
+          apply: (tr, pluginState, oldState, newState) => {
+            // Only recompute if selection or document changed
+            if (!tr.selectionSet && !tr.docChanged) {
+              return pluginState;
             }
 
-            const decoration = Decoration.node(
-              codeBlock.pos,
-              codeBlock.pos + codeBlock.node.nodeSize,
-              { class: "code-active" }
-            );
-            return DecorationSet.create(state.doc, [decoration]);
+            return createActiveCodeBlockDecoration(newState);
+          },
+        },
+        props: {
+          decorations(state) {
+            return this.getState(state);
           },
         },
       }),

shared/editor/nodes/TableCell.ts

@@ -1,6 +1,6 @@
 import { Token } from "markdown-it";
 import { NodeSpec, Slice } from "prosemirror-model";
-import { Plugin } from "prosemirror-state";
+import { EditorState, Plugin, PluginKey } from "prosemirror-state";
 import { DecorationSet, Decoration } from "prosemirror-view";
 import { addRowBefore, selectRow, selectTable } from "../commands/table";
 import { getCellAttrs, setCellAttrs } from "../lib/table";
@@ -69,8 +69,88 @@ export default class TableCell extends Node {
       );
     }
 
+    const createRowDecorations = (state: EditorState) => {
+      if (!this.editor.view?.editable) {
+        return DecorationSet.empty;
+      }
+
+      const { doc } = state;
+      const decorations: Decoration[] = [];
+      const rows = getCellsInColumn(0)(state);
+
+      if (rows) {
+        rows.forEach((pos, visualIndex) => {
+          const actualRowIndex = getRowIndexInMap(visualIndex, state);
+          const index = actualRowIndex !== -1 ? actualRowIndex : visualIndex;
+
+          if (index === 0) {
+            const className = cn(EditorStyleHelper.tableGrip, {
+              selected: isTableSelected(state),
+            });
+
+            decorations.push(
+              Decoration.widget(
+                pos + 1,
+                () => {
+                  const grip = document.createElement("a");
+                  grip.role = "button";
+                  grip.className = className;
+                  return grip;
+                },
+                {
+                  key: className,
+                }
+              )
+            );
+          }
+
+          const className = cn(EditorStyleHelper.tableGripRow, {
+            selected: isRowSelected(index)(state) || isTableSelected(state),
+            first: index === 0,
+            last: visualIndex === rows.length - 1,
+          });
+
+          decorations.push(
+            Decoration.widget(
+              pos + 1,
+              () => {
+                const grip = document.createElement("a");
+                grip.role = "button";
+                grip.className = className;
+                grip.dataset.index = index.toString();
+                return grip;
+              },
+              {
+                key: cn(className, index),
+              }
+            )
+          );
+
+          if (index === 0) {
+            decorations.push(buildAddRowDecoration(pos, index));
+          }
+
+          decorations.push(buildAddRowDecoration(pos, index + 1));
+        });
+      }
+
+      return DecorationSet.create(doc, decorations);
+    };
+
     return [
       new Plugin({
+        key: new PluginKey("table-cell-decorations"),
+        state: {
+          init: (_, state) => createRowDecorations(state),
+          apply: (tr, pluginState, oldState, newState) => {
+            // Only recompute if selection or document changed
+            if (!tr.selectionSet && !tr.docChanged) {
+              return pluginState;
+            }
+
+            return createRowDecorations(newState);
+          },
+        },
         props: {
           transformCopied: (slice) => {
             // check if the copied selection is a single table, with a single row, with a single cell. If so,
@@ -146,73 +226,8 @@ export default class TableCell extends Node {
               return false;
             },
           },
-          decorations: (state) => {
-            if (!this.editor.view?.editable) {
-              return;
-            }
-
-            const { doc } = state;
-            const decorations: Decoration[] = [];
-            const rows = getCellsInColumn(0)(state);
-
-            if (rows) {
-              rows.forEach((pos, visualIndex) => {
-                const actualRowIndex = getRowIndexInMap(visualIndex, state);
-                const index =
-                  actualRowIndex !== -1 ? actualRowIndex : visualIndex;
-                if (index === 0) {
-                  const className = cn(EditorStyleHelper.tableGrip, {
-                    selected: isTableSelected(state),
-                  });
-
-                  decorations.push(
-                    Decoration.widget(
-                      pos + 1,
-                      () => {
-                        const grip = document.createElement("a");
-                        grip.role = "button";
-                        grip.className = className;
-                        return grip;
-                      },
-                      {
-                        key: className,
-                      }
-                    )
-                  );
-                }
-
-                const className = cn(EditorStyleHelper.tableGripRow, {
-                  selected:
-                    isRowSelected(index)(state) || isTableSelected(state),
-                  first: index === 0,
-                  last: visualIndex === rows.length - 1,
-                });
-
-                decorations.push(
-                  Decoration.widget(
-                    pos + 1,
-                    () => {
-                      const grip = document.createElement("a");
-                      grip.role = "button";
-                      grip.className = className;
-                      grip.dataset.index = index.toString();
-                      return grip;
-                    },
-                    {
-                      key: cn(className, index),
-                    }
-                  )
-                );
-
-                if (index === 0) {
-                  decorations.push(buildAddRowDecoration(pos, index));
-                }
-
-                decorations.push(buildAddRowDecoration(pos, index + 1));
-              });
-            }
-
-            return DecorationSet.create(doc, decorations);
+          decorations(state) {
+            return this.getState(state);
           },
         },
       }),

shared/editor/nodes/TableHeader.ts

@@ -1,6 +1,6 @@
 import { Token } from "markdown-it";
 import { NodeSpec } from "prosemirror-model";
-import { Plugin } from "prosemirror-state";
+import { EditorState, Plugin, PluginKey } from "prosemirror-state";
 import { DecorationSet, Decoration, EditorView } from "prosemirror-view";
 import { addColumnBefore, selectColumn } from "../commands/table";
 import { getCellAttrs, setCellAttrs } from "../lib/table";
@@ -68,8 +68,64 @@ export default class TableHeader extends Node {
       );
     }
 
+    const createColumnDecorations = (state: EditorState) => {
+      if (!this.editor.view?.editable) {
+        return DecorationSet.empty;
+      }
+
+      const { doc } = state;
+      const decorations: Decoration[] = [];
+      const cols = getCellsInRow(0)(state);
+
+      if (cols) {
+        cols.forEach((pos, index) => {
+          const className = cn(EditorStyleHelper.tableGripColumn, {
+            selected: isColumnSelected(index)(state) || isTableSelected(state),
+            first: index === 0,
+            last: index === cols.length - 1,
+          });
+
+          decorations.push(
+            Decoration.widget(
+              pos + 1,
+              () => {
+                const grip = document.createElement("a");
+                grip.role = "button";
+                grip.className = className;
+                grip.dataset.index = index.toString();
+                return grip;
+              },
+              {
+                key: cn(className, index),
+              }
+            )
+          );
+
+          if (index === 0) {
+            decorations.push(buildAddColumnDecoration(pos, index));
+          }
+
+          decorations.push(buildAddColumnDecoration(pos, index + 1));
+        });
+      }
+
+      return DecorationSet.create(doc, decorations);
+    };
+
     return [
       new Plugin({
+        key: new PluginKey("table-header-decorations"),
+        state: {
+          init: (_, state) => createColumnDecorations(state),
+          apply: (tr, pluginState, oldState, newState) => {
+            // Only recompute if selection or document changed
+            if (!tr.selectionSet && !tr.docChanged) {
+              return pluginState;
+            }
+
+            return createColumnDecorations(newState);
+          },
+        },
         props: {
           handleDOMEvents: {
             mousedown: (view: EditorView, event: MouseEvent) => {
@@ -107,49 +163,8 @@ export default class TableHeader extends Node {
               return false;
             },
           },
-          decorations: (state) => {
-            if (!this.editor.view?.editable) {
-              return;
-            }
-
-            const { doc } = state;
-            const decorations: Decoration[] = [];
-            const cols = getCellsInRow(0)(state);
-
-            if (cols) {
-              cols.forEach((pos, index) => {
-                const className = cn(EditorStyleHelper.tableGripColumn, {
-                  selected:
-                    isColumnSelected(index)(state) || isTableSelected(state),
-                  first: index === 0,
-                  last: index === cols.length - 1,
-                });
-
-                decorations.push(
-                  Decoration.widget(
-                    pos + 1,
-                    () => {
-                      const grip = document.createElement("a");
-                      grip.role = "button";
-                      grip.className = className;
-                      grip.dataset.index = index.toString();
-                      return grip;
-                    },
-                    {
-                      key: cn(className, index),
-                    }
-                  )
-                );
-
-                if (index === 0) {
-                  decorations.push(buildAddColumnDecoration(pos, index));
-                }
-
-                decorations.push(buildAddColumnDecoration(pos, index + 1));
-              });
-            }
-
-            return DecorationSet.create(doc, decorations);
+          decorations(state) {
+            return this.getState(state);
           },
         },
       }),