fix: multiline code

index.js

@@ -1,12 +1,23 @@
 /**
- * @typedef {import('unist').Node & {properties: Object<any, any>}} Node
- * @typedef {import('unist').Parent & {properties: Object<any, any>}} Parent
+ * @typedef {import('hast').Node & {properties: Object<any, any>}} Node
+ * @typedef {import('hast').Parent & {properties: Object<any, any>}} Parent
+ * @typedef {import('hast').Root} Root
  * @typedef {import('unist-util-visit').Visitor<Node>} Visitor
+ * @typedef Options options
+ *   Configuration.
+ * @property {boolean} [showLineNumbers]
+ *   Set `showLineNumbers` to `true` to always display line number
+ * @property {boolean} [ignoreMissing]
+ *   Set `ignoreMissing` to `true` to ignore unsupported languages and line highlighting when no language is specified
  */
 
 import { visit } from 'unist-util-visit'
 import { toString } from 'hast-util-to-string'
 import { refractor } from 'refractor/lib/all.js'
+import { toHtml } from 'hast-util-to-html'
+import { filter } from 'unist-util-filter'
+import { unified } from 'unified'
+import parse from 'rehype-parse'
 import rangeParser from 'parse-numeric-range'
 
 /**
@@ -73,20 +84,68 @@ const splitLine = (text) => {
 }
 
 /**
- * Rehype plugin that highlights code blocks with refractor (prismjs)
- *
- * Set `showLineNumbers` to `true` to always display line number
+ * Split line to div node with className `code-line`
  *
- * Set `ignoreMissing` to `true` to ignore unsupported languages and line highlighting when no language is specified
+ * @param {import('refractor').RefractorRoot} ast
+ * @return {Root}
+ */
+const getNodePosition = (ast) => {
+  // @ts-ignore
+  let html = toHtml(ast)
+  const hast = unified().use(parse, { emitParseErrors: true, fragment: true }).parse(html)
+  return hast
+}
+
+/**
+ * Split multiline text nodes into individual nodes with positioning
  *
- * @typedef {{ showLineNumbers?: boolean, ignoreMissing?: boolean }} RehypePrismOptions
- * @param {RehypePrismOptions} options
- * @return {Visitor}
+ * @param {Parent['children']} ast
+ * @return {Parent['children']}
  */
-const rehypePrism = (options) => {
-  options = options || {}
+const splitTextByLine = (ast) => {
+  //@ts-ignore
+  return ast.reduce((result, node) => {
+    if (node.type === 'text') {
+      if (node.value.indexOf('\n') === -1) {
+        result.push(node)
+        return result
+      }
+
+      const lines = node.value.split('\n')
+      for (const [i, line] of lines.entries()) {
+        result.push({
+          type: 'text',
+          value: i === lines.length - 1 ? line : line + '\n',
+          position: {
+            start: { line: node.position.start.line + i },
+            end: { line: node.position.start.line + i },
+          },
+        })
+      }
+
+      return result
+    }
+
+    if (node.children) {
+      // @ts-ignore
+      node.children = splitTextByLine(node.children)
+      result.push(node)
+      return result
+    }
 
+    result.push(node)
+    return result
+  }, [])
+}
+
+/**
+ * Rehype plugin that highlights code blocks with refractor (prismjs)
+ *
+ * @type {import('unified').Plugin<[Options?], Root>}
+ */
+const rehypePrism = (options = {}) => {
   return (tree) => {
+    // @ts-ignore
     visit(tree, 'element', visitor)
   }
 
@@ -112,6 +171,25 @@ const rehypePrism = (options) => {
       meta = `${lang} ${meta}`
     }
 
+    let refractorRoot
+    let langError = false
+
+    // Syntax highlight
+    if (lang) {
+      try {
+        // @ts-ignore
+        refractorRoot = refractor.highlight(toString(node), lang)
+        refractorRoot = getNodePosition(refractorRoot)
+        refractorRoot.children = splitTextByLine(refractorRoot.children)
+      } catch (err) {
+        if (options.ignoreMissing && /Unknown language/.test(err.message)) {
+          langError = true
+        } else {
+          throw err
+        }
+      }
+    }
+
     const shouldHighlightLine = calculateLinesToHighlight(meta)
     // @ts-ignore
     const codeLineArray = splitLine(toString(node))
@@ -129,16 +207,12 @@ const rehypePrism = (options) => {
       }
 
       // Syntax highlight
-      if (lang && line.children) {
-        try {
-          line.children = refractor.highlight(line.children[0].value, lang).children
-        } catch (err) {
-          // eslint-disable-next-line no-empty
-          if (options.ignoreMissing && /Unknown language/.test(err.message)) {
-          } else {
-            throw err
-          }
-        }
+      if (lang && line.children && !langError) {
+        const treeExtract = filter(
+          refractorRoot,
+          (node) => node.position.start.line <= i + 1 && node.position.end.line >= i + 1
+        )
+        line.children = treeExtract.children
       }
     }
 

package.json

@@ -42,23 +42,27 @@
   },
   "homepage": "https://github.com/timlrx/rehype-prism-plus#readme",
   "dependencies": {
+    "hast-util-to-html": "^8.0.1",
     "hast-util-to-string": "^2.0.0",
-    "parse-numeric-range": "^1.2.0",
+    "parse-numeric-range": "^1.3.0",
     "refractor": "^4.1.1",
+    "rehype-parse": "^8.0.2",
+    "unified": "^10.1.0",
+    "unist-util-filter": "^4.0.0",
     "unist-util-visit": "^4.0.0"
   },
   "devDependencies": {
     "dedent": "^0.7.0",
-    "eslint": "^7.29.0",
+    "eslint": "^7.32.0",
     "eslint-config-prettier": "^8.3.0",
     "eslint-plugin-node": "^11.1.0",
     "husky": "^4.0.0",
-    "lint-staged": "^11.0.0",
+    "lint-staged": "^11.1.2",
     "prettier": "^2.3.2",
     "rehype": "^12.0.0",
-    "typescript": "^4.3.4",
+    "typescript": "^4.3.5",
     "uvu": "^0.5.1",
-    "vite": "^2.3.8"
+    "vite": "^2.5.0"
   },
   "prettier": {
     "printWidth": 100,

test.js

@@ -70,9 +70,12 @@ y
 </div>
 `).trim()
   const expected = dedent`
-    <div>
-    <pre class="language-py"><code class="language-py"><div class="code-line">x</div><div class="code-line">\n</div><div class="code-line">y</div></code></pre>
-    </div>
+  <div>
+  <pre class="language-py"><code class="language-py"><div class="code-line">x
+  </div><div class="code-line">
+  </div><div class="code-line">y
+  </div></code></pre>
+  </div>
     `
   assert.is(result, expected)
 })