1
import { findChildren } from "@tiptap/core";
2
import type { Node as PMNode } from "@tiptap/pm/model";
3
import { Plugin, PluginKey, type EditorState, type Transaction } from "@tiptap/pm/state";
4
import { Decoration, DecorationSet } from "@tiptap/pm/view";
5
import type {
6
  CodeBlockHighlighter,
7
  HighlightNode,
8
  HighlightResultNode,
9
} from "@/shared/editor/highlight/code-highlight-runtime";
10
 
11
interface LazyHighlightMeta {
12
  refresh?: boolean;
13
}
14
 
15
const lazyHighlightPluginKey = new PluginKey<DecorationSet>("lazyCodeBlockHighlight");
16
const PLAIN_TEXT_LANGUAGES = new Set(["text", "plaintext", "txt"]);
17
let highlightRuntime: CodeBlockHighlighter | null = null;
18
let highlightRuntimePromise: Promise<CodeBlockHighlighter> | null = null;
19
 
20
function parseNodes(nodes: HighlightNode[], className: string[] = []): Array<{ text: string; classes: string[] }> {
21
  return nodes.flatMap((node) => {
22
    const classes = [...className, ...readClassNames(node.properties?.className)];
23
    if (Array.isArray(node.children) && node.children.length > 0) {
24
      return parseNodes(node.children, classes);
25
    }
26
    return node.value ? [{ text: node.value, classes }] : [];
27
  });
28
}
29
 
30
function readClassNames(className: string[] | string | undefined): string[] {
31
  if (Array.isArray(className)) return className;
32
  return className ? [className] : [];
33
}
34
 
35
function getHighlightNodes(result: HighlightResultNode): HighlightNode[] {
36
  return result.value ?? result.children ?? [];
37
}
38
 
39
function isPlainTextLanguage(language: string | null | undefined): boolean {
40
  return language ? PLAIN_TEXT_LANGUAGES.has(language.trim().toLowerCase()) : false;
41
}
42
 
43
function getCodeBlockLanguage(
44
  language: string | null | undefined,
45
  defaultLanguage: string | null | undefined,
46
): string | null | undefined {
47
  if (language && language.trim() !== "") {
48
    return language.trim();
49
  }
50
 
51
  return defaultLanguage?.trim() || undefined;
52
}
53
 
54
function shouldLoadHighlightRuntime(
55
  doc: EditorState["doc"],
56
  name: string,
57
  defaultLanguage: string | null | undefined,
58
): boolean {
59
  return findChildren(doc, (node) => node.type.name === name).some(({ node }) => {
60
    const language = getCodeBlockLanguage(
61
      typeof node.attrs.language === "string" ? node.attrs.language : null,
62
      defaultLanguage,
63
    );
64
    return language == null || !isPlainTextLanguage(language);
65
  });
66
}
67
 
68
export function getCodeBlockSnapshot(doc: PMNode, name: string, defaultLanguage: string | null | undefined): string[] {
69
  return findChildren(doc, (node) => node.type.name === name).map(({ node }) => {
70
    const language = getCodeBlockLanguage(
71
      typeof node.attrs.language === "string" ? node.attrs.language : null,
72
      defaultLanguage,
73
    );
74
    return `${language ?? ""}\u0000${node.textContent}`;
75
  });
76
}
77
 
78
export function codeBlockSnapshotChanged(
79
  oldDoc: PMNode,
80
  newDoc: PMNode,
81
  name: string,
82
  defaultLanguage: string | null | undefined,
83
): boolean {
84
  const oldSnapshot = getCodeBlockSnapshot(oldDoc, name, defaultLanguage);
85
  const newSnapshot = getCodeBlockSnapshot(newDoc, name, defaultLanguage);
86
 
87
  if (oldSnapshot.length !== newSnapshot.length) {
88
    return true;
89
  }
90
 
91
  return oldSnapshot.some((entry, index) => entry !== newSnapshot[index]);
92
}
93
 
94
function transactionTouchesCodeBlockRanges(transaction: Transaction, doc: PMNode, name: string): boolean {
95
  const codeBlocks = findChildren(doc, (node) => node.type.name === name);
96
 
97
  return transaction.steps.some((step) => {
98
    const rangedStep = step as { from?: number; to?: number };
99
    if (typeof rangedStep.from !== "number" || typeof rangedStep.to !== "number") {
100
      return false;
101
    }
102
 
103
    return codeBlocks.some(
104
      (block) => block.pos >= rangedStep.from! && block.pos + block.node.nodeSize <= rangedStep.to!,
105
    );
106
  });
107
}
108
 
109
function getDecorations(opts: {
110
  doc: EditorState["doc"];
111
  name: string;
112
  lowlight: CodeBlockHighlighter;
113
  defaultLanguage: string | null | undefined;
114
}): DecorationSet {
115
  const decorations: Decoration[] = [];
116
 
117
  findChildren(opts.doc, (node) => node.type.name === opts.name).forEach((block) => {
118
    let from = block.pos + 1;
119
    const language = getCodeBlockLanguage(
120
      typeof block.node.attrs.language === "string" ? block.node.attrs.language : null,
121
      opts.defaultLanguage,
122
    );
123
    const highlightNodes =
124
      language && !isPlainTextLanguage(language)
125
        ? (() => {
126
            const languages = opts.lowlight.listLanguages();
127
            if (languages.includes(language) || opts.lowlight.registered?.(language)) {
128
              return getHighlightNodes(opts.lowlight.highlight(language, block.node.textContent));
129
            }
130
            return getHighlightNodes(opts.lowlight.highlightAuto(block.node.textContent));
131
          })()
132
        : [];
133
 
134
    parseNodes(highlightNodes).forEach((node) => {
135
      const to = from + node.text.length;
136
      if (node.classes.length > 0) {
137
        decorations.push(Decoration.inline(from, to, { class: node.classes.join(" ") }));
138
      }
139
      from = to;
140
    });
141
  });
142
 
143
  return DecorationSet.create(opts.doc, decorations);
144
}
145
 
146
function shouldRefreshDecorations(
147
  transaction: Transaction,
148
  oldState: EditorState,
149
  newState: EditorState,
150
  name: string,
151
  defaultLanguage: string | null | undefined,
152
): boolean {
153
  return (
154
    transaction.docChanged &&
155
    (codeBlockSnapshotChanged(oldState.doc, newState.doc, name, defaultLanguage) ||
156
      transactionTouchesCodeBlockRanges(transaction, oldState.doc, name))
157
  );
158
}
159
 
160
function loadHighlightRuntime(): Promise<CodeBlockHighlighter> {
161
  if (highlightRuntime) {
162
    return Promise.resolve(highlightRuntime);
163
  }
164
 
165
  if (!highlightRuntimePromise) {
166
    highlightRuntimePromise = import("@/shared/editor/highlight/code-highlight-runtime").then((mod) => {
167
      highlightRuntime = mod.codeBlockLowlight;
168
      return highlightRuntime;
169
    });
170
  }
171
 
172
  return highlightRuntimePromise!;
173
}
174
 
175
export function createLazyHighlightPlugin(name: string, defaultLanguage: string | null | undefined) {
176
  return new Plugin<DecorationSet>({
177
    key: lazyHighlightPluginKey,
178
    state: {
179
      init: (_, state) => {
180
        return highlightRuntime
181
          ? getDecorations({ doc: state.doc, name, lowlight: highlightRuntime, defaultLanguage })
182
          : DecorationSet.empty;
183
      },
184
      apply: (transaction, decorationSet, oldState, newState) => {
185
        const meta = transaction.getMeta(lazyHighlightPluginKey) as LazyHighlightMeta | undefined;
186
        if (meta?.refresh) {
187
          return highlightRuntime
188
            ? getDecorations({ doc: newState.doc, name, lowlight: highlightRuntime, defaultLanguage })
189
            : DecorationSet.empty;
190
        }
191
 
192
        if (highlightRuntime && shouldRefreshDecorations(transaction, oldState, newState, name, defaultLanguage)) {
193
          return getDecorations({ doc: transaction.doc, name, lowlight: highlightRuntime, defaultLanguage });
194
        }
195
 
196
        return decorationSet.map(transaction.mapping, transaction.doc);
197
      },
198
    },
199
    props: {
200
      decorations(state) {
201
        return lazyHighlightPluginKey.getState(state) ?? DecorationSet.empty;
202
      },
203
    },
204
    view(view) {
205
      let active = true;
206
 
207
      const maybeLoad = () => {
208
        if (highlightRuntime || !shouldLoadHighlightRuntime(view.state.doc, name, defaultLanguage)) {
209
          return;
210
        }
211
 
212
        void loadHighlightRuntime()
213
          .then(() => {
214
            if (!active) return;
215
            view.dispatch(
216
              view.state.tr.setMeta(lazyHighlightPluginKey, { refresh: true }).setMeta("addToHistory", false),
217
            );
218
          })
219
          .catch(() => {
220
            // Keep code blocks readable even if the highlight runtime fails to load.
221
          });
222
      };
223
 
224
      maybeLoad();
225
 
226
      return {
227
        update(nextView, previousState) {
228
          view = nextView;
229
          if (previousState.doc !== nextView.state.doc || previousState.selection !== nextView.state.selection) {
230
            maybeLoad();
231
          }
232
        },
233
        destroy() {
234
          active = false;
235
        },
236
      };
237
    },
238
  });
239
}