feat: store grammar state in weakmap (#804)

packages/core/src/constructors/bundle-factory.ts

@@ -190,7 +190,9 @@ export interface ShorthandsBundle<L extends string, T extends string> {
    * Shorthand for `getLastGrammarState` with auto-loaded theme and language.
    * A singleton highlighter it maintained internally.
    */
-  getLastGrammarState: (code: string, options: CodeToTokensBaseOptions<L, T>) => Promise<GrammarState>
+  getLastGrammarState:
+    | ((element: ThemedToken[][] | Root) => GrammarState)
+    | ((code: string, options: CodeToTokensBaseOptions<L, T>) => Promise<GrammarState>)
 }
 
 export function makeSingletonHighlighter<L extends string, T extends string>(

packages/core/src/constructors/highlighter.ts

@@ -20,7 +20,7 @@ export async function createHighlighterCore(options: HighlighterCoreOptions = {}
   const internal = await createShikiInternal(options)
 
   return {
-    getLastGrammarState: (code, options) => getLastGrammarState(internal, code, options),
+    getLastGrammarState: (...args: any[]) => getLastGrammarState(internal, ...args as [any])!,
     codeToTokensBase: (code, options) => codeToTokensBase(internal, code, options),
     codeToTokensWithThemes: (code, options) => codeToTokensWithThemes(internal, code, options),
     codeToTokens: (code, options) => codeToTokens(internal, code, options),
@@ -43,7 +43,7 @@ export function createHighlighterCoreSync(options: HighlighterCoreOptions<true>
   const internal = createShikiInternalSync(options)
 
   return {
-    getLastGrammarState: (code, options) => getLastGrammarState(internal, code, options),
+    getLastGrammarState: (...args: any[]) => getLastGrammarState(internal, ...args as [any, any]),
     codeToTokensBase: (code, options) => codeToTokensBase(internal, code, options),
     codeToTokensWithThemes: (code, options) => codeToTokensWithThemes(internal, code, options),
     codeToTokens: (code, options) => codeToTokens(internal, code, options),

packages/core/src/highlight/code-to-hast.ts

@@ -1,6 +1,7 @@
 import type {
   CodeToHastOptions,
   CodeToHastRenderOptions,
+  GrammarState,
   ShikiInternal,
   ShikiTransformerContext,
   ShikiTransformerContextCommon,
@@ -14,7 +15,7 @@ import type {
 } from 'hast'
 
 import { FontStyle } from '@shikijs/vscode-textmate'
-
+import { getLastGrammarStateFromMap, setLastGrammarStateToMap } from '../textmate/grammar-state'
 import { addClassToHast, getTokenStyleObject, stringifyTokenStyle } from '../utils'
 import { warnDeprecated } from '../warn'
 import { getTransformers } from './_get-transformers'
@@ -42,6 +43,7 @@ export function codeToHast(
     bg,
     themeName,
     rootStyle,
+    grammarState,
   } = codeToTokens(internal, input, options)
 
   const {
@@ -73,13 +75,15 @@ export function codeToHast(
       rootStyle,
     },
     contextSource,
+    grammarState,
   )
 }
 
 export function tokensToHast(
   tokens: ThemedToken[][],
   options: CodeToHastRenderOptions,
   transformerContext: ShikiTransformerContextSource,
+  grammarState: GrammarState | undefined = getLastGrammarStateFromMap(tokens),
 ): Root {
   const transformers = getTransformers(options)
 
@@ -220,6 +224,9 @@ export function tokensToHast(
   for (const transformer of transformers)
     result = transformer?.root?.call(context, result) || result
 
+  if (grammarState)
+    setLastGrammarStateToMap(result, grammarState)
+
   return result
 }
 

packages/core/src/highlight/code-to-tokens-base.ts

@@ -3,6 +3,8 @@
  *-------------------------------------------------------- */
 import type {
   CodeToTokensBaseOptions,
+  Grammar,
+  GrammarState,
   ShikiInternal,
   ThemedToken,
   ThemedTokenScopeExplanation,
@@ -11,15 +13,15 @@ import type {
 } from '@shikijs/types'
 import type {
   FontStyle,
-  IGrammar,
   IRawThemeSetting,
   StateStack,
 } from '@shikijs/vscode-textmate'
 
+import type { Root } from 'hast'
 import { ShikiError } from '@shikijs/types'
-import { EncodedTokenMetadata, INITIAL } from '@shikijs/vscode-textmate'
 
-import { getGrammarStack, GrammarState } from '../textmate/grammar-state'
+import { EncodedTokenMetadata, INITIAL } from '@shikijs/vscode-textmate'
+import { getGrammarStack, getLastGrammarStateFromMap, GrammarState as GrammarStateImpl, setLastGrammarStateToMap } from '../textmate/grammar-state'
 import { applyColorReplacements, isNoneTheme, isPlainLang, resolveColorReplacements, splitLines } from '../utils'
 import { tokenizeAnsiWithTheme } from './code-to-tokens-ansi'
 
@@ -50,19 +52,29 @@ export function codeToTokensBase(
     if (options.grammarState.lang !== _grammar.name) {
       throw new ShikiError(`Grammar state language "${options.grammarState.lang}" does not match highlight language "${_grammar.name}"`)
     }
-    if (options.grammarState.theme !== themeName) {
-      throw new ShikiError(`Grammar state theme "${options.grammarState.theme}" does not match highlight theme "${themeName}"`)
+    if (!options.grammarState.themes.includes(theme.name)) {
+      throw new ShikiError(`Grammar state themes "${options.grammarState.themes}" do not contain highlight theme "${theme.name}"`)
     }
   }
 
   return tokenizeWithTheme(code, _grammar, theme, colorMap, options)
 }
 
+export function getLastGrammarState(
+  internal: ShikiInternal,
+  element: ThemedToken[][] | Root
+): GrammarState | undefined
 export function getLastGrammarState(
   internal: ShikiInternal,
   code: string,
-  options: CodeToTokensBaseOptions = {},
-): GrammarState {
+  options?: CodeToTokensBaseOptions
+): GrammarState
+export function getLastGrammarState(...args: any[]): GrammarState | undefined {
+  if (args.length === 2) {
+    return getLastGrammarStateFromMap(args[1])
+  }
+
+  const [internal, code, options = {}] = args as [ShikiInternal, string, CodeToTokensBaseOptions]
   const {
     lang = 'text',
     theme: themeName = internal.getLoadedThemes()[0],
@@ -77,7 +89,7 @@ export function getLastGrammarState(
 
   const _grammar = internal.getLanguage(lang)
 
-  return new GrammarState(
+  return new GrammarStateImpl(
     _tokenizeWithTheme(code, _grammar, theme, colorMap, options).stateStack,
     _grammar.name,
     theme.name,
@@ -92,17 +104,27 @@ interface ThemeSettingsSelectors {
 
 export function tokenizeWithTheme(
   code: string,
-  grammar: IGrammar,
+  grammar: Grammar,
   theme: ThemeRegistrationResolved,
   colorMap: string[],
   options: TokenizeWithThemeOptions,
 ): ThemedToken[][] {
-  return _tokenizeWithTheme(code, grammar, theme, colorMap, options).tokens
+  const result = _tokenizeWithTheme(code, grammar, theme, colorMap, options)
+
+  const grammarState = new GrammarStateImpl(
+    _tokenizeWithTheme(code, grammar, theme, colorMap, options).stateStack,
+    grammar.name,
+    theme.name,
+  )
+
+  setLastGrammarStateToMap(result.tokens, grammarState)
+
+  return result.tokens
 }
 
 function _tokenizeWithTheme(
   code: string,
-  grammar: IGrammar,
+  grammar: Grammar,
   theme: ThemeRegistrationResolved,
   colorMap: string[],
   options: TokenizeWithThemeOptions,
@@ -120,7 +142,7 @@ function _tokenizeWithTheme(
   const lines = splitLines(code)
 
   let stateStack = options.grammarState
-    ? getGrammarStack(options.grammarState)
+    ? getGrammarStack(options.grammarState, theme.name) ?? INITIAL
     : options.grammarContextCode != null
       ? _tokenizeWithTheme(
         options.grammarContextCode,

packages/core/src/highlight/code-to-tokens-themes.ts

@@ -4,6 +4,7 @@ import type {
   ThemedToken,
   ThemedTokenWithVariants,
 } from '@shikijs/types'
+import { getLastGrammarStateFromMap, GrammarState, setLastGrammarStateToMap } from '../textmate/grammar-state'
 import { codeToTokensBase } from './code-to-tokens-base'
 
 /**
@@ -19,11 +20,24 @@ export function codeToTokensWithThemes(
     .filter(i => i[1])
     .map(i => ({ color: i[0], theme: i[1]! }))
 
-  const tokens = syncThemesTokenization(
-    ...themes.map(t => codeToTokensBase(internal, code, {
+  const themedTokens = themes.map((t) => {
+    const tokens = codeToTokensBase(internal, code, {
       ...options,
       theme: t.theme,
-    })),
+    })
+    const state = getLastGrammarStateFromMap(tokens)
+    const theme = typeof t.theme === 'string'
+      ? t.theme
+      : t.theme.name
+    return {
+      tokens,
+      state,
+      theme,
+    }
+  })
+
+  const tokens = syncThemesTokenization(
+    ...themedTokens.map(i => i.tokens),
   )
 
   const mergedTokens: ThemedTokenWithVariants[][] = tokens[0]
@@ -54,6 +68,15 @@ export function codeToTokensWithThemes(
       }),
     )
 
+  const mergedGrammarState = themedTokens[0].state
+    ? new GrammarState(
+      Object.fromEntries(themedTokens.map(s => [s.theme, s.state?.getInternalStack(s.theme)])),
+      themedTokens[0].state.lang,
+    )
+    : undefined
+  if (mergedGrammarState)
+    setLastGrammarStateToMap(mergedTokens, mergedGrammarState)
+
   return mergedTokens
 }
 

packages/core/src/highlight/code-to-tokens.ts

@@ -1,5 +1,6 @@
-import type { CodeToTokensOptions, ShikiInternal, ThemedToken, ThemedTokenWithVariants, TokensResult } from '@shikijs/types'
-import { ShikiError } from '../../../types/src/error'
+import type { CodeToTokensOptions, GrammarState, ShikiInternal, ThemedToken, ThemedTokenWithVariants, TokensResult } from '@shikijs/types'
+import { ShikiError } from '@shikijs/types'
+import { getLastGrammarStateFromMap, setLastGrammarStateToMap } from '../textmate/grammar-state'
 import { applyColorReplacements, getTokenStyleObject, resolveColorReplacements } from '../utils'
 import { codeToTokensBase } from './code-to-tokens-base'
 import { codeToTokensWithThemes } from './code-to-tokens-themes'
@@ -19,6 +20,7 @@ export function codeToTokens(
   let tokens: ThemedToken[][]
   let themeName: string
   let rootStyle: string | undefined
+  let grammarState: GrammarState | undefined
 
   if ('themes' in options) {
     const {
@@ -41,6 +43,8 @@ export function codeToTokens(
       options,
     )
 
+    grammarState = getLastGrammarStateFromMap(themeTokens)
+
     if (defaultColor && !themes.find(t => t.color === defaultColor))
       throw new ShikiError(`\`themes\` option must contain the defaultColor key \`${defaultColor}\``)
 
@@ -49,6 +53,9 @@ export function codeToTokens(
     tokens = themeTokens
       .map(line => line.map(token => mergeToken(token, themesOrder, cssVariablePrefix, defaultColor)))
 
+    if (grammarState)
+      setLastGrammarStateToMap(tokens, grammarState)
+
     const themeColorReplacements = themes.map(t => resolveColorReplacements(t.theme, options))
 
     fg = themes.map((t, idx) => (idx === 0 && defaultColor
@@ -73,6 +80,7 @@ export function codeToTokens(
     bg = applyColorReplacements(_theme.bg, colorReplacements)
     fg = applyColorReplacements(_theme.fg, colorReplacements)
     themeName = _theme.name
+    grammarState = getLastGrammarStateFromMap(tokens)
   }
   else {
     throw new ShikiError('Invalid options, either `theme` or `themes` must be provided')
@@ -84,6 +92,7 @@ export function codeToTokens(
     bg,
     themeName,
     rootStyle,
+    grammarState,
   }
 }