From f2b39561ffbf96ed12f2e65789747496097e85cf Mon Sep 17 00:00:00 2001
From: Anthony Fu <github@antfu.me>
Date: Fri, 28 Jun 2024 10:10:03 +0200
Subject: [PATCH] feat: expose `getLastGrammarState` shorthand

---
 packages/core/src/bundle-factory.ts | 18 ++++++++++++++++--
 packages/shiki/src/bundle-full.ts   |  1 +
 packages/shiki/src/bundle-web.ts    |  1 +
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/bundle-factory.ts b/packages/core/src/bundle-factory.ts
index 63f3ae122..d89c8036e 100644
--- a/packages/core/src/bundle-factory.ts
+++ b/packages/core/src/bundle-factory.ts
@@ -1,5 +1,5 @@
 import type { Root } from 'hast'
-import type { BundledHighlighterOptions, CodeToHastOptions, CodeToTokensBaseOptions, CodeToTokensOptions, CodeToTokensWithThemesOptions, HighlighterCoreOptions, HighlighterGeneric, LanguageInput, RequireKeys, SpecialLanguage, SpecialTheme, ThemeInput, ThemedToken, ThemedTokenWithVariants, TokensResult } from './types'
+import type { BundledHighlighterOptions, CodeToHastOptions, CodeToTokensBaseOptions, CodeToTokensOptions, CodeToTokensWithThemesOptions, GrammarState, HighlighterCoreOptions, HighlighterGeneric, LanguageInput, RequireKeys, SpecialLanguage, SpecialTheme, ThemeInput, ThemedToken, ThemedTokenWithVariants, TokensResult } from './types'
 import { isSpecialLang, isSpecialTheme } from './utils'
 import { createHighlighterCore } from './highlighter'
 import { ShikiError } from './error'
@@ -115,9 +115,15 @@ export interface ShorthandsBundle<L extends string, T extends string> {
   codeToTokensWithThemes: (code: string, options: RequireKeys<CodeToTokensWithThemesOptions<L, T>, 'themes' | 'lang'>) => Promise<ThemedTokenWithVariants[][]>
 
   /**
-   * Get internal singleton highlighter.
+   * Get the singleton highlighter.
    */
   getSingletonHighlighter: (options?: Partial<BundledHighlighterOptions<L, T>>) => Promise<HighlighterGeneric<L, T>>
+
+  /**
+   * Shorthand for `getLastGrammarState` with auto-loaded theme and language.
+   * A singleton highlighter it maintained internally.
+   */
+  getLastGrammarState: (code: string, options: CodeToTokensBaseOptions<L, T>) => Promise<GrammarState>
 }
 
 export function makeSingletonHighlighter<L extends string, T extends string>(createHighlighter: CreateHighlighterFactory<L, T>) {
@@ -196,5 +202,13 @@ export function createSingletonShorthands<L extends string, T extends string >(
       })
       return shiki.codeToTokensWithThemes(code, options)
     },
+
+    async getLastGrammarState(code, options) {
+      const shiki = await getSingletonHighlighter({
+        langs: [options.lang as L],
+        themes: [options.theme as T],
+      })
+      return shiki.getLastGrammarState(code, options)
+    },
   }
 }
diff --git a/packages/shiki/src/bundle-full.ts b/packages/shiki/src/bundle-full.ts
index 7dc5b750d..ab741d502 100644
--- a/packages/shiki/src/bundle-full.ts
+++ b/packages/shiki/src/bundle-full.ts
@@ -40,6 +40,7 @@ export const {
   codeToTokensBase,
   codeToTokensWithThemes,
   getSingletonHighlighter,
+  getLastGrammarState,
 } = /* @__PURE__ */ createSingletonShorthands<
   BundledLanguage,
   BundledTheme
diff --git a/packages/shiki/src/bundle-web.ts b/packages/shiki/src/bundle-web.ts
index 025bd48dd..822c04f40 100644
--- a/packages/shiki/src/bundle-web.ts
+++ b/packages/shiki/src/bundle-web.ts
@@ -40,6 +40,7 @@ export const {
   codeToTokens,
   codeToTokensWithThemes,
   getSingletonHighlighter,
+  getLastGrammarState,
 } = /* @__PURE__ */ createSingletonShorthands<
   BundledLanguage,
   BundledTheme