feat(language-core): reimplement multi-source mapping

packages/eslint/index.ts

@@ -31,7 +31,7 @@ export function createProcessor(
 	const documents = new FileMap<{
 		sourceDocument: TextDocument;
 		embeddedDocuments: TextDocument[];
-		codes: VirtualCode[];
+		codes: VirtualCode<string>[];
 	}>(caseSensitive);
 	return {
 		supportsAutofix,

packages/language-core/index.ts

@@ -7,15 +7,15 @@ export * from './lib/utils';
 import { SourceMap } from '@volar/source-map';
 import type * as ts from 'typescript';
 import { LinkedCodeMap } from './lib/linkedCodeMap';
-import type { CodeInformation, Language, LanguagePlugin, SourceScript, VirtualCode } from './lib/types';
+import type { CodeInformation, CodegenContext, Language, LanguagePlugin, SourceScript, VirtualCode } from './lib/types';
 
 export function createLanguage<T>(
 	plugins: LanguagePlugin<T>[],
 	scriptRegistry: Map<T, SourceScript<T>>,
 	sync: (id: T) => void
 ): Language<T> {
-	const virtualCodeToSourceScriptMap = new WeakMap<VirtualCode, SourceScript<T>>();
-	const virtualCodeToSourceMap = new WeakMap<ts.IScriptSnapshot, [ts.IScriptSnapshot, SourceMap<CodeInformation>]>();
+	const virtualCodeToSourceScriptMap = new WeakMap<VirtualCode<T>, SourceScript<T>>();
+	const virtualCodeToSourceMap = new WeakMap<ts.IScriptSnapshot, WeakMap<ts.IScriptSnapshot, SourceMap<CodeInformation>>>();
 	const virtualCodeToLinkedCodeMap = new WeakMap<ts.IScriptSnapshot, [ts.IScriptSnapshot, LinkedCodeMap | undefined]>();
 
 	return {
@@ -48,14 +48,15 @@ export function createLanguage<T>(
 						this.delete(id);
 						return this.set(id, snapshot, languageId);
 					}
-					else if (sourceScript.snapshot !== snapshot) {
+					else if (sourceScript.isAssociationDirty || sourceScript.snapshot !== snapshot) {
 						// snapshot updated
 						sourceScript.snapshot = snapshot;
+						const codegenCtx = prepareCreateVirtualCode(sourceScript);
 						if (sourceScript.generated) {
 							const { updateVirtualCode, createVirtualCode } = sourceScript.generated.languagePlugin;
 							const newVirtualCode = updateVirtualCode
-								? updateVirtualCode(id, sourceScript.generated.root, snapshot)
-								: createVirtualCode?.(id, languageId, snapshot);
+								? updateVirtualCode(id, sourceScript.generated.root, snapshot, codegenCtx)
+								: createVirtualCode?.(id, languageId, snapshot, codegenCtx);
 							if (newVirtualCode) {
 								sourceScript.generated.root = newVirtualCode;
 								sourceScript.generated.embeddedCodes.clear();
@@ -70,6 +71,7 @@ export function createLanguage<T>(
 								return;
 							}
 						}
+						triggerTargetsDirty(sourceScript);
 					}
 					else {
 						// not changed
@@ -78,10 +80,17 @@ export function createLanguage<T>(
 				}
 				else {
 					// created
-					const sourceScript: SourceScript<T> = { id, languageId, snapshot };
+					const sourceScript: SourceScript<T> = {
+						id: id,
+						languageId,
+						snapshot,
+						associatedIds: new Set(),
+						targetIds: new Set(),
+					};
 					scriptRegistry.set(id, sourceScript);
+
 					for (const languagePlugin of _plugins) {
-						const virtualCode = languagePlugin.createVirtualCode?.(id, languageId, snapshot);
+						const virtualCode = languagePlugin.createVirtualCode?.(id, languageId, snapshot, prepareCreateVirtualCode(sourceScript));
 						if (virtualCode) {
 							sourceScript.generated = {
 								root: virtualCode,
@@ -95,36 +104,58 @@ export function createLanguage<T>(
 							break;
 						}
 					}
+
 					return sourceScript;
 				}
 			},
 			delete(id) {
-				const value = scriptRegistry.get(id);
-				if (value) {
-					if (value.generated) {
-						value.generated.languagePlugin.disposeVirtualCode?.(id, value.generated.root);
-					}
+				const sourceScript = scriptRegistry.get(id);
+				if (sourceScript) {
+					sourceScript.generated?.languagePlugin.disposeVirtualCode?.(id, sourceScript.generated.root);
 					scriptRegistry.delete(id);
+					triggerTargetsDirty(sourceScript);
 				}
 			},
 		},
 		maps: {
 			get(virtualCode) {
-				const sourceScript = virtualCodeToSourceScriptMap.get(virtualCode)!;
+				for (const map of this.forEach(virtualCode)) {
+					return map[2];
+				}
+				throw `no map found for ${virtualCode.id}`;
+			},
+			*forEach(virtualCode) {
 				let mapCache = virtualCodeToSourceMap.get(virtualCode.snapshot);
-				if (mapCache?.[0] !== sourceScript.snapshot) {
-					if (virtualCode.mappings.some(mapping => mapping.source)) {
-						throw 'not implemented';
-					}
+				if (!mapCache) {
 					virtualCodeToSourceMap.set(
 						virtualCode.snapshot,
-						mapCache = [
-							sourceScript.snapshot,
-							new SourceMap(virtualCode.mappings),
-						]
+						mapCache = new WeakMap()
 					);
 				}
-				return mapCache[1];
+
+				const sourceScript = virtualCodeToSourceScriptMap.get(virtualCode)!;
+				if (!mapCache.has(sourceScript.snapshot)) {
+					mapCache.set(
+						sourceScript.snapshot,
+						new SourceMap(virtualCode.mappings)
+					);
+				}
+				yield [sourceScript.id, sourceScript.snapshot, mapCache.get(sourceScript.snapshot)!];
+
+				if (virtualCode.associatedScriptMappings) {
+					for (const [relatedScriptId, relatedMappings] of virtualCode.associatedScriptMappings) {
+						const relatedSourceScript = scriptRegistry.get(relatedScriptId);
+						if (relatedSourceScript) {
+							if (!mapCache.has(relatedSourceScript.snapshot)) {
+								mapCache.set(
+									relatedSourceScript.snapshot,
+									new SourceMap(relatedMappings)
+								);
+							}
+							yield [relatedSourceScript.id, relatedSourceScript.snapshot, mapCache.get(relatedSourceScript.snapshot)!];
+						}
+					}
+				}
 			},
 		},
 		linkedCodeMaps: {
@@ -146,9 +177,37 @@ export function createLanguage<T>(
 			},
 		},
 	};
+
+	function triggerTargetsDirty(sourceScript: SourceScript<T>) {
+		sourceScript.targetIds.forEach(id => {
+			const sourceScript = scriptRegistry.get(id);
+			if (sourceScript) {
+				sourceScript.isAssociationDirty = true;
+			}
+		});
+	}
+
+	function prepareCreateVirtualCode(sourceScript: SourceScript<T>): CodegenContext<T> {
+		for (const id of sourceScript.associatedIds) {
+			scriptRegistry.get(id)?.targetIds.delete(sourceScript.id);
+		}
+		sourceScript.associatedIds.clear();
+		sourceScript.isAssociationDirty = false;
+		return {
+			getAssociatedScript(id) {
+				sync(id);
+				const relatedSourceScript = scriptRegistry.get(id);
+				if (relatedSourceScript) {
+					relatedSourceScript.targetIds.add(sourceScript.id);
+					sourceScript.associatedIds.add(relatedSourceScript.id);
+				}
+				return relatedSourceScript;
+			},
+		};
+	}
 }
 
-export function* forEachEmbeddedCode(virtualCode: VirtualCode): Generator<VirtualCode> {
+export function* forEachEmbeddedCode<T>(virtualCode: VirtualCode<T>): Generator<VirtualCode<T>> {
 	yield virtualCode;
 	if (virtualCode.embeddedCodes) {
 		for (const embeddedCode of virtualCode.embeddedCodes) {

packages/language-core/lib/linkedCodeMap.ts

@@ -1,6 +1,6 @@
 import { SourceMap } from '@volar/source-map';
 
-export class LinkedCodeMap extends SourceMap {
+export class LinkedCodeMap extends SourceMap<any> {
 	*getLinkedOffsets(start: number) {
 		for (const mapped of this.getGeneratedOffsets(start)) {
 			yield mapped[0];

packages/language-core/lib/types.ts

@@ -2,19 +2,20 @@ import type { Mapping, SourceMap } from '@volar/source-map';
 import type * as ts from 'typescript';
 import type { LinkedCodeMap } from './linkedCodeMap';
 
-export interface Language<T> {
+export interface Language<T = unknown> {
 	plugins: LanguagePlugin<T>[];
 	scripts: {
 		get(id: T): SourceScript<T> | undefined;
 		set(id: T, snapshot: ts.IScriptSnapshot, languageId?: string, plugins?: LanguagePlugin<T>[]): SourceScript<T> | undefined;
 		delete(id: T): void;
-		fromVirtualCode(virtualCode: VirtualCode): SourceScript<T>;
+		fromVirtualCode(virtualCode: VirtualCode<T>): SourceScript<T>;
 	};
 	maps: {
-		get(virtualCode: VirtualCode): SourceMap<CodeInformation>;
+		get(virtualCode: VirtualCode<T>): SourceMap<CodeInformation>;
+		forEach(virtualCode: VirtualCode<T>): Generator<[id: T, snapshot: ts.IScriptSnapshot, map: SourceMap<CodeInformation>]>;
 	};
 	linkedCodeMaps: {
-		get(virtualCode: VirtualCode): LinkedCodeMap | undefined;
+		get(virtualCode: VirtualCode<T>): LinkedCodeMap | undefined;
 	};
 	typescript?: {
 		configFileName: string | undefined;
@@ -23,31 +24,35 @@ export interface Language<T> {
 			sync?(): Promise<number>;
 		};
 		languageServiceHost: ts.LanguageServiceHost;
-		getExtraServiceScript(fileName: string): TypeScriptExtraServiceScript | undefined;
+		getExtraServiceScript(fileName: string): TypeScriptExtraServiceScript<T> | undefined;
 		asScriptId(fileName: string): T;
 		asFileName(scriptId: T): string;
 	};
 }
 
-export interface SourceScript<T> {
+export interface SourceScript<T = unknown> {
 	id: T;
 	languageId: string;
 	snapshot: ts.IScriptSnapshot;
+	targetIds: Set<T>;
+	associatedIds: Set<T>;
+	isAssociationDirty?: boolean;
 	generated?: {
-		root: VirtualCode;
+		root: VirtualCode<T>;
 		languagePlugin: LanguagePlugin<T>;
-		embeddedCodes: Map<string, VirtualCode>;
+		embeddedCodes: Map<string, VirtualCode<T>>;
 	};
 }
 
 export type CodeMapping = Mapping<CodeInformation>;
 
-export interface VirtualCode {
+export interface VirtualCode<T = unknown> {
 	id: string;
 	languageId: string;
 	snapshot: ts.IScriptSnapshot;
 	mappings: CodeMapping[];
-	embeddedCodes?: VirtualCode[];
+	associatedScriptMappings?: Map<T, CodeMapping[]>;
+	embeddedCodes?: VirtualCode<T>[];
 	linkedCodeMappings?: Mapping[];
 }
 
@@ -77,51 +82,55 @@ export interface CodeInformation {
 	format?: boolean;
 }
 
-export interface TypeScriptServiceScript {
-	code: VirtualCode;
+export interface TypeScriptServiceScript<T = unknown> {
+	code: VirtualCode<T>;
 	extension: '.ts' | '.js' | '.mts' | '.mjs' | '.cjs' | '.cts' | '.d.ts' | string;
 	scriptKind: ts.ScriptKind;
 	/** See #188 */
 	preventLeadingOffset?: boolean;
 }
 
-export interface TypeScriptExtraServiceScript extends TypeScriptServiceScript {
+export interface TypeScriptExtraServiceScript<T = unknown> extends TypeScriptServiceScript<T> {
 	fileName: string;
 }
 
-export interface LanguagePlugin<T, K extends VirtualCode = VirtualCode> {
+export interface LanguagePlugin<T = unknown, K extends VirtualCode<T> = VirtualCode<T>> {
 	/**
 	 * For files that are not opened in the IDE, the language ID will not be synchronized to the language server, so a hook is needed to parse the language ID of files that are known extension but not opened in the IDE.
 	 */
 	getLanguageId(scriptId: T): string | undefined;
 	/**
 	 * Generate a virtual code.
 	 */
-	createVirtualCode?(scriptId: T, languageId: string, snapshot: ts.IScriptSnapshot): K | undefined;
+	createVirtualCode?(scriptId: T, languageId: string, snapshot: ts.IScriptSnapshot, ctx: CodegenContext<T>): K | undefined;
 	/**
 	 * Incremental update a virtual code. If not provide, call createVirtualCode again.
 	 */
-	updateVirtualCode?(scriptId: T, virtualCode: K, newSnapshot: ts.IScriptSnapshot): K | undefined;
+	updateVirtualCode?(scriptId: T, virtualCode: K, newSnapshot: ts.IScriptSnapshot, ctx: CodegenContext<T>): K | undefined;
 	/**
 	 * Cleanup a virtual code.
 	 */
 	disposeVirtualCode?(scriptId: T, virtualCode: K): void;
-	typescript?: TypeScriptGenericOptions<K> & TypeScriptNonTSPluginOptions<K>;
+	typescript?: TypeScriptGenericOptions<T, K> & TypeScriptNonTSPluginOptions<T, K>;
+}
+
+export interface CodegenContext<T = unknown> {
+	getAssociatedScript(scriptId: T): SourceScript<T> | undefined;
 }
 
 /**
  * The following options available to all situations.
  */
-interface TypeScriptGenericOptions<T> {
+interface TypeScriptGenericOptions<T, K> {
 	extraFileExtensions: ts.FileExtensionInfo[];
 	resolveHiddenExtensions?: boolean;
-	getServiceScript(rootVirtualCode: T): TypeScriptServiceScript | undefined;
+	getServiceScript(root: K): TypeScriptServiceScript<T> | undefined;
 }
 
 /**
  * The following options will not be available in TS plugin.
  */
-interface TypeScriptNonTSPluginOptions<T> {
-	getExtraServiceScripts?(fileName: string, rootVirtualCode: T): TypeScriptExtraServiceScript[];
+interface TypeScriptNonTSPluginOptions<T, K> {
+	getExtraServiceScripts?(fileName: string, rootVirtualCode: K): TypeScriptExtraServiceScript<T>[];
 	resolveLanguageServiceHost?(host: ts.LanguageServiceHost): ts.LanguageServiceHost;
 }

packages/language-server/lib/register/registerEditorFeatures.ts

@@ -92,8 +92,9 @@ export function registerEditorFeatures(server: LanguageServer) {
 		const virtualCode = sourceScript?.generated?.embeddedCodes.get(params.virtualCodeId);
 		if (virtualCode) {
 			const mappings: Record<string, CodeMapping[]> = {};
-			const map = languageService.context.documents.getSourceMap(virtualCode);
-			mappings[map.sourceDocument.uri] = map.map.mappings;
+			for (const map of languageService.context.documents.getMaps(virtualCode)) {
+				mappings[map.sourceDocument.uri] = map.map.mappings;
+			}
 			return {
 				content: virtualCode.snapshot.getText(0, virtualCode.snapshot.getLength()),
 				mappings,

packages/language-server/protocol.ts

@@ -155,7 +155,6 @@ export namespace GetVirtualCodeRequest {
 	};
 	export type ResponseType = {
 		content: string;
-		// TODO: Simplify this, no map required
 		mappings: Record<string, CodeMapping[]>;
 	};
 	export type ErrorType = never;

packages/language-service/lib/features/provideCallHierarchyItems.ts

@@ -186,19 +186,22 @@ export function register(context: LanguageServiceContext) {
 			return [tsItem, tsRanges];
 		}
 
-		const map = context.documents.getSourceMap(virtualCode);
-
-		let range = map.getSourceRange(tsItem.range);
-		if (!range) {
-			// TODO: <script> range
-			range = {
-				start: map.sourceDocument.positionAt(0),
-				end: map.sourceDocument.positionAt(map.sourceDocument.getText().length),
-			};
-		}
+		for (const map of context.documents.getMaps(virtualCode)) {
+
+			let range = map.getSourceRange(tsItem.range);
+			if (!range) {
+				// TODO: <script> range
+				range = {
+					start: map.sourceDocument.positionAt(0),
+					end: map.sourceDocument.positionAt(map.sourceDocument.getText().length),
+				};
+			}
+
+			const selectionRange = map.getSourceRange(tsItem.selectionRange);
+			if (!selectionRange) {
+				continue;
+			}
 
-		const selectionRange = map.getSourceRange(tsItem.selectionRange);
-		if (selectionRange) {
 			const vueRanges = tsRanges.map(tsRange => map.getSourceRange(tsRange)).filter(notEmpty);
 			const vueItem: vscode.CallHierarchyItem = {
 				...tsItem,