feat(transformers): Support word highlighting across multiple tokens

packages/shikiji-transformers/package.json

@@ -35,6 +35,7 @@
     "prepublishOnly": "nr build"
   },
   "dependencies": {
+    "hast-util-to-string": "^3.0.0",
     "shikiji": "workspace:*"
   }
 }

packages/shikiji-transformers/src/shared/highlight-word.ts

@@ -1,35 +1,66 @@
-import type { Element } from 'hast'
+import type { Element, ElementContent, Text } from 'hast'
 import { addClassToHast } from 'shikiji/core'
+import { toString as hastToString } from 'hast-util-to-string'
 
 export function highlightWordInLine(line: Element, ignoredElement: Element | null, word: string, className: string): void {
-  line.children = line.children.flatMap((span) => {
-    if (span.type !== 'element' || span.tagName !== 'span' || span === ignoredElement)
-      return span
+  const content = hastToString(line)
+  let index = content.indexOf(word)
 
-    const textNode = span.children[0]
+  while (index !== -1) {
+    highlightRange(line.children, ignoredElement, index, word.length, className)
 
+    index = content.indexOf(word, index + 1)
+  }
+}
+
+/**
+ * @param elements
+ * @param ignoredElement
+ * @param index highlight beginning index
+ * @param len highlight length
+ */
+function highlightRange(elements: ElementContent[], ignoredElement: Element | null, index: number, len: number, className: string) {
+  let currentIdx = 0
+
+  for (let i = 0; i < elements.length; i++) {
+    const element = elements[i]
+    if (element.type !== 'element' || element.tagName !== 'span' || element === ignoredElement)
+      continue
+    const textNode = element.children[0]
     if (textNode.type !== 'text')
-      return span
+      continue
 
-    return replaceSpan(span, textNode.value, word, className) ?? span
-  })
-}
+    // check if it is overlapped with highlight range
+    if (hasOverlap([currentIdx, currentIdx + textNode.value.length - 1], [index, index + len])) {
+      const start = Math.max(0, index - currentIdx)
+      const length = len - Math.max(0, currentIdx - index)
 
-function inheritElement(original: Element, overrides: Partial<Element>): Element {
-  return {
-    ...original,
-    properties: {
-      ...original.properties,
-    },
-    ...overrides,
+      if (length === 0)
+        continue
+
+      const separated = separateToken(element, textNode, start, length)
+      addClassToHast(separated[1], className)
+
+      // insert
+      const output = separated.filter(Boolean) as Element[]
+      elements.splice(i, 1, ...output)
+      i += output.length - 1
+    }
+
+    currentIdx += textNode.value.length
   }
 }
 
-function replaceSpan(span: Element, text: string, word: string, className: string): Element[] | undefined {
-  const index = text.indexOf(word)
+function hasOverlap(range1: [number, number], range2: [ number, number]): boolean {
+  return (range1[0] <= range2[1]) && (range1[1]) >= range2[0]
+}
 
-  if (index === -1)
-    return
+function separateToken(span: Element, textNode: Text, index: number, len: number): [
+  before: Element | undefined,
+  med: Element,
+  after: Element | undefined,
+] {
+  const text = textNode.value
 
   const createNode = (value: string) => inheritElement(span, {
     children: [
@@ -40,17 +71,19 @@ function replaceSpan(span: Element, text: string, word: string, className: strin
     ],
   })
 
-  const nodes: Element[] = []
-
-  if (index > 0)
-    nodes.push(createNode(text.slice(0, index)))
-
-  const highlightedNode = createNode(word)
-  addClassToHast(highlightedNode, className)
-  nodes.push(highlightedNode)
-
-  if (index + word.length < text.length)
-    nodes.push(createNode(text.slice(index + word.length)))
+  return [
+    index > 0 ? createNode(text.slice(0, index)) : undefined,
+    createNode(text.slice(index, index + len)),
+    index + len < text.length ? createNode(text.slice(index + len)) : undefined,
+  ]
+}
 
-  return nodes
+function inheritElement(original: Element, overrides: Partial<Element>): Element {
+  return {
+    ...original,
+    properties: {
+      ...original.properties,
+    },
+    ...overrides,
+  }
 }

packages/shikiji-transformers/test/fixtures/highlight-word/complex.js

@@ -1,9 +1,9 @@
 export function transformerNotationFocus(
   // [!code word:options.a]
-  options = {}, // [!code word:console.log:2]
+  options = {}, // [!code word:console.log:3]
 ) {
   const options = 'options'
-  console.log(options) // TODO: cross-token highlighting should be supported
+  console.log(options)
   options.a = "HELLO"
   console.log('// [!code word:options.a]')
 }

packages/shikiji-transformers/test/fixtures/highlight-word/complex.js.output.html

@@ -1,4 +1,4 @@
-<pre class="shiki github-dark" style="background-color:#24292e;color:#e1e4e8" tabindex="0"><code><span class="line"><span style="color:#F97583">export</span><span style="color:#F97583"> function</span><span style="color:#B392F0"> transformerNotationFocus</span><span style="color:#E1E4E8">(</span></span><span class="line"><span style="color:#FFAB70">  options</span><span style="color:#F97583"> =</span><span style="color:#E1E4E8"> {}, </span></span><span class="line"><span style="color:#E1E4E8">) {</span></span><span class="line"><span style="color:#F97583">  const</span><span style="color:#79B8FF"> options</span><span style="color:#F97583"> =</span><span style="color:#9ECBFF"> 'options'</span></span><span class="line"><span style="color:#E1E4E8">  console.</span><span style="color:#B392F0">log</span><span style="color:#E1E4E8">(options) </span><span style="color:#6A737D">// TODO: cross-token highlighting should be supported</span></span><span class="line"><span style="color:#E1E4E8">  </span><span style="color:#E1E4E8" class="highlighted-word">options.a</span><span style="color:#E1E4E8"> </span><span style="color:#F97583">=</span><span style="color:#9ECBFF"> "HELLO"</span></span><span class="line"><span style="color:#E1E4E8">  console.</span><span style="color:#B392F0">log</span><span style="color:#E1E4E8">(</span><span style="color:#9ECBFF">'// [!code word:</span><span style="color:#9ECBFF" class="highlighted-word">options.a</span><span style="color:#9ECBFF">]'</span><span style="color:#E1E4E8">)</span></span><span class="line"><span style="color:#E1E4E8">}</span></span><span class="line"></span></code></pre>
+<pre class="shiki github-dark" style="background-color:#24292e;color:#e1e4e8" tabindex="0"><code><span class="line"><span style="color:#F97583">export</span><span style="color:#F97583"> function</span><span style="color:#B392F0"> transformerNotationFocus</span><span style="color:#E1E4E8">(</span></span><span class="line"><span style="color:#FFAB70">  options</span><span style="color:#F97583"> =</span><span style="color:#E1E4E8"> {}, </span></span><span class="line"><span style="color:#E1E4E8">) {</span></span><span class="line"><span style="color:#F97583">  const</span><span style="color:#79B8FF"> options</span><span style="color:#F97583"> =</span><span style="color:#9ECBFF"> 'options'</span></span><span class="line"><span style="color:#E1E4E8">  </span><span style="color:#E1E4E8" class="highlighted-word">console.</span><span style="color:#B392F0" class="highlighted-word">log</span><span style="color:#E1E4E8">(options)</span></span><span class="line"><span style="color:#E1E4E8">  </span><span style="color:#E1E4E8" class="highlighted-word">options.a</span><span style="color:#E1E4E8"> </span><span style="color:#F97583">=</span><span style="color:#9ECBFF"> "HELLO"</span></span><span class="line"><span style="color:#E1E4E8">  console.</span><span style="color:#B392F0">log</span><span style="color:#E1E4E8">(</span><span style="color:#9ECBFF">'// [!code word:</span><span style="color:#9ECBFF" class="highlighted-word">options.a</span><span style="color:#9ECBFF">]'</span><span style="color:#E1E4E8">)</span></span><span class="line"><span style="color:#E1E4E8">}</span></span><span class="line"></span></code></pre>
 <style>
 body { margin: 0; }
 .shiki { padding: 1em; }

packages/shikiji-transformers/test/fixtures/highlight-word/occurrence.js

@@ -1,5 +1,5 @@
 export function transformerNotationFocus(
-  // [!code word:options:4]
+  // [!code word:'options':4]
   options = {},
 ) {
   const options = 'options'

packages/shikiji-transformers/test/fixtures/highlight-word/occurrence.js.output.html

@@ -1,4 +1,4 @@
-<pre class="shiki github-dark" style="background-color:#24292e;color:#e1e4e8" tabindex="0"><code><span class="line"><span style="color:#F97583">export</span><span style="color:#F97583"> function</span><span style="color:#B392F0"> transformerNotationFocus</span><span style="color:#E1E4E8">(</span></span><span class="line"><span style="color:#FFAB70">  </span><span style="color:#FFAB70" class="highlighted-word">options</span><span style="color:#F97583"> =</span><span style="color:#E1E4E8"> {},</span></span><span class="line"><span style="color:#E1E4E8">) {</span></span><span class="line"><span style="color:#F97583">  const</span><span style="color:#79B8FF"> </span><span style="color:#79B8FF" class="highlighted-word">options</span><span style="color:#F97583"> =</span><span style="color:#9ECBFF"> '</span><span style="color:#9ECBFF" class="highlighted-word">options</span><span style="color:#9ECBFF">'</span></span><span class="line"><span style="color:#E1E4E8">  console.</span><span style="color:#B392F0">log</span><span style="color:#E1E4E8">(</span><span style="color:#E1E4E8" class="highlighted-word">options</span><span style="color:#E1E4E8">)</span></span><span class="line"><span style="color:#E1E4E8">  options.a </span><span style="color:#F97583">=</span><span style="color:#9ECBFF"> "HELLO"</span><span style="color:#6A737D"> // should not be highlighted</span></span><span class="line"><span style="color:#E1E4E8">}</span></span><span class="line"></span></code></pre>
+<pre class="shiki github-dark" style="background-color:#24292e;color:#e1e4e8" tabindex="0"><code><span class="line"><span style="color:#F97583">export</span><span style="color:#F97583"> function</span><span style="color:#B392F0"> transformerNotationFocus</span><span style="color:#E1E4E8">(</span></span><span class="line"><span style="color:#FFAB70">  options</span><span style="color:#F97583"> =</span><span style="color:#E1E4E8"> {},</span></span><span class="line"><span style="color:#E1E4E8">) {</span></span><span class="line"><span style="color:#F97583">  const</span><span style="color:#79B8FF"> options</span><span style="color:#F97583"> =</span><span style="color:#9ECBFF"> </span><span style="color:#9ECBFF" class="highlighted-word">'options'</span></span><span class="line"><span style="color:#E1E4E8">  console.</span><span style="color:#B392F0">log</span><span style="color:#E1E4E8">(options)</span></span><span class="line"><span style="color:#E1E4E8">  options.a </span><span style="color:#F97583">=</span><span style="color:#9ECBFF"> "HELLO"</span><span style="color:#6A737D"> // should not be highlighted</span></span><span class="line"><span style="color:#E1E4E8">}</span></span><span class="line"></span></code></pre>
 <style>
 body { margin: 0; }
 .shiki { padding: 1em; }