fix(core): preserve chars highlights across token trees

Author: atomiks

packages/core/src/chars/charsHighlighter.ts

@@ -2,19 +2,12 @@ import type { Element } from 'hast';
 import type { CharsHighlighterOptions, CharsElement } from '../types';
 import { getElementsToHighlight } from './getElementsToHighlight';
 import { wrapHighlightedChars } from './wrapHighlightedChars';
-import { toString as hastToString } from 'hast-util-to-string';
 import { isElement } from '../utils';
 
-/**
- * Stands in for a node excluded from the remaining text, so the fragments
- * around it stay separated. NUL cannot occur in the highlighted source.
- */
-const BOUNDARY = '\0';
-
 /**
  * Loops through the child nodes and finds the nodes that make up the chars.
  * If the chars cross node boundaries, those nodes are wrapped with
- * <span data-highlighted-chars-mark>, and that node is passed to
+ * <mark data-highlighted-chars-mark>, and that node is passed to
  * onVisitHighlightedChars.
  *
  * If a node partially matches the chars, its content is replaced with the
@@ -30,83 +23,42 @@ export function charsHighlighter(
   ) => void,
 ) {
   const { ranges = [] } = options;
-  const textContent = hastToString(element);
 
   // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: char matching is inherently branchy
   charsList.forEach((chars, index) => {
-    if (chars && textContent?.includes(chars)) {
-      let textContent = hastToString(element);
-      let startIndex = 0;
-
-      while (textContent.includes(chars)) {
-        // Snapshot the remaining text so we can bail out if an iteration
-        // fails to make progress. The recomputed `textContent` excludes
-        // already-highlighted nodes, so a productive iteration always makes
-        // it strictly shorter; if it does not shrink, no occurrence was
-        // consumed and continuing would loop forever.
-        const previousTextContent = textContent;
-        const currentCharsRange = ranges[index] || [];
-        const id = `${chars}-${index}`;
+    if (!chars) {
+      return;
+    }
 
-        const elementsToWrap = getElementsToHighlight(
-          element,
-          chars,
-          startIndex,
-        );
+    const currentCharsRange = ranges[index] || [];
+    const id = `${chars}-${index}`;
+    let startIndex = 0;
 
-        // maybe throw / notify due to failure here
-        if (elementsToWrap.length === 0) break;
+    while (true) {
+      const elementsToWrap = getElementsToHighlight(element, chars, startIndex);
 
-        // Count the occurrence only once it has resolved to real elements.
-        // Counting before the match is confirmed lets a candidate that never
-        // materializes consume a number, so a later genuine occurrence is
-        // tested against the wrong index and a range like `/lo/2` misses it.
-        const occurrence = (options.counterMap.get(id) || 0) + 1;
-        options.counterMap.set(id, occurrence);
+      if (elementsToWrap.length === 0) break;
 
-        const ignoreChars =
-          currentCharsRange.length > 0 &&
-          !currentCharsRange.includes(occurrence);
+      const matchStartIndex = element.children.indexOf(elementsToWrap[0]);
+      if (matchStartIndex === -1) break;
 
-        wrapHighlightedChars(
-          element,
-          elementsToWrap,
-          options,
-          ignoreChars,
-          onVisitHighlightedChars,
-        );
+      const occurrence = (options.counterMap.get(id) || 0) + 1;
+      const ignoreChars =
+        currentCharsRange.length > 0 && !currentCharsRange.includes(occurrence);
 
-        // re-start from the 'last' node (the chars or part of them may exist
-        // multiple times in the same node)
-        // account for possible extra nodes added from split with - 2
-        startIndex = Math.max(
-          elementsToWrap[elementsToWrap.length - 1].index - 2,
-          0,
-        );
+      const didWrap = wrapHighlightedChars(
+        element,
+        elementsToWrap,
+        chars,
+        options,
+        ignoreChars,
+        onVisitHighlightedChars,
+      );
 
-        textContent = element.children
-          .map((childNode) => {
-            const props = isElement(childNode) ? childNode.properties : {};
-            if (
-              props &&
-              !Object.hasOwn(props, 'rehype-pretty-code-visited') &&
-              !Object.hasOwn(props, 'data-highlighted-chars-mark')
-            ) {
-              return hastToString(childNode);
-            }
-            // An excluded node is a gap in the text, not an absence. Dropping
-            // it would join its neighbours: once `balloon` is split around an
-            // ignored `lo`, `bal` + `on;` reads as `balon;`, which contains an
-            // occurrence of `lo` the code never had. Keep a boundary so
-            // fragments on either side cannot spell a synthetic match.
-            return BOUNDARY;
-          })
-          .join('');
+      if (!didWrap) break;
 
-        // Safety guard: if the remaining text is unchanged, this iteration
-        // consumed nothing, so stop instead of spinning forever.
-        if (textContent === previousTextContent) break;
-      }
+      options.counterMap.set(id, occurrence);
+      startIndex = matchStartIndex;
     }
   });