perf: support parsing and caching line by line based on shiki grammar state

Author: KazariEX

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.0.7 (2024-10-06)
+
+### Perf
+
+- support parsing and caching line by line based on shiki grammar state
+
 ## 0.0.6 (2024-10-06)
 
 ### Perf

eslint.config.js

@@ -16,7 +16,7 @@ export default antfu({
         ...zin.stylistic,
         ...zin.vue,
         ...zin.patch,
-        "no-return-assign": "off",
+        "no-multi-assign": "off",
         "unicorn/prefer-dom-node-text-content": "off"
     }
 });

src/index.ts

@@ -1,4 +1,4 @@
-import type { BundledLanguage, BundledTheme, CodeToTokensWithThemesOptions, HighlighterCore, ThemedTokenWithVariants } from "shiki";
+import type { BundledLanguage, BundledTheme, CodeToTokensWithThemesOptions, GrammarState, HighlighterCore, ThemedTokenWithVariants } from "shiki";
 import { debounce } from "./utils";
 
 export interface MountPlainShikiOptions {
@@ -40,11 +40,17 @@ export interface MountPlainShikiOptions {
 
 export type CreatePlainShikiReturns = ReturnType<typeof createPlainShiki>;
 
-interface ColorLoads {
+interface ColorLoad {
     token: ThemedTokenWithVariants;
     range: Range;
 }
 
+interface LoadLine {
+    text: string;
+    lastGrammarState: GrammarState;
+    loads: ColorLoad[];
+}
+
 export function createPlainShiki(shiki: HighlighterCore) {
     const isSupported = () => "CSS" in globalThis && "highlights" in CSS;
 
@@ -65,6 +71,7 @@ export function createPlainShiki(shiki: HighlighterCore) {
         const stylesheet = new CSSStyleSheet();
         document.adoptedStyleSheets.push(stylesheet);
 
+        const loadLines: LoadLine[] = [];
         const names = new Set<string>();
 
         if (isSupported()) {
@@ -79,16 +86,28 @@ export function createPlainShiki(shiki: HighlighterCore) {
             document.adoptedStyleSheets.splice(idx, 1);
         }
 
-        function patch(loads: ColorLoads[]) {
-            for (const name of names) {
-                const highlight = CSS.highlights.get(name);
-                highlight?.clear();
+        function resolveName(theme: string, color: string) {
+            return `shiki-${theme}-${color.slice(1).toLowerCase()}`;
+        }
+
+        function patch(loads: ColorLoad[], oldLoads: ColorLoad[]) {
+            for (const { token, range } of oldLoads) {
+                for (const theme in token.variants) {
+                    const { color } = token.variants[theme];
+                    const name = resolveName(theme, color!);
+
+                    const highlight = CSS.highlights.get(name);
+                    if (!highlight) {
+                        continue;
+                    }
+                    highlight.delete(range);
+                }
             }
 
             for (const { token, range } of loads) {
                 for (const theme in token.variants) {
                     const { color } = token.variants[theme];
-                    const name = `shiki-${theme}-${color!.slice(1).toLowerCase()}`;
+                    const name = resolveName(theme, color!);
                     const isDefault = theme === "light";
 
                     let highlight = CSS.highlights.get(name);
@@ -110,6 +129,21 @@ export function createPlainShiki(shiki: HighlighterCore) {
             }
         }
 
+        function diff(textLines: string[]) {
+            let i = 0;
+            for (; i < textLines.length; i++) {
+                if (textLines[i] !== loadLines[i]?.text) {
+                    return i;
+                }
+                for (const { range } of loadLines[i]?.loads ?? []) {
+                    if (range.collapsed) {
+                        return i;
+                    }
+                }
+            }
+            return i;
+        }
+
         function update() {
             const walker = document.createTreeWalker(el, NodeFilter.SHOW_TEXT);
 
@@ -120,33 +154,52 @@ export function createPlainShiki(shiki: HighlighterCore) {
             }
 
             const { innerText } = el;
-            const tokenLines = shiki.codeToTokensWithThemes(innerText, {
-                lang,
-                themes
-            });
+            const textLines = innerText.split("\n");
 
-            let i = 0;
+            let i = diff(textLines);
+            let lineOffset = textLines.slice(0, i).reduce((res, text) => res + text.length + 1, 0);
+
+            let j = 0;
             let offset = 0;
             let isCorrect = false;
+            findNodeAndOffset(lineOffset);
+
+            for (; i < textLines.length; i++) {
+                const text = textLines[i];
 
-            const loads: ColorLoads[] = [];
-            for (const tokens of tokenLines) {
-                for (const token of tokens) {
-                    const [startNode, startOffset] = findNodeAndOffset(token.offset);
-                    const [endNode, endOffset] = findNodeAndOffset(token.offset + token.content.length);
+                const tokenLines = shiki.codeToTokensWithThemes(text, {
+                    lang,
+                    themes
+                });
+
+                const loads: ColorLoad[] = [];
+                for (const token of tokenLines[0]) {
+                    const [startNode, startOffset] = findNodeAndOffset(lineOffset + token.offset);
+                    const [endNode, endOffset] = findNodeAndOffset(lineOffset + token.offset + token.content.length);
 
                     const range = document.createRange();
                     range.setStart(startNode, startOffset);
                     range.setEnd(endNode, endOffset);
-
                     loads.push({ token, range });
                 }
+
+                const loadLine = loadLines[i] ??= {} as LoadLine;
+                loadLine.text = text;
+
+                patch(loads, loadLine.loads ?? []);
+                loadLine.loads = loads;
+
+                const lastGrammarState = shiki.getLastGrammarState(text ?? "", { lang });
+                if (loadLine.lastGrammarState !== lastGrammarState) {
+                    loadLine.lastGrammarState = lastGrammarState;
+                    lineOffset += text.length + 1;
+                }
+                else break;
             }
-            patch(loads);
 
             function findNodeAndOffset(tokenOffset: number) {
-                for (; i < textNodes.length; i++) {
-                    const node = textNodes[i];
+                for (; j < textNodes.length; j++) {
+                    const node = textNodes[j];
                     const { textContent } = node;
                     if (!textContent) {
                         continue;
@@ -157,15 +210,13 @@ export function createPlainShiki(shiki: HighlighterCore) {
                         isCorrect = true;
                     }
 
-                    if (offset + textContent.length >= tokenOffset) {
-                        return [node, tokenOffset - offset] as const;
-                    }
-                    else {
+                    if (offset + textContent.length < tokenOffset) {
                         offset += textContent.length;
                         isCorrect = false;
                     }
+                    else break;
                 }
-                throw new RangeError(`[Plain Shiki] cannot find node at offset ${tokenOffset}.`);
+                return [textNodes[j], tokenOffset - offset] as const;
             }
         }