feat: code highlight for tiptap

composables/shiki.ts

@@ -1,11 +1,11 @@
 import type { Highlighter, Lang } from 'shiki-es'
 
-export const shiki = ref<Highlighter>()
+const shiki = ref<Highlighter>()
 
 const registeredLang = ref(new Map<string, boolean>())
 let shikiImport: Promise<void> | undefined
 
-export function highlightCode(code: string, lang: Lang) {
+export function useHightlighter(lang: Lang) {
   if (!shikiImport) {
     shikiImport = import('shiki-es')
       .then(async (r) => {
@@ -25,7 +25,7 @@ export function highlightCode(code: string, lang: Lang) {
   }
 
   if (!shiki.value)
-    return code
+    return undefined
 
   if (!registeredLang.value.get(lang)) {
     shiki.value.loadLanguage(lang)
@@ -37,11 +37,27 @@ export function highlightCode(code: string, lang: Lang) {
         console.error(e)
         registeredLang.value.set(lang, false)
       })
-    return code
+    return undefined
   }
 
-  return shiki.value.codeToHtml(code, {
+  return shiki.value
+}
+
+export function useShikiTheme() {
+  return isDark.value ? 'vitesse-dark' : 'vitesse-light'
+}
+
+export function highlightCode(code: string, lang: Lang) {
+  const shiki = useHightlighter(lang)
+  if (!shiki)
+    return code
+
+  return shiki.codeToHtml(code, {
     lang,
-    theme: isDark.value ? 'vitesse-dark' : 'vitesse-light',
+    theme: useShikiTheme(),
   })
 }
+
+export function useShiki() {
+  return shiki
+}

composables/tiptap.ts

@@ -1,4 +1,4 @@
-import { Extension, VueNodeViewRenderer, useEditor } from '@tiptap/vue-3'
+import { Extension, useEditor } from '@tiptap/vue-3'
 import Placeholder from '@tiptap/extension-placeholder'
 import Document from '@tiptap/extension-document'
 import Paragraph from '@tiptap/extension-paragraph'
@@ -10,11 +10,10 @@ import Bold from '@tiptap/extension-bold'
 import Italic from '@tiptap/extension-italic'
 import Code from '@tiptap/extension-code'
 import { Plugin } from 'prosemirror-state'
-import CodeBlock from '@tiptap/extension-code-block'
 
 import type { Ref } from 'vue'
 import { HashSuggestion, MentionSuggestion } from './tiptap/suggestion'
-import TiptapCodeBlock from '~/components/tiptap/TiptapCodeBlock.vue'
+import { CodeBlockShiki } from './tiptap/shiki'
 
 export interface UseTiptapOptions {
   content: Ref<string | undefined>
@@ -56,12 +55,7 @@ export function useTiptap(options: UseTiptapOptions) {
       CharacterCount.configure({
         limit: characterLimit.value,
       }),
-      CodeBlock
-        .extend({
-          addNodeView() {
-            return VueNodeViewRenderer(TiptapCodeBlock)
-          },
-        }),
+      CodeBlockShiki,
       Extension.create({
         name: 'api',
         addKeyboardShortcuts() {

composables/tiptap/shiki.ts

@@ -0,0 +1,129 @@
+import type { CodeBlockOptions } from '@tiptap/extension-code-block'
+import CodeBlock from '@tiptap/extension-code-block'
+import { VueNodeViewRenderer } from '@tiptap/vue-3'
+
+import { findChildren } from '@tiptap/core'
+import type { Node as ProsemirrorNode } from 'prosemirror-model'
+import { Plugin, PluginKey } from 'prosemirror-state'
+import { Decoration, DecorationSet } from 'prosemirror-view'
+import TiptapCodeBlock from '~/components/tiptap/TiptapCodeBlock.vue'
+
+export interface CodeBlockShikiOptions extends CodeBlockOptions {
+  defaultLanguage: string | null | undefined
+}
+
+export const CodeBlockShiki = CodeBlock.extend<CodeBlockShikiOptions>({
+  addOptions() {
+    return {
+      ...this.parent?.(),
+      defaultLanguage: null,
+    }
+  },
+
+  addProseMirrorPlugins() {
+    return [
+      ...this.parent?.() || [],
+      ProseMirrorShikiPlugin({
+        name: this.name,
+      }),
+    ]
+  },
+
+  addNodeView() {
+    return VueNodeViewRenderer(TiptapCodeBlock)
+  },
+})
+
+function getDecorations({
+  doc,
+  name,
+}: { doc: ProsemirrorNode; name: string }) {
+  const decorations: Decoration[] = []
+
+  findChildren(doc, node => node.type.name === name)
+    .forEach((block) => {
+      let from = block.pos + 1
+      const language = block.node.attrs.language || 'text'
+
+      const shiki = useHightlighter(language)
+
+      if (!shiki)
+        return
+
+      const lines = shiki.codeToThemedTokens(block.node.textContent, language, useShikiTheme())
+
+      lines.forEach((line) => {
+        line.forEach((token) => {
+          const decoration = Decoration.inline(from, from + token.content.length, {
+            style: `color: ${token.color}`,
+          })
+
+          decorations.push(decoration)
+          from += token.content.length
+        })
+        from += 1
+      })
+    })
+
+  return DecorationSet.create(doc, decorations)
+}
+
+function ProseMirrorShikiPlugin({ name }: { name: string }) {
+  const plugin: Plugin<any> = new Plugin({
+    key: new PluginKey('shiki'),
+
+    state: {
+      init: (_, { doc }) => getDecorations({
+        doc,
+        name,
+      }),
+      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-expect-error cast
+              return step.from !== undefined
+                // @ts-expect-error cast
+                && step.to !== undefined
+                && oldNodes.some((node) => {
+                  // @ts-expect-error cast
+                  return node.pos >= step.from
+                    // @ts-expect-error cast
+                    && node.pos + node.node.nodeSize <= step.to
+                })
+            })
+          )
+        ) {
+          return getDecorations({
+            doc: transaction.doc,
+            name,
+          })
+        }
+
+        return decorationSet.map(transaction.mapping, transaction.doc)
+      },
+    },
+
+    props: {
+      decorations(state) {
+        return plugin.getState(state)
+      },
+    },
+  })
+
+  return plugin
+}