perf: store single source-map segments directly

__tests__/source-map/segment-mapping.spec.ts

@@ -0,0 +1,56 @@
+import {
+  appendSegment,
+  compactSegments,
+  segmentAt,
+  segmentCount,
+} from '../../src/source-map/segment-mapping';
+import type {
+  MarkdownSourceMapSegment,
+  SegmentMapping,
+} from '../../src/source-map/types';
+
+const first: MarkdownSourceMapSegment = {
+  valueStart: 0,
+  valueEnd: 1,
+  sourceStart: 0,
+  sourceEnd: 1,
+  kind: 'literal',
+};
+
+const second: MarkdownSourceMapSegment = {
+  valueStart: 1,
+  valueEnd: 2,
+  sourceStart: 1,
+  sourceEnd: 3,
+  kind: 'escape',
+};
+
+describe('segment mapping storage', () => {
+  it('stores one segment directly and promotes the second segment', () => {
+    const mappings = new WeakMap<object, SegmentMapping>();
+    const node = {};
+
+    appendSegment(mappings, node, first);
+    expect(mappings.get(node)).toBe(first);
+
+    appendSegment(mappings, node, second);
+    expect(mappings.get(node)).toEqual([first, second]);
+  });
+
+  it('compacts completed single-segment arrays', () => {
+    expect(compactSegments([first])).toBe(first);
+
+    const multiple = [first, second];
+    expect(compactSegments(multiple)).toBe(multiple);
+  });
+
+  it('reads direct and array mappings without normalization', () => {
+    expect(segmentCount(first)).toBe(1);
+    expect(segmentAt(first, 0)).toBe(first);
+    expect(segmentAt(first, 1)).toBeUndefined();
+
+    const multiple = [first, second];
+    expect(segmentCount(multiple)).toBe(2);
+    expect(segmentAt(multiple, 1)).toBe(second);
+  });
+});

scripts/profile-source-map-heap.mjs

@@ -1,7 +1,7 @@
 // Heap-profile source-map retained structures.
 //
-// Run:  pnpm run profile:source-map -- <fixture> <phase>
-// e.g. pnpm run profile:source-map -- segments build
+// Run:  pnpm run profile:source-map <fixture> <phase>
+// e.g. pnpm run profile:source-map segments build
 //
 // Fixtures:
 //   many-nodes   – 10k small text nodes (flat list)
@@ -88,7 +88,12 @@ function collectMappedNodes(root) {
 
 // ── Main ────────────────────────────────────────────────────────────────
 
-const [, , fixtureName, phase] = process.argv;
+const args = process.argv.slice(2);
+
+// Accept the npm-style separator because pnpm passes it to the script.
+if (args[0] === '--') args.shift();
+
+const [fixtureName, phase] = args;
 
 if (!fixtureName || !phase || !FIXTURES[fixtureName] || !['build', 'raw', 'range'].includes(phase)) {
   console.error(

src/source-map/build-source-map.ts

@@ -20,10 +20,12 @@ import {
   SourceMapUnavailableError,
 } from './errors';
 import { recordingExtension } from './recording-extension';
+import { compactSegments, segmentAt, segmentCount } from './segment-mapping';
 import type {
   MarkdownSourceMap,
   MarkdownSourceMapSegment,
   ParsedMarkdownDocument,
+  SegmentMapping,
   SourceSpan,
 } from './types';
 
@@ -35,15 +37,15 @@ const { micromarkExtensions, fromMarkdownExtensions } = getParserExtensions();
 
 interface RecordingState {
   /** node -> ordered, gap-free, non-overlapping segments. */
-  segments: WeakMap<object, MarkdownSourceMapSegment[]>
+  segments: WeakMap<object, SegmentMapping>
   /** inlineCode node -> value segments (see buildInlineCodeSegments). */
-  inlineCodeSegments: WeakMap<object, MarkdownSourceMapSegment[]>
+  inlineCodeSegments: WeakMap<object, SegmentMapping>
   /** code node -> value segments (see buildCodeSegments). */
-  codeSegments: WeakMap<object, MarkdownSourceMapSegment[]>
+  codeSegments: WeakMap<object, SegmentMapping>
   /** code node -> source point for an empty value. */
   emptyCodeOffsets: WeakMap<object, number>
   /** link / definition node -> normalized URL segments. */
-  urlSegments: WeakMap<object, MarkdownSourceMapSegment[]>
+  urlSegments: WeakMap<object, SegmentMapping>
   /** link / definition node -> source point for an empty URL. */
   emptyUrlOffsets: WeakMap<object, number>
   /** link / definition node -> parser-confirmed destination content span. */
@@ -326,38 +328,40 @@ function buildUrlSegments(
  * repeatedly.
  */
 function findSegmentIndexAt(
-  segs: MarkdownSourceMapSegment[],
+  segs: SegmentMapping,
   valueIndex: number,
 ): number | undefined {
   let lo = 0;
-  let hi = segs.length - 1;
+  let hi = segmentCount(segs) - 1;
   while (lo < hi) {
     const mid = (lo + hi + 1) >> 1;
-    if (segs[mid].valueStart <= valueIndex)
+    if (segmentAt(segs, mid)!.valueStart <= valueIndex)
       lo = mid;
     else hi = mid - 1;
   }
-  const seg = segs[lo];
+  const seg = segmentAt(segs, lo);
   return seg && valueIndex >= seg.valueStart && valueIndex < seg.valueEnd
     ? lo
     : undefined;
 }
 
 function findSegmentAt(
-  segs: MarkdownSourceMapSegment[],
+  segs: SegmentMapping,
   valueIndex: number,
 ): MarkdownSourceMapSegment | undefined {
   const index = findSegmentIndexAt(segs, valueIndex);
-  return index === undefined ? undefined : segs[index];
+  return index === undefined ? undefined : segmentAt(segs, index);
 }
 
 /** Prefix count of source gaps before each segment. */
-function buildSourceGapPrefix(segs: MarkdownSourceMapSegment[]): number[] {
+function buildSourceGapPrefix(segs: SegmentMapping): number[] {
   const prefix = [0];
-  for (let index = 0; index + 1 < segs.length; index++) {
+  for (let index = 0; index + 1 < segmentCount(segs); index++) {
+    const current = segmentAt(segs, index)!;
+    const next = segmentAt(segs, index + 1)!;
     prefix.push(
       prefix[index]
-      + Number(segs[index].sourceEnd !== segs[index + 1].sourceStart),
+      + Number(current.sourceEnd !== next.sourceStart),
     );
   }
   return prefix;
@@ -372,7 +376,7 @@ interface RangeResolutionMessages {
 }
 
 interface SegmentRangeOptions {
-  segments: MarkdownSourceMapSegment[]
+  segments: SegmentMapping
   valueLength: number
   valueStart: number
   valueEnd: number
@@ -441,15 +445,16 @@ function resolveEmptyRange(
     emptyOffset,
     messages,
   } = options;
-  if (segments.length === 0) {
+  const count = segmentCount(segments);
+  if (count === 0) {
     if (valueLength === 0 && valueStart === 0 && emptyOffset !== undefined)
       return emptyOffset;
     throw new RangeError(messages.incomplete);
   }
   if (valueStart === 0)
-    return segments[0].sourceStart;
+    return segmentAt(segments, 0)!.sourceStart;
   if (valueStart === valueLength)
-    return segments[segments.length - 1].sourceEnd;
+    return segmentAt(segments, count - 1)!.sourceEnd;
   const segment = findSegmentAt(segments, valueStart);
   if (segment && valueStart === segment.valueStart)
     return segment.sourceStart;
@@ -476,7 +481,7 @@ function resolveSegmentRange(options: SegmentRangeOptions): ParsedPosition {
 
   if (valueStart === valueEnd)
     return pointRange(resolveEmptyRange(options));
-  if (segments.length === 0)
+  if (segmentCount(segments) === 0)
     throw new RangeError(messages.incomplete);
 
   const startSegmentIndex = findSegmentIndexAt(segments, valueStart);
@@ -489,8 +494,8 @@ function resolveSegmentRange(options: SegmentRangeOptions): ParsedPosition {
       throw new RangeError(messages.nonContiguous);
   }
 
-  const startSegment = segments[startSegmentIndex];
-  const endSegment = segments[endSegmentIndex];
+  const startSegment = segmentAt(segments, startSegmentIndex)!;
+  const endSegment = segmentAt(segments, endSegmentIndex)!;
 
   let startOffset: number;
   if (startSegment.kind !== 'literal')
@@ -553,7 +558,7 @@ export const parseMdWithSourceMap = (md: string): ParsedMarkdownDocument => {
   // Snapshot offsets before consumers can mutate positions.
   // getRaw() must describe the source that originally produced each node.
   const originalOffsets = new WeakMap<object, readonly [number, number]>();
-  const sourceGapPrefixes = new WeakMap<MarkdownSourceMapSegment[], number[]>();
+  const sourceGapPrefixes = new WeakMap<SegmentMapping, number[]>();
 
   function indexNode(node: TraversableNode): void {
     // The standard mdast handler compiles inline code.
@@ -564,7 +569,7 @@ export const parseMdWithSourceMap = (md: string): ParsedMarkdownDocument => {
     ) {
       const segments = buildInlineCodeSegments(md, node);
       if (segments)
-        state.inlineCodeSegments.set(node, segments);
+        state.inlineCodeSegments.set(node, compactSegments(segments));
     }
 
     if (
@@ -573,7 +578,7 @@ export const parseMdWithSourceMap = (md: string): ParsedMarkdownDocument => {
     ) {
       const mapping = buildCodeSegments(md, node);
       if (mapping) {
-        state.codeSegments.set(node, mapping.segments);
+        state.codeSegments.set(node, compactSegments(mapping.segments));
         if (mapping.emptyOffset !== undefined) {
           state.emptyCodeOffsets.set(node, mapping.emptyOffset);
         }
@@ -587,7 +592,7 @@ export const parseMdWithSourceMap = (md: string): ParsedMarkdownDocument => {
       const bounds = state.urlSourceSpans.get(node);
       const segments = bounds ? buildUrlSegments(md, node, bounds) : undefined;
       if (segments) {
-        state.urlSegments.set(node, segments.segments);
+        state.urlSegments.set(node, compactSegments(segments.segments));
         if (segments.emptyOffset !== undefined)
           state.emptyUrlOffsets.set(node, segments.emptyOffset);
       }
@@ -651,13 +656,14 @@ export const parseMdWithSourceMap = (md: string): ParsedMarkdownDocument => {
         );
       }
       const segs = state.segments.get(node as object);
-      if (segs && segs.length > 0) {
+      if (segs && segmentCount(segs) > 0) {
         assertUnmodified(node as object);
         // Text nodes with a source map: use the full recorded outer-token
         // span, which covers the complete raw source that produced the value
         // (e.g. '&#0;' includes the trailing ';' even though the parser
         // positions the text node one code unit earlier).
-        return md.slice(segs[0].sourceStart, segs[segs.length - 1].sourceEnd);
+        const last = segmentCount(segs) - 1;
+        return md.slice(segmentAt(segs, 0)!.sourceStart, segmentAt(segs, last)!.sourceEnd);
       }
       if (
         state.inlineCodeSegments.has(node as object)

src/source-map/recording-extension.ts

@@ -1,10 +1,11 @@
 import { decodeNamedCharacterReference } from 'decode-named-character-reference';
 import { decodeNumericCharacterReference } from 'micromark-util-decode-numeric-character-reference';
 import type { ParsedPoint } from '../types';
-import type { MarkdownSourceMapSegment, SourceSpan } from './types';
+import { appendSegment } from './segment-mapping';
+import type { MarkdownSourceMapSegment, SegmentMapping, SourceSpan } from './types';
 
 interface RecordingState {
-  segments: WeakMap<object, MarkdownSourceMapSegment[]>
+  segments: WeakMap<object, SegmentMapping>
   urlSourceSpans: WeakMap<object, SourceSpan>
 }
 
@@ -86,9 +87,6 @@ export function recordingExtension(state: RecordingState) {
       node.children.push(tail);
     }
     this.stack.push(tail);
-    if (!state.segments.has(tail)) {
-      state.segments.set(tail, []);
-    }
   };
 
   // For escapes and character references the decoded value is appended by the
@@ -116,14 +114,11 @@ export function recordingExtension(state: RecordingState) {
     const valueStart = tail.value.length;
     tail.value += slice;
     tail.position.end = point(token.end);
-    const segs = state.segments.get(tail);
-    if (segs) {
-      segs.push({
-        valueStart,
-        valueEnd: valueStart + slice.length,
-        ...metadata,
-      });
-    }
+    appendSegment(state.segments, tail, {
+      valueStart,
+      valueEnd: valueStart + slice.length,
+      ...metadata,
+    });
   };
 
   const onexitcharacterreferencevalue = function (this: CompileContext, token: any) {
@@ -151,15 +146,12 @@ export function recordingExtension(state: RecordingState) {
     const valueStart = tail.value.length;
     tail.value += value;
     tail.position.end = point(token.end);
-    const segs = state.segments.get(tail);
-    if (segs) {
-      segs.push({
-        valueStart,
-        valueEnd: valueStart + value.length,
-        ...construct,
-        kind,
-      });
-    }
+    appendSegment(state.segments, tail, {
+      valueStart,
+      valueEnd: valueStart + value.length,
+      ...construct,
+      kind,
+    });
   };
 
   const onexitlineending = function (this: CompileContext, token: any) {

src/source-map/segment-mapping.ts

@@ -0,0 +1,35 @@
+import type { MarkdownSourceMapSegment, SegmentMapping } from './types';
+
+export function appendSegment(
+  mappings: WeakMap<object, SegmentMapping>,
+  node: object,
+  segment: MarkdownSourceMapSegment,
+): void {
+  const current = mappings.get(node);
+  if (!current) {
+    mappings.set(node, segment);
+  }
+  else if (Array.isArray(current)) {
+    current.push(segment);
+  }
+  else {
+    mappings.set(node, [current, segment]);
+  }
+}
+
+export function compactSegments(segments: MarkdownSourceMapSegment[]): SegmentMapping {
+  return segments.length === 1 ? segments[0] : segments;
+}
+
+export function segmentCount(mapping: SegmentMapping): number {
+  return Array.isArray(mapping) ? mapping.length : 1;
+}
+
+export function segmentAt(
+  mapping: SegmentMapping,
+  index: number,
+): MarkdownSourceMapSegment | undefined {
+  return Array.isArray(mapping)
+    ? mapping[index]
+    : index === 0 ? mapping : undefined;
+}

src/source-map/types.ts

@@ -65,6 +65,15 @@ export interface MarkdownSourceMapSegment {
   kind: SourceMapSegmentKind
 }
 
+/**
+ * Stores one segment without an array allocation.
+ *
+ * @internal
+ */
+export type SegmentMapping =
+  | MarkdownSourceMapSegment
+  | MarkdownSourceMapSegment[];
+
 /**
  * Sidecar source map produced alongside a parse. Maps supported value nodes to
  * compressed segments that reconstruct supported normalized fields from the