trace-mapping: add operations needed for remapping

* Adjust rangeSegments to map to range mapping info such as the end
  point, not just a boolean state of being a range or not.
* Add isRange to allow querying the range mapping data for a segment
* Add traceSegmentsInRange for finding segments in a range

Author: takikawa

packages/trace-mapping/src/by-source.ts

@@ -15,7 +15,7 @@ export type Source = {
 export default function buildBySources(
   decoded: readonly SourceMapSegment[][],
   memos: unknown[],
-  rangeSegments: Set<SourceMapSegment>,
+  rangeSegments: Map<SourceMapSegment, any>,
 ): Source[] {
   const sources: Source[] = memos.map(() => ({ lines: [], rangeSegments: new Set() }));
 

packages/trace-mapping/src/trace-mapping.ts

@@ -1,4 +1,9 @@
-import { encode, decode, encodeRangeMappings, decodeRangeMappings } from '@jridgewell/sourcemap-codec';
+import {
+  encode,
+  decode,
+  encodeRangeMappings,
+  decodeRangeMappings,
+} from '@jridgewell/sourcemap-codec';
 
 import resolver from './resolve';
 import maybeSort from './sort';
@@ -39,6 +44,7 @@ import type {
   XInput,
   SectionedSourceMap,
   Ro,
+  RangeInfo,
 } from './types';
 import type { Source } from './by-source';
 import type { MemoState } from './binary-search';
@@ -156,9 +162,9 @@ export class TraceMap implements SourceMap {
   declare private _decodedRangeMappings: number[][] | undefined;
 
   /**
-   * A set of segments that are range mappings.
+   * A map of segments that are range mappings to info about the range.
    */
-  declare private _rangeSegments: Set<SourceMapSegment> | undefined;
+  declare private _rangeSegments: Map<SourceMapSegment, RangeInfo> | undefined;
 
   constructor(map: Ro<SourceMapInput>, mapUrl?: string | null) {
     const isString = typeof map === 'string';
@@ -272,6 +278,77 @@ export function traceSegment(
   );
 }
 
+/**
+ * Find all the segments that are in a given range. Used to find
+ * segments that are relevant for composing a range mapping with
+ * another range.
+ */
+export function traceSegmentsInRange(
+  map: TraceMap,
+  startLine: number,
+  startColumn: number,
+  endLine: number,
+  endColumn: number,
+): Readonly<SourceMapSegment>[] {
+  const lines = decodedMappings(map);
+  const segments = startLine < lines.length ? lines[startLine] : [];
+  const memo = cast(map)._decodedMemo;
+
+  if (startLine >= lines.length || endLine >= lines.length) return [];
+
+  const segmentsInRange = [];
+
+  let startIndex = memoizedBinarySearchSegments(segments, startColumn, memo, startLine);
+  if (bsFound) {
+    startIndex = lowerBound(segments, startColumn, startIndex);
+  } else {
+    startIndex++;
+  }
+  if (startIndex === segments.length) return [];
+
+  const range = (start: number, end: number) => {
+    return Array.from({ length: end - start + 1 }, (_, i: number) => i + start);
+  };
+  const gen = (line: number) => (index: number) => {
+    return lines[line][index];
+  };
+
+  function previousPosition(line: number, column: number) {
+    if (column === 0) {
+      line--;
+      column = Infinity;
+    } else {
+      column--;
+    }
+
+    return { line, column };
+  }
+
+  if (startLine == endLine) {
+    // We need to use the position decremented by one position because we
+    // want the index of a true lower bound, not an equal position. That is,
+    // the end point of the range should be exclusive and not inclusive.
+    const { line, column } = previousPosition(endLine, endColumn);
+    let endIndex = memoizedBinarySearchSegments(segments, column, memo, line);
+    if (bsFound) {
+      endIndex = upperBound(segments, column, endIndex);
+    }
+    segmentsInRange.push(...range(startIndex, endIndex).map(gen(startLine)));
+  } else {
+    segmentsInRange.push(...range(startIndex, segments.length - 1).map(gen(startLine)));
+    for (let i = startLine + 1; i < endLine; i++) {
+      segmentsInRange.push(...range(0, lines[i].length - 1).map(gen(i)));
+    }
+    let endIndex = memoizedBinarySearchSegments(lines[endLine], endColumn, memo, endLine);
+    if (bsFound) {
+      endIndex = upperBound(lines[endLine], endColumn, endIndex);
+    }
+    segmentsInRange.push(...range(0, endIndex).map(gen(endLine)));
+  }
+
+  return segmentsInRange;
+}
+
 /**
  * A higher-level API to find the source/line/column associated with a generated line/column
  * (think, from a stack trace). Line is 1-based, but column is 0-based, due to legacy behavior in
@@ -397,6 +474,16 @@ export function isIgnored(map: TraceMap, source: string): boolean {
   return index === -1 ? false : ignoreList.includes(index);
 }
 
+/**
+ * Determines if a segment is for a range mapping, and if so returns
+ * some metadata about the range.
+ */
+export function isRange(map: TraceMap, segment: Readonly<SourceMapSegment>): RangeInfo | false {
+  const decoded = decodedMappings(map);
+  const rangeSegments = initRangeSegments(map, decoded);
+  return rangeSegments.get(segment as SourceMapSegment) || false;
+}
+
 /**
  * A helper that skips sorting of the input map's mappings array, which can be expensive for larger
  * maps.
@@ -485,23 +572,23 @@ function GMapping(
  */
 function traceSegmentInternal<T extends SourceMapSegment | ReverseSegment>(
   lines: readonly T[][],
-  rangeSegments: Set<T>,
+  rangeSegments: Set<T> | Map<T, RangeInfo>,
   memo: MemoState,
   line: number,
   column: number,
   bias: Bias,
 ): T | null;
 function traceSegmentInternal<T extends SourceMapSegment | ReverseSegment>(
   lines: readonly T[][],
-  rangeSegments: Set<T>,
+  rangeSegments: Set<T> | Map<T, RangeInfo>,
   memo: MemoState,
   line: number,
   column: number,
   bias: Bias,
 ): T | null;
 function traceSegmentInternal<T extends SourceMapSegment | ReverseSegment>(
   lines: readonly T[][],
-  rangeSegments: Set<T>,
+  rangeSegments: Set<T> | Map<T, RangeInfo>,
   memo: MemoState,
   line: number,
   column: number,
@@ -700,7 +787,7 @@ function initRangeSegments(map: TraceMap, decoded: readonly SourceMapSegment[][]
   const existing = cast(map)._rangeSegments;
   if (existing != null) return existing;
 
-  const set = new Set<SourceMapSegment>();
+  const set = new Map<SourceMapSegment, RangeInfo>();
   cast(map)._rangeSegments = set;
   const rangeMappings = decodedRangeMappings(map);
   if (rangeMappings == null) return set;
@@ -710,13 +797,37 @@ function initRangeSegments(map: TraceMap, decoded: readonly SourceMapSegment[][]
     const ranges = rangeMappings[i];
     for (let j = 0; j < ranges.length; j++) {
       const seg = line[ranges[j]];
-      set.add(seg);
+      const { line: endLine, index: endIndex } = findNextSegment(decoded, i, ranges[j]);
+      const endSegment = endLine && endIndex ? decoded[endLine][endIndex] : null;
+      set.set(seg, { line: i, endLine, endSegment });
     }
   }
 
   return set;
 }
 
+/**
+ * Find the index of the next segment in this line or in subsequent lines.
+ * Return null if there was no next segment.
+ */
+function findNextSegment<T extends SourceMapSegment>(
+  lines: readonly T[][],
+  line: number,
+  index: number,
+): { line: number | null; index: number | null } {
+  if (index + 1 < lines[line].length) {
+    return { line, index: index + 1 };
+  } else {
+    for (let i = line + 1; i < lines.length; i++) {
+      for (let j = 0; j < lines[i].length; j++) {
+        return { line: i, index: j };
+      }
+    }
+  }
+
+  return { line: null, index: null };
+}
+
 /**
  * If we didn't find a match on this line, back searches to find the previous
  * line that has a segment.

packages/trace-mapping/src/types.ts

@@ -91,6 +91,12 @@ export type EachMapping =
       name: string | null;
     };
 
+export type RangeInfo = {
+  line: number;
+  endLine: number | null;
+  endSegment: Readonly<SourceMapSegment> | null;
+};
+
 export abstract class SourceMap {
   declare version: SourceMapV3['version'];
   declare file: SourceMapV3['file'];

packages/trace-mapping/test/trace-mapping.test.ts

@@ -8,6 +8,7 @@ import {
   encodedMappings,
   decodedMappings,
   traceSegment,
+  traceSegmentsInRange,
   originalPositionFor,
   generatedPositionFor,
   presortedDecodedMap,
@@ -446,6 +447,25 @@ describe('TraceMap', () => {
           [{ line: 1, column: 13 }],
         );
       });
+      describe('rangeMappings operations', () => {
+        // 2nd line, 5th mapping is a range mapping
+        const mapWithRangeMappings = replaceField(map, 'rangeMappings', ';F;');
+
+        it('traceSegmentsInRange', () => {
+          const tracer = new TraceMap(mapWithRangeMappings);
+          assert.deepEqual(traceSegmentsInRange(tracer, 0, 9, 0, 16), [
+            [9, 0, 0, 9, 0],
+            [12, 0, 0, 0],
+            [13, 0, 0, 13, 1],
+          ]);
+          assert.deepEqual(traceSegmentsInRange(tracer, 0, 16, 1, 9), [
+            [16, 0, 0, 0],
+            [18, 0, 0, 33],
+            [4, 0, 1, 4],
+            [8, 0, 1, 10],
+          ]);
+        });
+      });
     };
   }