fix(stream): align ShikiStreamTokenizer offsets and CRLF handling with codeToTokens (#1335)

Author: xyrolle

packages/stream/src/tokenizer.ts

@@ -1,5 +1,6 @@
 import type { GrammarState, ThemedToken } from '@shikijs/core'
 import type { ShikiStreamTokenizerEnqueueResult, ShikiStreamTokenizerOptions } from './types'
+import { splitLines } from '@shikijs/core'
 
 export class ShikiStreamTokenizer {
   public readonly options: ShikiStreamTokenizerOptions
@@ -8,6 +9,10 @@ export class ShikiStreamTokenizer {
   public tokensUnstable: ThemedToken[] = []
 
   public lastUnstableCodeChunk: string = ''
+  /**
+   * Offset of `lastUnstableCodeChunk` in the whole streamed code
+   */
+  public lastUnstableCodeOffset: number = 0
   public lastStableGrammarState: GrammarState | undefined
 
   constructor(
@@ -20,22 +25,26 @@ export class ShikiStreamTokenizer {
    * Enqueue a chunk of code to the buffer.
    */
   async enqueue(chunk: string): Promise<ShikiStreamTokenizerEnqueueResult> {
-    const chunkLines = (this.lastUnstableCodeChunk + chunk).split('\n')
+    const offset = this.lastUnstableCodeOffset
+    const chunkLines = splitLines(this.lastUnstableCodeChunk + chunk)
 
     const stable: ThemedToken[] = []
     let unstable: ThemedToken[] = []
     const recall = this.tokensUnstable.length
 
-    chunkLines.forEach((line, i) => {
+    chunkLines.forEach(([line, lineOffset], i) => {
       const isLastLine = i === chunkLines.length - 1
 
       const result = this.options.highlighter.codeToTokens(line, {
         ...this.options,
         grammarState: this.lastStableGrammarState,
       })
       const tokens = result.tokens[0] // only one line
+      for (const token of tokens)
+        token.offset += offset + lineOffset
+      // the `\n` is the last character before the next line, after any `\r`
       if (!isLastLine)
-        tokens.push({ content: '\n', offset: 0 })
+        tokens.push({ content: '\n', offset: offset + chunkLines[i + 1][1] - 1 })
 
       if (!isLastLine) {
         this.lastStableGrammarState = result.grammarState
@@ -44,6 +53,7 @@ export class ShikiStreamTokenizer {
       else {
         unstable = tokens
         this.lastUnstableCodeChunk = line
+        this.lastUnstableCodeOffset = offset + lineOffset
       }
     })
 
@@ -61,6 +71,7 @@ export class ShikiStreamTokenizer {
     const stable = this.tokensUnstable
     this.tokensUnstable = []
     this.lastUnstableCodeChunk = ''
+    this.lastUnstableCodeOffset = 0
     this.lastStableGrammarState = undefined
     return {
       stable,
@@ -71,6 +82,7 @@ export class ShikiStreamTokenizer {
     this.tokensStable = []
     this.tokensUnstable = []
     this.lastUnstableCodeChunk = ''
+    this.lastUnstableCodeOffset = 0
     this.lastStableGrammarState = undefined
   }
 
@@ -79,6 +91,7 @@ export class ShikiStreamTokenizer {
       this.options,
     )
     clone.lastUnstableCodeChunk = this.lastUnstableCodeChunk
+    clone.lastUnstableCodeOffset = this.lastUnstableCodeOffset
     clone.tokensUnstable = this.tokensUnstable
     clone.tokensStable = this.tokensStable
     clone.lastStableGrammarState = this.lastStableGrammarState

packages/stream/test/tokenizer.test.ts

@@ -1,3 +1,4 @@
+import type { ThemedToken } from 'shiki'
 import type { ShikiStreamTokenizerOptions } from '../src'
 import { createHighlighter } from 'shiki'
 import { expect, it } from 'vitest'
@@ -64,3 +65,113 @@ it('exported', async () => {
 
   expect.soft(merged).toBe(onepass)
 })
+
+const tokensOptions = {
+  lang: 'vue',
+  themes: {
+    dark: 'vitesse-dark',
+    light: 'vitesse-light',
+  },
+} as const
+
+async function tokenizeInChunks(options: ShikiStreamTokenizerOptions, code: string, size: number): Promise<ThemedToken[]> {
+  const streamer = new ShikiStreamTokenizer(options)
+  for (let i = 0; i < code.length; i += size)
+    await streamer.enqueue(code.slice(i, i + size))
+  return [...streamer.tokensStable, ...streamer.close().stable]
+}
+
+// The streamed tokens, apart from the inserted `\n` ones, match the one-pass tokens,
+// every token's offset points at its content in the code, and each inserted `\n` sits at its own LF
+function expectSameTokens(options: ShikiStreamTokenizerOptions, code: string, tokens: ThemedToken[]): void {
+  expect(tokens.filter(token => token.content !== '\n')).toEqual(options.highlighter.codeToTokens(code, options).tokens.flat())
+  for (const token of tokens)
+    expect(code.slice(token.offset, token.offset + token.content.length)).toBe(token.content)
+  expect(tokens.filter(token => token.content === '\n').map(token => token.offset))
+    .toEqual([...code.matchAll(/\n/g)].map(match => match.index))
+}
+
+it('handles CRLF line endings like codeToTokens', async () => {
+  const highlighter = await createHighlighter({
+    langs: ['vue'],
+    themes: ['vitesse-dark', 'vitesse-light'],
+  })
+  const code = fixture.trimEnd().replaceAll('\n', '\r\n')
+  const onepass = tokensToHtml(
+    highlighter.codeToTokens(code, tokensOptions)
+      .tokens
+      .flatMap((i, idx) => idx === 0 ? i : [{ content: '\n', offset: 0 }, ...i]),
+  )
+
+  // A chunk size of 1 splits every `\r\n` across two chunks
+  for (const size of [1, 5, code.length]) {
+    const tokens = await tokenizeInChunks({ highlighter, ...tokensOptions }, code, size)
+    expect(tokens.filter(token => token.content.includes('\r'))).toEqual([])
+    expect(tokensToHtml(tokens)).toBe(onepass)
+  }
+})
+
+it('reports token offsets relative to the whole code', async () => {
+  const highlighter = await createHighlighter({
+    langs: ['vue'],
+    themes: ['vitesse-dark', 'vitesse-light'],
+  })
+
+  const options: ShikiStreamTokenizerOptions = { highlighter, ...tokensOptions }
+
+  for (const code of [fixture.trimEnd(), fixture.trimEnd().replaceAll('\n', '\r\n')]) {
+    for (const size of [1, 5, code.length])
+      expectSameTokens(options, code, await tokenizeInChunks(options, code, size))
+  }
+})
+
+it('reports whole-code offsets after every chunk', async () => {
+  const highlighter = await createHighlighter({
+    langs: ['ts'],
+    themes: ['vitesse-light'],
+  })
+  const options: ShikiStreamTokenizerOptions = { highlighter, lang: 'ts', theme: 'vitesse-light' }
+  const streamer = new ShikiStreamTokenizer(options)
+
+  // `\r`, an empty chunk, `\n` and the two halves of the emoji arrive separately
+  const chunks = ['const a = 1', '\r', '', '\n', 'const b = \'\uD83D', '\uDE00\'\r\n', 'b']
+  let code = ''
+  const received: ThemedToken[] = []
+  for (const chunk of chunks) {
+    code += chunk
+    const { stable, unstable } = await streamer.enqueue(chunk)
+    received.push(...stable)
+    expectSameTokens(options, code, [...received, ...unstable])
+  }
+  expectSameTokens(options, code, [...received, ...streamer.close().stable])
+})
+
+it('continues offsets in a clone', async () => {
+  const highlighter = await createHighlighter({
+    langs: ['ts'],
+    themes: ['vitesse-light'],
+  })
+  const options: ShikiStreamTokenizerOptions = { highlighter, lang: 'ts', theme: 'vitesse-light' }
+  const streamer = new ShikiStreamTokenizer(options)
+  await streamer.enqueue('const a = 1\r\nconst b')
+
+  const clone = streamer.clone()
+  await clone.enqueue(' = 2\nc')
+  expectSameTokens(options, 'const a = 1\r\nconst b = 2\nc', [...clone.tokensStable, ...clone.close().stable])
+})
+
+it('restarts offsets after clear() and close()', async () => {
+  const highlighter = await createHighlighter({
+    langs: ['ts'],
+    themes: ['vitesse-light'],
+  })
+  const options: ShikiStreamTokenizerOptions = { highlighter, lang: 'ts', theme: 'vitesse-light' }
+
+  for (const reset of ['clear', 'close'] as const) {
+    const streamer = new ShikiStreamTokenizer(options)
+    await streamer.enqueue('const a = 1\nconst b')
+    streamer[reset]()
+    const { stable, unstable } = await streamer.enqueue('let c\nd')
+    expectSameTokens(options, 'let c\nd', [...stable, ...unstable])
+  }
+})