How to optimize custom Tiptap extension (React)

1

I noticed huge performance depreciation when started to use below plugin. I’m pretty new in React and maybe you have some suggestions what could I change to “fix” it? Its an extension that adds ability to add annotations to the editor.

import {
  generateHTML,
  getNodeType,
  mergeAttributes,
  Node,
  NodeWithPos,
} from '@tiptap/core';
import { DOMOutputSpec, NodeType, Node as PMNode } from '@tiptap/pm/model';
import { Editor } from '@tiptap/react';
import { Plugin, PluginKey, Transaction } from '@tiptap/pm/state';
import { getAllNodesOfType } from '@/components/editor/utils/utils';

declare module '@tiptap/core' {
  interface Commands<ReturnType> {
    annotation: {
      setAnnotation: (
        content?: PMNode,
        type?: 'annotation' | 'remark'
      ) => ReturnType;
      deleteAnnotation: () => ReturnType;
    };
  }
}

const ALPHABET = 'abcdefghijklmnopqrstuvwxyz';

export const getAnnotationNumber = (
  editor: Editor,
  node: PMNode,
  nodes?: PMNode[]
) => {
  if (!nodes) {
    let annotations = getAllNodesOfType(editor, 'annotation');
    nodes = [...annotations.map((el) => el.node)];
  }

  nodes = nodes.filter((x) => x.attrs['type'] === node.attrs['type']);
  let number = nodes.indexOf(node);
  if (node.attrs['type'] === 'annotation') return number + 1;
  return ALPHABET[number];
};

const updateAnnotations = (state: {
  editor: Editor;
  storage: AnnotationStorage;
  onUpdate: any;
}) => {
  const { editor, onUpdate } = state;
  let data: AnnotationData = [];

  // Find all annotation nodes and its positions within the doc
  const nodesWithPos = getAllNodesOfType(editor, 'annotation');

  // Convert annotation nodes to annotation data
  nodesWithPos.forEach((nodeWithPos, index) => {
    const element = editor.view.domAtPos(nodeWithPos.pos + 1)
      .node as HTMLElement;

    data.push({
      label: nodeWithPos.node.attrs['label'],
      id: nodeWithPos.node.attrs['id'],
      dom: element,
      type: nodeWithPos.node.attrs['type'],
    });
  });

  if (onUpdate) {
    const isCreate = 0 === state.storage.content.length;
    onUpdate(data, isCreate);
  }

  state.storage.content = data;
  // editor.state.tr.setMeta('annotation', data);
  // editor.view.dispatch(editor.state.tr);
};

export type AnnotationData = Array<AnnotationDataItem>;

export type AnnotationDataItem = {
  dom: HTMLElement;
  id: string;
  label: string | number;
  type: string;
};

export type AnnotationStorage = {
  content: AnnotationData;
  elements: Array<HTMLElement>;
};

export interface AnnotationOptions {
  HTMLAttributes: Record<string, any>;
  onUpdate?: (data: AnnotationData, isCreate?: boolean) => void;
}

export const Annotation = Node.create<AnnotationOptions, AnnotationStorage>({
  name: 'annotation',
  group: 'inline',
  content: 'block*',
  inline: true,
  selectable: true,
  atom: true,

  addStorage: () => ({
    content: [],
    elements: [],
  }),

  addAttributes() {
    return {
      id: {
        default: null,
        renderHTML: (attributes) => ({
          id: attributes['id'],
        }),
        parseHTML: (element: HTMLElement) => element.getAttribute('id'),
      },
      label: {
        default: null,
      },
      type: {
        default: 'annotation',
      },
    };
  },

  parseHTML() {
    return [
      {
        tag: 'span[data-type="annotation"]',
      },
    ];
  },

  renderHTML({ HTMLAttributes, node }) {
    let attrs = mergeAttributes(this.options.HTMLAttributes, HTMLAttributes, {
      'data-type': 'annotation',
      class: 'inline-block',
      contentEditable: false,
    });

    let content: DOMOutputSpec[];

    content = [
      // Label (number of annotation)
      [
        'span',
        { 'data-type': 'annotation-label' },
        `[${getAnnotationNumber(this.editor as Editor, node)}]`,
      ],
      // Content (any `block`)
      ['div', { 'data-type': 'annotation-content', class: 'hidden' }, 0],
    ];

    return ['div', attrs, ...content];
  },

  addCommands() {
    return {
      setAnnotation:
        (content, type) =>
        ({ tr, editor, chain }) => {
          return chain()
            .command(({ tr }) => {
              tr.replaceSelectionWith(
                this.type.create({ type }, content ? content : undefined)
              );
              return true;
            })
            .setMeta('refresh-annotation', true)
            .run();
        },
      deleteAnnotation:
        () =>
        ({ tr, editor, chain }) => {
          return (
            chain()
              // .command(({ tr }) => {
              //   tr.replaceSelectionWith(this.type.create());
              //   return true;
              // })
              .deleteNode('annotation')
              .setMeta('refresh-annotation', true)
              .run()
          );
        },
    };
  },

  onUpdate() {
    updateAnnotations({
      editor: this.editor as Editor,
      storage: this.storage,
      onUpdate:
        this.options.onUpdate === null || this.options.onUpdate === void 0
          ? void 0
          : this.options.onUpdate.bind(this),
    });
  },

  addProseMirrorPlugins() {
    let editor = this.editor;

    let plugin = new Plugin({
      key: new PluginKey('annotation'),
      appendTransaction(transactions, oldState, newState) {
        let { tr } = newState;
        if (transactions.some((tr) => tr.docChanged)) {
          let annotations = getAllNodesOfType(
            editor as Editor,
            'annotation',
            tr
          );

          let nodes = annotations.map((el) => el.node);

          annotations.forEach((item) => {
            tr.setNodeMarkup(item.pos, void 0, {
              ...item.node.attrs,
              label: getAnnotationNumber(editor as Editor, item.node, nodes),
            });
          });
        }

        return tr;
      },
    });

    return [plugin];
  },
});

UPDATE It seems that onUpdate method is the reason. It essentailly calls setAnnotations in the top-level component (there were useEditor is called as well).