refactor(source-map): partial application #206

Co-Authored-By: Piotr Tomiak <[email protected]>

packages/language-service/lib/documents.ts

@@ -60,13 +60,13 @@ export class SourceMapWithDocuments {
 	}
 
 	public * getSourcePositions(position: vscode.Position, filter: (data: CodeInformation) => boolean = () => true) {
-		for (const mapped of this.findPositions(position, filter, this.embeddedDocument, this.sourceDocument, 'generatedOffsets', 'sourceOffsets')) {
+		for (const mapped of this.findPositions(position, filter, this.embeddedDocument, this.sourceDocument, 'generatedOffsets')) {
 			yield mapped[0];
 		}
 	}
 
 	public * getGeneratedPositions(position: vscode.Position, filter: (data: CodeInformation) => boolean = () => true) {
-		for (const mapped of this.findPositions(position, filter, this.sourceDocument, this.embeddedDocument, 'sourceOffsets', 'generatedOffsets')) {
+		for (const mapped of this.findPositions(position, filter, this.sourceDocument, this.embeddedDocument, 'sourceOffsets')) {
 			yield mapped[0];
 		}
 	}
@@ -76,10 +76,9 @@ export class SourceMapWithDocuments {
 		filter: (data: CodeInformation) => boolean,
 		fromDoc: TextDocument,
 		toDoc: TextDocument,
-		from: CodeRangeKey,
-		to: CodeRangeKey
+		from: CodeRangeKey
 	) {
-		for (const mapped of this.map.findMatching(fromDoc.offsetAt(position), from, to)) {
+		for (const mapped of this.map.findMatchingOffsets(fromDoc.offsetAt(position), from)) {
 			if (!filter(mapped[1].data)) {
 				continue;
 			}

packages/source-map/lib/sourceMap.ts

@@ -24,38 +24,40 @@ export class SourceMap<Data = unknown> {
 	constructor(public readonly mappings: Mapping<Data>[]) { }
 
 	getSourceStartEnd(generatedStart: number, generatedEnd: number) {
-		return this.findMatchingStartEnd(generatedStart, generatedEnd, 'generatedOffsets', 'sourceOffsets');
+		return this.findMatchingStartEnd(generatedStart, generatedEnd, 'generatedOffsets');
 	}
 
 	getGeneratedStartEnd(sourceStart: number, sourceEnd: number) {
-		return this.findMatchingStartEnd(sourceStart, sourceEnd, 'sourceOffsets', 'generatedOffsets');
+		return this.findMatchingStartEnd(sourceStart, sourceEnd, 'sourceOffsets');
 	}
 
-	* findMatchingStartEnd(start: number, end: number, fromRange: CodeRangeKey, toRange: CodeRangeKey) {
-		for (const [mappedStart, mapping] of this.findMatching(start, fromRange, toRange)) {
+	getSourceOffsets(generatedOffset: number) {
+		return this.findMatchingOffsets(generatedOffset, 'generatedOffsets');
+	}
+
+	getGeneratedOffsets(sourceOffset: number) {
+		return this.findMatchingOffsets(sourceOffset, 'sourceOffsets');
+	}
+
+	* findMatchingStartEnd(start: number, end: number, fromRange: CodeRangeKey) {
+		const toRange: CodeRangeKey = fromRange == 'sourceOffsets' ? 'generatedOffsets' : 'sourceOffsets';
+		for (const [mappedStart, mapping] of this.findMatchingOffsets(start, fromRange)) {
 			const mappedEnd = translateOffset(end, mapping[fromRange], mapping[toRange], getLengths(mapping, fromRange), getLengths(mapping, toRange));
 			if (mappedEnd !== undefined) {
 				yield [mappedStart, mappedEnd, mapping] as const;
 			}
 		};
 	}
 
-	getSourceOffsets(generatedOffset: number) {
-		return this.findMatching(generatedOffset, 'generatedOffsets', 'sourceOffsets');
-	}
-
-	getGeneratedOffsets(sourceOffset: number) {
-		return this.findMatching(sourceOffset, 'sourceOffsets', 'generatedOffsets');
-	}
-
-	* findMatching(offset: number, fromRange: CodeRangeKey, toRange: CodeRangeKey) {
+	* findMatchingOffsets(offset: number, fromRange: CodeRangeKey) {
 		const memo = this.getMemoBasedOnRange(fromRange);
 		if (memo.offsets.length === 0) {
 			return;
 		}
 
 		const { low: start, high: end } = binarySearch(memo.offsets, offset);
 		const skip = new Set<Mapping>();
+		const toRange: CodeRangeKey = fromRange == 'sourceOffsets' ? 'generatedOffsets' : 'sourceOffsets';
 
 		for (let i = start; i <= end; i++) {
 			for (const mapping of memo.mappings[i]) {