From 5d112d0e37c0059aa6fbe6f574c4450afad8af47 Mon Sep 17 00:00:00 2001
From: pengzhanbo <volodymyr@foxmail.com>
Date: Wed, 2 Oct 2024 01:29:28 +0800
Subject: [PATCH] feat(plugin-shikiji): auto scan code blocks language (#243)

---
 .../src/node/highlight/index.ts               |  1 +
 .../src/node/highlight/scanLanguages.ts       | 34 +++++++++++++++++++
 .../plugin-shikiji/src/node/shikiPlugin.ts    | 15 +++++++-
 .../plugin-shikiji/src/node/utils/index.ts    |  4 +++
 4 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 plugins/plugin-shikiji/src/node/highlight/scanLanguages.ts

diff --git a/plugins/plugin-shikiji/src/node/highlight/index.ts b/plugins/plugin-shikiji/src/node/highlight/index.ts
index 80f96a38..e01cb0d3 100644
--- a/plugins/plugin-shikiji/src/node/highlight/index.ts
+++ b/plugins/plugin-shikiji/src/node/highlight/index.ts
@@ -1 +1,2 @@
 export * from './highlight.js'
+export * from './scanLanguages.js'
diff --git a/plugins/plugin-shikiji/src/node/highlight/scanLanguages.ts b/plugins/plugin-shikiji/src/node/highlight/scanLanguages.ts
new file mode 100644
index 00000000..219cc0a9
--- /dev/null
+++ b/plugins/plugin-shikiji/src/node/highlight/scanLanguages.ts
@@ -0,0 +1,34 @@
+import type { App } from 'vuepress'
+import fs from 'node:fs/promises'
+import path from 'node:path'
+import fg from 'fast-glob'
+import { bundledLanguages } from 'shiki'
+
+const languages = Object.keys(bundledLanguages)
+const RE_FENCE = /`{3,}[^`]*?(\s|$)/g
+
+export async function scanLanguages(app: App): Promise<string[]> {
+  const source = app.dir.source()
+  const pattern = ['**/*.md', '!.vuepress', '!node_modules']
+  const files = await fg(pattern, {
+    cwd: source,
+  })
+  const langs = new Set<string>()
+
+  for (const file of files) {
+    const filepath = path.join(source, file)
+    const content = await fs.readFile(filepath, 'utf-8')
+    const matched = content.match(RE_FENCE)
+    if (matched) {
+      for (const match of matched) {
+        let lang = match.replace(/`{3,}/, '').trim()
+        if (lang.includes(':'))
+          lang = lang.split(':')[0]
+        if (lang)
+          langs.add(lang)
+      }
+    }
+  }
+
+  return Array.from(langs).filter(lang => languages.includes(lang))
+}
diff --git a/plugins/plugin-shikiji/src/node/shikiPlugin.ts b/plugins/plugin-shikiji/src/node/shikiPlugin.ts
index fda2082c..8d85d839 100644
--- a/plugins/plugin-shikiji/src/node/shikiPlugin.ts
+++ b/plugins/plugin-shikiji/src/node/shikiPlugin.ts
@@ -7,7 +7,7 @@ import type {
 } from './types.js'
 import { isPlainObject } from 'vuepress/shared'
 import { copyCodeButtonPlugin } from './copy-code-button/index.js'
-import { highlight } from './highlight/index.js'
+import { highlight, scanLanguages } from './highlight/index.js'
 import {
   collapsedLinesPlugin,
   highlightLinesPlugin,
@@ -15,6 +15,7 @@ import {
   preWrapperPlugin,
 } from './markdown/index.js'
 import { prepareClientConfigFile } from './prepareClientConfigFile.js'
+import { logger } from './utils/index.js'
 
 export interface ShikiPluginOptions
   extends HighlighterOptions, LineNumberOptions, PreWrapperOptions {
@@ -49,9 +50,21 @@ export function shikiPlugin({
     }),
 
     extendsMarkdown: async (md, app) => {
+      const start = performance.now()
       const theme = options.theme ?? { light: 'github-light', dark: 'github-dark' }
 
+      if (!options.languages || !options.languages.length) {
+        options.languages = await scanLanguages(app)
+        if (app.env.isDebug) {
+          logger.info(`scan languages: ${JSON.stringify(options.languages)}`)
+          logger.info(`scan languages in: ${(performance.now() - start).toFixed(2)}ms`)
+        }
+      }
+
       md.options.highlight = await highlight(theme, options)
+      if (app.env.isDebug) {
+        logger.info(`highlight Loaded in: ${(performance.now() - start).toFixed(2)}ms`)
+      }
 
       md.use(highlightLinesPlugin)
       md.use(preWrapperPlugin, { preWrapper })
diff --git a/plugins/plugin-shikiji/src/node/utils/index.ts b/plugins/plugin-shikiji/src/node/utils/index.ts
index 290d390c..a8a899dd 100644
--- a/plugins/plugin-shikiji/src/node/utils/index.ts
+++ b/plugins/plugin-shikiji/src/node/utils/index.ts
@@ -1,3 +1,7 @@
+import { Logger } from '@vuepress/helper'
+
+export const logger = new Logger('@vuepress-plume/plugin-shikiji')
+
 export * from './attrsToLines.js'
 export * from './collapsedLines.js'
 export * from './lru.js'