feat(core): expose dispose function (#707)

antfu · web-flow · commit 2c5b38737eea · 2024-06-21T16:20:58.000+02:00
diff --git a/packages/core/src/internal.ts b/packages/core/src/internal.ts
@@ -24,7 +24,9 @@ let instancesCount = 0
 export async function getShikiInternal(options: HighlighterCoreOptions = {}): Promise<ShikiInternal> {
   instancesCount += 1
   if (options.warnings !== false && instancesCount >= 10 && instancesCount % 10 === 0)
-    console.warn(`[Shiki] ${instancesCount} instances have been created. Shiki is supposed to be used as a singleton, consider refactoring your code to cache your highlighter instance.`)
+    console.warn(`[Shiki] ${instancesCount} instances have been created. Shiki is supposed to be used as a singleton, consider refactoring your code to cache your highlighter instance; Or call \`highlighter.dispose()\` to release unused instances.`)
+
+  let isDisposed = false
 
   async function normalizeGetter<T>(p: MaybeGetter<T>): Promise<T> {
     return Promise.resolve(typeof p === 'function' ? (p as any)() : p).then(r => r.default || r)
@@ -67,6 +69,7 @@ export async function getShikiInternal(options: HighlighterCoreOptions = {}): Pr
   let _lastTheme: string | ThemeRegistrationAny
 
   function getLanguage(name: string | LanguageRegistration) {
+    ensureNotDisposed()
     const _lang = _registry.getGrammar(typeof name === 'string' ? name : name.name)
     if (!_lang)
       throw new ShikiError(`Language \`${name}\` not found, you may need to load it first`)
@@ -76,13 +79,15 @@ export async function getShikiInternal(options: HighlighterCoreOptions = {}): Pr
   function getTheme(name: string | ThemeRegistrationAny): ThemeRegistrationResolved {
     if (name === 'none')
       return { bg: '', fg: '', name: 'none', settings: [], type: 'dark' }
+    ensureNotDisposed()
     const _theme = _registry.getTheme(name)
     if (!_theme)
       throw new ShikiError(`Theme \`${name}\` not found, you may need to load it first`)
     return _theme
   }
 
   function setTheme(name: string | ThemeRegistrationAny) {
+    ensureNotDisposed()
     const theme = getTheme(name)
     if (_lastTheme !== name) {
       _registry.setTheme(theme)
@@ -96,18 +101,22 @@ export async function getShikiInternal(options: HighlighterCoreOptions = {}): Pr
   }
 
   function getLoadedThemes() {
+    ensureNotDisposed()
     return _registry.getLoadedThemes()
   }
 
   function getLoadedLanguages() {
+    ensureNotDisposed()
     return _registry.getLoadedLanguages()
   }
 
   async function loadLanguage(...langs: (LanguageInput | SpecialLanguage)[]) {
+    ensureNotDisposed()
     await _registry.loadLanguages(await resolveLangs(langs))
   }
 
   async function loadTheme(...themes: (ThemeInput | SpecialTheme)[]) {
+    ensureNotDisposed()
     await Promise.all(
       themes.map(async theme =>
         isSpecialTheme(theme)
@@ -117,6 +126,19 @@ export async function getShikiInternal(options: HighlighterCoreOptions = {}): Pr
     )
   }
 
+  function ensureNotDisposed() {
+    if (isDisposed)
+      throw new ShikiError('Shiki instance has been disposed')
+  }
+
+  function dispose() {
+    if (isDisposed)
+      return
+    isDisposed = true
+    _registry.dispose()
+    instancesCount -= 1
+  }
+
   return {
     setTheme,
     getTheme,
@@ -125,5 +147,7 @@ export async function getShikiInternal(options: HighlighterCoreOptions = {}): Pr
     getLoadedLanguages,
     loadLanguage,
     loadTheme,
+    dispose,
+    [Symbol.dispose]: dispose,
   }
 }
diff --git a/packages/core/src/registry.ts b/packages/core/src/registry.ts
@@ -6,9 +6,9 @@ import { normalizeTheme } from './normalize'
 import { ShikiError } from './error'
 
 export class Registry extends TextMateRegistry {
-  private _resolvedThemes: Record<string, ThemeRegistrationResolved> = {}
-  private _resolvedGrammars: Record<string, IGrammar> = {}
-  private _langMap: Record<string, LanguageRegistration> = {}
+  private _resolvedThemes: Map<string, ThemeRegistrationResolved> = new Map()
+  private _resolvedGrammars: Map<string, IGrammar> = new Map()
+  private _langMap: Map<string, LanguageRegistration> = new Map()
   private _langGraph: Map<string, LanguageRegistration> = new Map()
 
   private _textmateThemeCache = new WeakMap<IRawTheme, TextMateTheme>()
@@ -29,15 +29,15 @@ export class Registry extends TextMateRegistry {
 
   public getTheme(theme: ThemeRegistrationAny | string) {
     if (typeof theme === 'string')
-      return this._resolvedThemes[theme]
+      return this._resolvedThemes.get(theme)
     else
       return this.loadTheme(theme)
   }
 
   public loadTheme(theme: ThemeRegistrationAny): ThemeRegistrationResolved {
     const _theme = normalizeTheme(theme)
     if (_theme.name) {
-      this._resolvedThemes[_theme.name] = _theme
+      this._resolvedThemes.set(_theme.name, _theme)
       // Reset cache
       this._loadedThemesCache = null
     }
@@ -46,7 +46,7 @@ export class Registry extends TextMateRegistry {
 
   public getLoadedThemes() {
     if (!this._loadedThemesCache)
-      this._loadedThemesCache = Object.keys(this._resolvedThemes)
+      this._loadedThemesCache = [...this._resolvedThemes.keys()]
     return this._loadedThemesCache
   }
 
@@ -76,14 +76,17 @@ export class Registry extends TextMateRegistry {
         resolved.add(name)
       }
     }
-    return this._resolvedGrammars[name]
+    return this._resolvedGrammars.get(name)
   }
 
   public async loadLanguage(lang: LanguageRegistration) {
     if (this.getGrammar(lang.name))
       return
 
-    const embeddedLazilyBy = new Set(Object.values(this._langMap).filter(i => i.embeddedLangsLazy?.includes(lang.name)))
+    const embeddedLazilyBy = new Set(
+      [...this._langMap.values()]
+        .filter(i => i.embeddedLangsLazy?.includes(lang.name)),
+    )
 
     this._resolver.addLanguage(lang)
 
@@ -95,7 +98,7 @@ export class Registry extends TextMateRegistry {
     // @ts-expect-error Private members, set this to override the previous grammar (that can be a stub)
     this._syncRegistry._rawGrammars.set(lang.scopeName, lang)
     const g = await this.loadGrammarWithConfiguration(lang.scopeName, 1, grammarConfig)
-    this._resolvedGrammars[lang.name] = g!
+    this._resolvedGrammars.set(lang.name, g!)
     if (lang.aliases) {
       lang.aliases.forEach((alias) => {
         this._alias[alias] = lang.name
@@ -107,14 +110,14 @@ export class Registry extends TextMateRegistry {
     // If there is a language that embeds this language lazily, we need to reload it
     if (embeddedLazilyBy.size) {
       for (const e of embeddedLazilyBy) {
-        delete this._resolvedGrammars[e.name]
+        this._resolvedGrammars.delete(e.name)
         // Reset cache
         this._loadedLanguagesCache = null
         // @ts-expect-error clear cache
         this._syncRegistry?._injectionGrammars?.delete(e.scopeName)
         // @ts-expect-error clear cache
         this._syncRegistry?._grammars?.delete(e.scopeName)
-        await this.loadLanguage(this._langMap[e.name])
+        await this.loadLanguage(this._langMap.get(e.name)!)
       }
     }
   }
@@ -124,6 +127,15 @@ export class Registry extends TextMateRegistry {
     await this.loadLanguages(this._langs)
   }
 
+  public override dispose(): void {
+    super.dispose()
+    this._resolvedThemes.clear()
+    this._resolvedGrammars.clear()
+    this._langMap.clear()
+    this._langGraph.clear()
+    this._loadedThemesCache = null
+  }
+
   public async loadLanguages(langs: LanguageRegistration[]) {
     for (const lang of langs)
       this.resolveEmbeddedLanguages(lang)
@@ -146,17 +158,20 @@ export class Registry extends TextMateRegistry {
   }
 
   public getLoadedLanguages() {
-    if (!this._loadedLanguagesCache)
-      this._loadedLanguagesCache = Object.keys({ ...this._resolvedGrammars, ...this._alias })
+    if (!this._loadedLanguagesCache) {
+      this._loadedLanguagesCache = [
+        ...new Set([...this._resolvedGrammars.keys(), ...Object.keys(this._alias)]),
+      ]
+    }
     return this._loadedLanguagesCache
   }
 
   private resolveEmbeddedLanguages(lang: LanguageRegistration) {
-    this._langMap[lang.name] = lang
+    this._langMap.set(lang.name, lang)
     this._langGraph.set(lang.name, lang)
     if (lang.embeddedLangs) {
       for (const embeddedLang of lang.embeddedLangs)
-        this._langGraph.set(embeddedLang, this._langMap[embeddedLang])
+        this._langGraph.set(embeddedLang, this._langMap.get(embeddedLang)!)
     }
   }
 }
diff --git a/packages/core/src/types/highlighter.ts b/packages/core/src/types/highlighter.ts
@@ -48,6 +48,14 @@ export interface ShikiInternal<BundledLangKeys extends string = never, BundledTh
    * Special-handled themes like `none` are not included.
    */
   getLoadedThemes: () => string[]
+  /**
+   * Dispose the internal registry and release resources
+   */
+  dispose: () => void
+  /**
+   * Dispose the internal registry and release resources
+   */
+  [Symbol.dispose]: () => void
 }
 
 /**
diff --git a/packages/shiki/test/core.test.ts b/packages/shiki/test/core.test.ts
@@ -179,4 +179,19 @@ describe('errors', () => {
     await expect(() => shiki.codeToHtml('console.log("Hi")', { lang: 'mylang', theme: 'nord' }))
       .toThrowErrorMatchingInlineSnapshot(`[ShikiError: Circular alias \`mylang -> mylang2 -> mylang\`]`)
   })
+
+  it('throw on using disposed instance', async () => {
+    const shiki = await getHighlighterCore({
+      themes: [nord],
+      langs: [js as any],
+    })
+
+    expect(shiki.codeToHtml('console.log("Hi")', { lang: 'javascript', theme: 'nord' }))
+      .toContain('console')
+
+    shiki.dispose()
+
+    expect(() => shiki.codeToHtml('console.log("Hi")', { lang: 'javascript', theme: 'nord' }))
+      .toThrowErrorMatchingInlineSnapshot(`[ShikiError: Shiki instance has been disposed]`)
+  })
 })
