tiptap: fix select in codeBlock

packages/client/src/tiptap/core/extensions/code-block.ts

@@ -1,249 +1,157 @@
-import { mergeAttributes, Node, textblockTypeInputRule } from '@tiptap/core';
+import { findChildren } from '@tiptap/core';
+import BuiltInCodeBlock, { CodeBlockOptions } from '@tiptap/extension-code-block';
 import { ReactNodeViewRenderer } from '@tiptap/react';
 import { lowlight } from 'lowlight/lib/all';
-import { Plugin, PluginKey, TextSelection } from 'prosemirror-state';
+import { Node as ProsemirrorNode } from 'prosemirror-model';
+import { Plugin, PluginKey } from 'prosemirror-state';
+import { Decoration, DecorationSet } from 'prosemirror-view';
 import { CodeBlockWrapper } from 'tiptap/core/wrappers/code-block';
-import { LowlightPlugin } from 'tiptap/prose-utils';
 
-export interface CodeBlockOptions {
+function parseNodes(nodes: any[], className: string[] = []): { text: string; classes: string[] }[] {
-  /**
+  return nodes
-   * Adds a prefix to language classes that are applied to code tags.
+    .map((node) => {
-   * Defaults to `'language-'`.
+      const classes = [...className, ...(node.properties ? node.properties.className : [])];
-   */
+
-  languageClassPrefix: string;
+      if (node.children) {
-  /**
+        return parseNodes(node.children, classes);
-   * Define whether the node should be exited on triple enter.
-   * Defaults to `true`.
-   */
-  exitOnTripleEnter: boolean;
-  /**
-   * Define whether the node should be exited on arrow down if there is no node after it.
-   * Defaults to `true`.
-   */
-  exitOnArrowDown: boolean;
-  /**
-   * Custom HTML attributes that should be added to the rendered HTML tag.
-   */
-  HTMLAttributes: Record<string, any>;
       }
 
-declare module '@tiptap/core' {
-  interface Commands<ReturnType> {
-    codeBlock: {
-      /**
-       * Set a code block
-       */
-      setCodeBlock: (attributes?: { language: string }) => ReturnType;
-      /**
-       * Toggle a code block
-       */
-      toggleCodeBlock: (attributes?: { language: string }) => ReturnType;
-    };
-  }
-}
-
-export const backtickInputRegex = /^```([a-z]+)?[\s\n]$/;
-export const tildeInputRegex = /^~~~([a-z]+)?[\s\n]$/;
-
-export const BuiltInCodeBlock = Node.create<CodeBlockOptions>({
-  name: 'codeBlock',
-  draggable: true,
-
-  addOptions() {
       return {
-      languageClassPrefix: 'language-',
+        text: node.value,
-      exitOnTripleEnter: true,
+        classes,
-      exitOnArrowDown: true,
-      HTMLAttributes: {},
       };
-  },
-
-  content: 'text*',
-
-  marks: '',
-
-  group: 'block',
-
-  code: true,
-
-  defining: true,
-
-  addAttributes() {
-    return {
-      language: {
-        default: null,
-        parseHTML: (element) => {
-          const { languageClassPrefix } = this.options;
-          const classNames = Array.from(element.firstElementChild?.classList || element.classList || []);
-          const languages = classNames
-            .filter((className) => className.startsWith(languageClassPrefix))
-            .map((className) => className.replace(languageClassPrefix, ''));
-          const language = languages[0];
-
-          if (!language) {
-            return null;
-          }
-
-          return language;
-        },
-        rendered: false,
-      },
-    };
-  },
-
-  parseHTML() {
-    return [
-      {
-        tag: 'pre',
-        preserveWhitespace: 'full',
-      },
-    ];
-  },
-
-  renderHTML({ node, HTMLAttributes }) {
-    return [
-      'pre',
-      mergeAttributes(this.options.HTMLAttributes, HTMLAttributes),
-      [
-        'code',
-        {
-          class: node.attrs.language ? this.options.languageClassPrefix + node.attrs.language : null,
-        },
-        0,
-      ],
-    ];
-  },
-
-  addCommands() {
-    return {
-      setCodeBlock:
-        (attributes) =>
-        ({ commands }) => {
-          return commands.setNode(this.name, attributes);
-        },
-      toggleCodeBlock:
-        (attributes) =>
-        ({ commands }) => {
-          return commands.toggleNode(this.name, 'paragraph', attributes);
-        },
-    };
-  },
-
-  addKeyboardShortcuts() {
-    return {
-      'Mod-Alt-c': () => this.editor.commands.toggleCodeBlock(),
-
-      // remove code block when at start of document or code block is empty
-      'Backspace': () => {
-        const { empty, $anchor } = this.editor.state.selection;
-        const isAtStart = $anchor.pos === 1;
-
-        if (!empty || $anchor.parent.type.name !== this.name) {
-          return false;
-        }
-
-        if (isAtStart || !$anchor.parent.textContent.length) {
-          return this.editor.commands.clearNodes();
-        }
-
-        return false;
-      },
-
-      // exit node on triple enter
-      'Enter': ({ editor }) => {
-        if (!this.options.exitOnTripleEnter) {
-          return false;
-        }
-
-        const { state } = editor;
-        const { selection } = state;
-        const { $from, empty } = selection;
-
-        if (!empty || $from.parent.type !== this.type) {
-          return false;
-        }
-
-        const isAtEnd = $from.parentOffset === $from.parent.nodeSize - 2;
-        const endsWithDoubleNewline = $from.parent.textContent.endsWith('\n\n');
-
-        if (!isAtEnd || !endsWithDoubleNewline) {
-          return false;
-        }
-
-        return editor
-          .chain()
-          .command(({ tr }) => {
-            tr.delete($from.pos - 2, $from.pos);
-
-            return true;
     })
-          .exitCode()
+    .flat();
-          .run();
+}
+
+function getHighlightNodes(result: any) {
+  // `.value` for lowlight v1, `.children` for lowlight v2
+  return result.value || result.children || [];
+}
+
+function getDecorations({
+  doc,
+  name,
+  lowlight,
+  defaultLanguage,
+}: {
+  doc: ProsemirrorNode;
+  name: string;
+  lowlight: any;
+  defaultLanguage: string | null | undefined;
+}) {
+  const decorations: Decoration[] = [];
+
+  findChildren(doc, (node) => node.type.name === name).forEach((block) => {
+    let from = block.pos + 1;
+    const language = block.node.attrs.language || defaultLanguage;
+    const languages = lowlight.listLanguages();
+    const nodes =
+      language && languages.includes(language)
+        ? getHighlightNodes(lowlight.highlight(language, block.node.textContent))
+        : getHighlightNodes(lowlight.highlightAuto(block.node.textContent));
+
+    parseNodes(nodes).forEach((node) => {
+      const to = from + node.text.length;
+
+      if (node.classes.length) {
+        const decoration = Decoration.inline(from, to, {
+          class: node.classes.join(' '),
+        });
+
+        decorations.push(decoration);
+      }
+
+      from = to;
+    });
+  });
+
+  return DecorationSet.create(doc, decorations);
+}
+
+function isFunction(param: Function) {
+  return typeof param === 'function';
+}
+
+export function LowlightPlugin({
+  name,
+  lowlight,
+  defaultLanguage,
+}: {
+  name: string;
+  lowlight: any;
+  defaultLanguage: string | null | undefined;
+}) {
+  if (!['highlight', 'highlightAuto', 'listLanguages'].every((api) => isFunction(lowlight[api]))) {
+    throw Error('You should provide an instance of lowlight to use the code-block-lowlight extension');
+  }
+
+  const lowlightPlugin: Plugin<any> = new Plugin({
+    key: new PluginKey('lowlight'),
+
+    state: {
+      init: (_, { doc }) =>
+        getDecorations({
+          doc,
+          name,
+          lowlight,
+          defaultLanguage,
+        }),
+      apply: (transaction, decorationSet, oldState, newState) => {
+        const oldNodeName = oldState.selection.$head.parent.type.name;
+        const newNodeName = newState.selection.$head.parent.type.name;
+        const oldNodes = findChildren(oldState.doc, (node) => node.type.name === name);
+        const newNodes = findChildren(newState.doc, (node) => node.type.name === name);
+
+        if (
+          transaction.docChanged &&
+          // Apply decorations if:
+          // selection includes named node,
+          ([oldNodeName, newNodeName].includes(name) ||
+            // OR transaction adds/removes named node,
+            newNodes.length !== oldNodes.length ||
+            // OR transaction has changes that completely encapsulte a node
+            // (for example, a transaction that affects the entire document).
+            // Such transactions can happen during collab syncing via y-prosemirror, for example.
+            transaction.steps.some((step) => {
+              // @ts-ignore
+              return (
+                step.from !== undefined &&
+                // @ts-ignore
+                step.to !== undefined &&
+                oldNodes.some((node) => {
+                  // @ts-ignore
+                  return (
+                    node.pos >= step.from &&
+                    // @ts-ignore
+                    node.pos + node.node.nodeSize <= step.to
+                  );
+                })
+              );
+            }))
+        ) {
+          return getDecorations({
+            doc: transaction.doc,
+            name,
+            lowlight,
+            defaultLanguage,
+          });
+        }
+
+        return decorationSet.map(transaction.mapping, transaction.doc);
+      },
     },
 
-      // exit node on arrow down
+    props: {
-      'ArrowDown': ({ editor }) => {
+      decorations(state) {
-        if (!this.options.exitOnArrowDown) {
+        return lowlightPlugin.getState(state);
-          return false;
-        }
-
-        const { state } = editor;
-        const { selection, doc } = state;
-        const { $from, empty } = selection;
-
-        if (!empty || $from.parent.type !== this.type) {
-          return false;
-        }
-
-        const isAtEnd = $from.parentOffset === $from.parent.nodeSize - 2;
-
-        if (!isAtEnd) {
-          return false;
-        }
-
-        const after = $from.after();
-
-        if (after === undefined) {
-          return false;
-        }
-
-        const nodeAfter = doc.nodeAt(after);
-
-        if (nodeAfter) {
-          return false;
-        }
-
-        return editor.commands.exitCode();
       },
-
-      'Tab': ({ editor }) => {
-        const { selection } = this.editor.state;
-        const { $anchor } = selection;
-        return editor.chain().insertContentAt($anchor.pos, '  ').run();
-      },
-    };
-  },
-
-  addInputRules() {
-    return [
-      textblockTypeInputRule({
-        find: backtickInputRegex,
-        type: this.type,
-        getAttributes: (match) => ({
-          language: match[1],
-        }),
-      }),
-      textblockTypeInputRule({
-        find: tildeInputRegex,
-        type: this.type,
-        getAttributes: (match) => ({
-          language: match[1],
-        }),
-      }),
-    ];
     },
   });
 
+  return lowlightPlugin;
+}
+
 export interface CodeBlockLowlightOptions extends CodeBlockOptions {
   lowlight: any;
   defaultLanguage: string | null | undefined;
@@ -253,11 +161,15 @@ export const CodeBlock = BuiltInCodeBlock.extend<CodeBlockLowlightOptions>({
   addOptions() {
     return {
       ...this.parent?.(),
-      lowlight,
+      lowlight: {},
       defaultLanguage: null,
     };
   },
 
+  addNodeView() {
+    return ReactNodeViewRenderer(CodeBlockWrapper);
+  },
+
   addProseMirrorPlugins() {
     return [
       ...(this.parent?.() || []),
@@ -268,10 +180,6 @@ export const CodeBlock = BuiltInCodeBlock.extend<CodeBlockLowlightOptions>({
       }),
     ];
   },
-
-  addNodeView() {
-    return ReactNodeViewRenderer(CodeBlockWrapper);
-  },
 }).configure({
   lowlight,
   defaultLanguage: 'auto',