perf(core): token split performance improvements (#1301)

Author: amadeus

bench/decorations/single-file.bench.ts

@@ -0,0 +1,41 @@
+import type { DecorationItem } from 'shiki'
+import { createHighlighter, createOnigurumaEngine } from 'shiki'
+import { afterAll, bench, describe } from 'vitest'
+
+const lineCount = 5_000
+const code = Array.from(
+  { length: lineCount },
+  (_, line) => `const value${line.toString().padStart(5, '0')} = compute(${line}, 'before');`,
+).join('\n')
+const decorations: DecorationItem[] = Array.from(
+  { length: lineCount },
+  (_, line) => ({
+    start: { line, character: 8 },
+    end: { line, character: 13 },
+    properties: { class: 'highlighted' },
+  }),
+)
+
+const engine = await createOnigurumaEngine(() => import('shiki/wasm'))
+const highlighter = await createHighlighter({
+  engine,
+  langs: ['typescript'],
+  themes: ['vitesse-dark'],
+})
+
+afterAll(() => highlighter.dispose())
+
+describe(`single file (${lineCount} lines and decorations)`, () => {
+  bench('default decoration rendering', () => {
+    highlighter.codeToHast(code, {
+      decorations,
+      lang: 'typescript',
+      theme: 'vitesse-dark',
+    })
+  }, {
+    iterations: 3,
+    time: 0,
+    warmupIterations: 1,
+    warmupTime: 0,
+  })
+})

packages/core/src/utils/tokens.ts

@@ -39,6 +39,22 @@ export function splitToken<
   return tokens
 }
 
+// Find the first candidate so each token only scans its contained breakpoints.
+function findFirstBreakpointAfter(breakpoints: number[], offset: number): number {
+  let start = 0
+  let end = breakpoints.length
+
+  while (start < end) {
+    const middle = start + ((end - start) >> 1)
+    if (breakpoints[middle] <= offset)
+      start = middle + 1
+    else
+      end = middle
+  }
+
+  return start
+}
+
 /**
  * Split 2D tokens array by given breakpoints.
  */
@@ -55,14 +71,20 @@ export function splitTokens<
 
   return tokens.map((line) => {
     return line.flatMap((token) => {
-      const breakpointsInToken = sorted
-        .filter(i => token.offset < i && i < token.offset + token.content.length)
-        .map(i => i - token.offset)
-        .sort((a, b) => a - b)
+      const tokenEnd = token.offset + token.content.length
+      const start = findFirstBreakpointAfter(sorted, token.offset)
+      let end = start
+
+      while (end < sorted.length && sorted[end] < tokenEnd)
+        end++
 
-      if (!breakpointsInToken.length)
+      if (start === end)
         return token
 
+      const breakpointsInToken = sorted.slice(start, end)
+      for (let index = 0; index < breakpointsInToken.length; index++)
+        breakpointsInToken[index] -= token.offset
+
       return splitToken(token, breakpointsInToken)
     })
   })

packages/core/test/tokens.test.ts

@@ -1,6 +1,23 @@
 import { createJavaScriptRegexEngine } from 'shiki'
 import { describe, expect, it } from 'vitest'
-import { codeToHtml, codeToTokens, codeToTokensBase, createShikiPrimitiveAsync } from '../src'
+import { codeToHtml, codeToTokens, codeToTokensBase, createShikiPrimitiveAsync, splitTokens } from '../src'
+
+describe('splitTokens', () => {
+  it('splits at contained breakpoints while preserving token order and fields', () => {
+    const tokens = [[
+      { content: 'wxyz', marker: 'later', offset: 10 },
+      { content: 'abcd', marker: 'earlier', offset: 0 },
+    ]]
+
+    expect(splitTokens(tokens, [14, 2, 13, 11, 10, 4, 0, 2])).toEqual([[
+      { content: 'w', marker: 'later', offset: 10 },
+      { content: 'xy', marker: 'later', offset: 11 },
+      { content: 'z', marker: 'later', offset: 13 },
+      { content: 'ab', marker: 'earlier', offset: 0 },
+      { content: 'cd', marker: 'earlier', offset: 2 },
+    ]])
+  })
+})
 
 it('includeExplanation', async () => {
   using engine = await createShikiPrimitiveAsync({