From 725359dbd6e5fb68800c6b78cd8f9c64e26985ca Mon Sep 17 00:00:00 2001
From: Johnson Chu <johnsoncodehk@gmail.com>
Date: Tue, 4 Jun 2024 11:59:07 +0800
Subject: [PATCH 1/5] Revert "refactor: no longer support mapping to different
 source files"

This reverts commit c4b297b2426ec8b032eec16c5396263bc03e2163.
---
 .../lib/register/registerEditorFeatures.ts    |   5 +-
 packages/language-server/protocol.ts          |   1 -
 .../lib/features/provideCallHierarchyItems.ts |  27 +--
 .../lib/features/provideCompletionItems.ts    |  65 +++----
 .../lib/features/provideDefinition.ts         |  52 +++---
 .../lib/features/provideDiagnostics.ts        |  21 +--
 .../provideDocumentFormattingEdits.ts         |  47 +++--
 .../lib/features/provideFileReferences.ts     |  13 +-
 .../lib/features/provideReferences.ts         |  15 +-
 .../lib/features/provideWorkspaceSymbols.ts   |   9 +-
 .../lib/features/resolveCompletionItem.ts     |  19 +-
 .../language-service/lib/languageService.ts   |  22 +--
 packages/language-service/lib/types.ts        |   2 +-
 .../lib/utils/featureWorkers.ts               |  10 +-
 .../language-service/lib/utils/transform.ts   | 168 ++++++++++--------
 packages/typescript/lib/node/utils.ts         |   2 +-
 .../typescript/lib/protocol/createProject.ts  |  12 +-
 17 files changed, 263 insertions(+), 227 deletions(-)

diff --git a/packages/language-server/lib/register/registerEditorFeatures.ts b/packages/language-server/lib/register/registerEditorFeatures.ts
index 76a3521c..8e155e55 100644
--- a/packages/language-server/lib/register/registerEditorFeatures.ts
+++ b/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,
diff --git a/packages/language-server/protocol.ts b/packages/language-server/protocol.ts
index 3143f464..8c650cf5 100644
--- a/packages/language-server/protocol.ts
+++ b/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;
diff --git a/packages/language-service/lib/features/provideCallHierarchyItems.ts b/packages/language-service/lib/features/provideCallHierarchyItems.ts
index c3f5087c..eda7b6d6 100644
--- a/packages/language-service/lib/features/provideCallHierarchyItems.ts
+++ b/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,
diff --git a/packages/language-service/lib/features/provideCompletionItems.ts b/packages/language-service/lib/features/provideCompletionItems.ts
index a1280d5b..70849d0d 100644
--- a/packages/language-service/lib/features/provideCompletionItems.ts
+++ b/packages/language-service/lib/features/provideCompletionItems.ts
@@ -63,39 +63,40 @@ export function register(context: LanguageServiceContext) {
 						continue;
 					}
 
-					const map = context.documents.getSourceMap(virtualCode);
-
-					for (const mapped of map.getGeneratedPositions(position, data => isCompletionEnabled(data))) {
-
-						if (!cacheData.plugin.provideCompletionItems) {
-							continue;
-						}
-
-						cacheData.list = await cacheData.plugin.provideCompletionItems(map.embeddedDocument, mapped, completionContext, token);
-
-						if (!cacheData.list) {
-							continue;
+					for (const map of context.documents.getMaps(virtualCode)) {
+
+						for (const mapped of map.getGeneratedPositions(position, data => isCompletionEnabled(data))) {
+
+							if (!cacheData.plugin.provideCompletionItems) {
+								continue;
+							}
+
+							cacheData.list = await cacheData.plugin.provideCompletionItems(map.embeddedDocument, mapped, completionContext, token);
+
+							if (!cacheData.list) {
+								continue;
+							}
+
+							for (const item of cacheData.list.items) {
+								item.data = {
+									uri: uri.toString(),
+									original: {
+										additionalTextEdits: item.additionalTextEdits,
+										textEdit: item.textEdit,
+										data: item.data,
+									},
+									pluginIndex: pluginIndex,
+									embeddedDocumentUri: map.embeddedDocument.uri,
+								} satisfies ServiceCompletionData;
+							}
+
+							cacheData.list = transformCompletionList(
+								cacheData.list,
+								range => map.getSourceRange(range),
+								map.embeddedDocument,
+								context
+							);
 						}
-
-						for (const item of cacheData.list.items) {
-							item.data = {
-								uri: uri.toString(),
-								original: {
-									additionalTextEdits: item.additionalTextEdits,
-									textEdit: item.textEdit,
-									data: item.data,
-								},
-								pluginIndex: pluginIndex,
-								embeddedDocumentUri: map.embeddedDocument.uri,
-							} satisfies ServiceCompletionData;
-						}
-
-						cacheData.list = transformCompletionList(
-							cacheData.list,
-							range => map.getSourceRange(range),
-							map.embeddedDocument,
-							context
-						);
 					}
 				}
 				else {
diff --git a/packages/language-service/lib/features/provideDefinition.ts b/packages/language-service/lib/features/provideDefinition.ts
index 9bbb90b7..a4034308 100644
--- a/packages/language-service/lib/features/provideDefinition.ts
+++ b/packages/language-service/lib/features/provideDefinition.ts
@@ -104,15 +104,23 @@ export function register(
 					link.originSelectionRange = originSelectionRange;
 				}
 
+				let foundTargetSelectionRange = false;
+
 				const decoded = context.decodeEmbeddedDocumentUri(URI.parse(link.targetUri));
 				const sourceScript = decoded && context.language.scripts.get(decoded[0]);
 				const targetVirtualFile = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 				if (targetVirtualFile) {
 
-					const targetSourceMap = context.documents.getSourceMap(targetVirtualFile);
-					const targetSelectionRange = targetSourceMap.getSourceRange(link.targetSelectionRange);
-					if (targetSelectionRange) {
+					for (const targetSourceMap of context.documents.getMaps(targetVirtualFile)) {
+
+						const targetSelectionRange = targetSourceMap.getSourceRange(link.targetSelectionRange);
+						if (!targetSelectionRange) {
+							continue;
+						}
+
+						foundTargetSelectionRange = true;
+
 						let targetRange = targetSourceMap.getSourceRange(link.targetRange);
 
 						link.targetUri = targetSourceMap.sourceDocument.uri;
@@ -120,28 +128,26 @@ export function register(
 						link.targetRange = targetRange ?? targetSelectionRange;
 						link.targetSelectionRange = targetSelectionRange;
 					}
-					else if (apiName === 'provideDefinition') {
-						const targetMap = context.documents.getSourceMap(targetVirtualFile);
-						// cross file definition
-						if (targetMap.sourceDocument.uri !== uri.toString()) {
-							return {
-								...link,
-								targetUri: targetMap.sourceDocument.uri,
-								targetRange: {
-									start: { line: 0, character: 0 },
-									end: { line: 0, character: 0 },
-								},
-								targetSelectionRange: {
-									start: { line: 0, character: 0 },
-									end: { line: 0, character: 0 },
-								},
-							};
-						}
-						else {
-							return;
+
+					if (apiName === 'provideDefinition' && !foundTargetSelectionRange) {
+						for (const targetMap of context.documents.getMaps(targetVirtualFile)) {
+							if (targetMap && targetMap.sourceDocument.uri !== uri.toString()) {
+								return {
+									...link,
+									targetUri: targetMap.sourceDocument.uri,
+									targetRange: {
+										start: { line: 0, character: 0 },
+										end: { line: 0, character: 0 },
+									},
+									targetSelectionRange: {
+										start: { line: 0, character: 0 },
+										end: { line: 0, character: 0 },
+									},
+								};
+							}
 						}
+						return;
 					}
-
 				}
 
 				return link;
diff --git a/packages/language-service/lib/features/provideDiagnostics.ts b/packages/language-service/lib/features/provideDiagnostics.ts
index 53384129..55ee0c57 100644
--- a/packages/language-service/lib/features/provideDiagnostics.ts
+++ b/packages/language-service/lib/features/provideDiagnostics.ts
@@ -303,16 +303,17 @@ export function register(context: LanguageServiceContext) {
 					const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 					if (virtualCode) {
-						const map = context.documents.getSourceMap(virtualCode);
-						const range = map.getSourceRange(info.location.range, filter);
-						if (range) {
-							relatedInfos.push({
-								location: {
-									uri: map.sourceDocument.uri,
-									range,
-								},
-								message: info.message,
-							});
+						for (const map of context.documents.getMaps(virtualCode)) {
+							const range = map.getSourceRange(info.location.range, filter);
+							if (range) {
+								relatedInfos.push({
+									location: {
+										uri: map.sourceDocument.uri,
+										range,
+									},
+									message: info.message,
+								});
+							}
 						}
 					}
 					else {
diff --git a/packages/language-service/lib/features/provideDocumentFormattingEdits.ts b/packages/language-service/lib/features/provideDocumentFormattingEdits.ts
index dcf86d35..c602283d 100644
--- a/packages/language-service/lib/features/provideDocumentFormattingEdits.ts
+++ b/packages/language-service/lib/features/provideDocumentFormattingEdits.ts
@@ -1,4 +1,4 @@
-import { SourceMap, SourceScript, VirtualCode, forEachEmbeddedCode, isFormattingEnabled } from '@volar/language-core';
+import { CodeInformation, SourceMap, SourceScript, VirtualCode, forEachEmbeddedCode, isFormattingEnabled, updateVirtualCodeMapOfMap } from '@volar/language-core';
 import type * as ts from 'typescript';
 import type * as vscode from 'vscode-languageserver-protocol';
 import { TextDocument } from 'vscode-languageserver-textdocument';
@@ -8,6 +8,7 @@ import type { EmbeddedCodeFormattingOptions, LanguageServiceContext } from '../t
 import { NoneCancellationToken } from '../utils/cancellation';
 import { findOverlapCodeRange, stringToSnapshot } from '../utils/common';
 import { getEmbeddedFilesByLevel as getEmbeddedCodesByLevel } from '../utils/featureWorkers';
+import { createUriMap } from '../utils/uriMap';
 
 export function register(context: LanguageServiceContext) {
 
@@ -252,24 +253,32 @@ export function register(context: LanguageServiceContext) {
 	};
 
 	function createDocMap(virtualCode: VirtualCode, documentUri: URI, sourceLanguageId: string, _sourceSnapshot: ts.IScriptSnapshot) {
-		const map = new SourceMap(virtualCode.mappings);
-		const version = fakeVersion++;
-		return new SourceMapWithDocuments(
-			TextDocument.create(
-				documentUri.toString(),
-				sourceLanguageId,
-				version,
-				_sourceSnapshot.getText(0, _sourceSnapshot.getLength())
-			),
-			TextDocument.create(
-				context.encodeEmbeddedDocumentUri(documentUri, virtualCode.id).toString(),
-				virtualCode.languageId,
-				version,
-				virtualCode.snapshot.getText(0, virtualCode.snapshot.getLength())
-			),
-			map,
-			virtualCode,
-		);
+		const mapOfMap = createUriMap<[ts.IScriptSnapshot, SourceMap<CodeInformation>]>();
+		updateVirtualCodeMapOfMap(virtualCode, mapOfMap, sourceFileUri2 => {
+			if (!sourceFileUri2) {
+				return [documentUri, _sourceSnapshot];
+			}
+		});
+		if (mapOfMap.has(documentUri) && mapOfMap.get(documentUri)![0] === _sourceSnapshot) {
+			const map = mapOfMap.get(documentUri)!;
+			const version = fakeVersion++;
+			return new SourceMapWithDocuments(
+				TextDocument.create(
+					documentUri.toString(),
+					sourceLanguageId,
+					version,
+					_sourceSnapshot.getText(0, _sourceSnapshot.getLength())
+				),
+				TextDocument.create(
+					context.encodeEmbeddedDocumentUri(documentUri, virtualCode.id).toString(),
+					virtualCode.languageId,
+					version,
+					virtualCode.snapshot.getText(0, virtualCode.snapshot.getLength())
+				),
+				map[1],
+				virtualCode,
+			);
+		}
 	}
 }
 
diff --git a/packages/language-service/lib/features/provideFileReferences.ts b/packages/language-service/lib/features/provideFileReferences.ts
index 85f915e6..4eda2559 100644
--- a/packages/language-service/lib/features/provideFileReferences.ts
+++ b/packages/language-service/lib/features/provideFileReferences.ts
@@ -32,12 +32,13 @@ export function register(context: LanguageServiceContext) {
 						return reference;
 					}
 
-					const map = context.documents.getSourceMap(virtualCode);
-					const range = map.getSourceRange(reference.range, isReferencesEnabled);
-					if (range) {
-						reference.uri = map.sourceDocument.uri;
-						reference.range = range;
-						return reference;
+					for (const map of context.documents.getMaps(virtualCode)) {
+						const range = map.getSourceRange(reference.range, isReferencesEnabled);
+						if (range) {
+							reference.uri = map.sourceDocument.uri;
+							reference.range = range;
+							return reference;
+						}
 					}
 				})
 				.filter(notEmpty),
diff --git a/packages/language-service/lib/features/provideReferences.ts b/packages/language-service/lib/features/provideReferences.ts
index 4b4baa2f..97823faa 100644
--- a/packages/language-service/lib/features/provideReferences.ts
+++ b/packages/language-service/lib/features/provideReferences.ts
@@ -87,13 +87,14 @@ export function register(context: LanguageServiceContext) {
 					const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 					if (virtualCode) {
-						const map = context.documents.getSourceMap(virtualCode);
-						const range = map.getSourceRange(reference.range, isReferencesEnabled);
-						if (range) {
-							results.push({
-								uri: map.sourceDocument.uri,
-								range,
-							});
+						for (const map of context.documents.getMaps(virtualCode)) {
+							const range = map.getSourceRange(reference.range, isReferencesEnabled);
+							if (range) {
+								results.push({
+									uri: map.sourceDocument.uri,
+									range,
+								});
+							}
 						}
 					}
 					else {
diff --git a/packages/language-service/lib/features/provideWorkspaceSymbols.ts b/packages/language-service/lib/features/provideWorkspaceSymbols.ts
index 21e394f2..0baacd69 100644
--- a/packages/language-service/lib/features/provideWorkspaceSymbols.ts
+++ b/packages/language-service/lib/features/provideWorkspaceSymbols.ts
@@ -32,10 +32,11 @@ export function register(context: LanguageServiceContext) {
 				const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 				if (virtualCode) {
-					const map = context.documents.getSourceMap(virtualCode);
-					const range = map.getSourceRange(loc.range);
-					if (range) {
-						return { uri: map.sourceDocument.uri, range };
+					for (const map of context.documents.getMaps(virtualCode)) {
+						const range = map.getSourceRange(loc.range);
+						if (range) {
+							return { uri: map.sourceDocument.uri, range };
+						}
 					}
 				}
 				else {
diff --git a/packages/language-service/lib/features/resolveCompletionItem.ts b/packages/language-service/lib/features/resolveCompletionItem.ts
index 7d2575a4..502e2b57 100644
--- a/packages/language-service/lib/features/resolveCompletionItem.ts
+++ b/packages/language-service/lib/features/resolveCompletionItem.ts
@@ -28,14 +28,17 @@ export function register(context: LanguageServiceContext) {
 				const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 				if (virtualCode) {
-					const map = context.documents.getSourceMap(virtualCode);
-					item = await plugin[1].resolveCompletionItem(item, token);
-					item = plugin[1].transformCompletionItem?.(item) ?? transformCompletionItem(
-						item,
-						embeddedRange => map.getSourceRange(embeddedRange),
-						map.embeddedDocument,
-						context
-					);
+
+					for (const map of context.documents.getMaps(virtualCode)) {
+
+						item = await plugin[1].resolveCompletionItem(item, token);
+						item = plugin[1].transformCompletionItem?.(item) ?? transformCompletionItem(
+							item,
+							embeddedRange => map.getSourceRange(embeddedRange),
+							map.embeddedDocument,
+							context
+						);
+					}
 				}
 			}
 			else {
diff --git a/packages/language-service/lib/languageService.ts b/packages/language-service/lib/languageService.ts
index d44cbf26..90767bfa 100644
--- a/packages/language-service/lib/languageService.ts
+++ b/packages/language-service/lib/languageService.ts
@@ -72,23 +72,19 @@ export function createLanguageService(
 				}
 				return map.get(uri)!;
 			},
-			getSourceMap(virtualCode) {
-				const map = context.language.maps.get(virtualCode);
-				let result = map2DocMap.get(map);
-				if (!result) {
-					const sourceScript = context.language.scripts.fromVirtualCode(virtualCode);
-					const embeddedUri = context.encodeEmbeddedDocumentUri(sourceScript.id, virtualCode.id);
-					map2DocMap.set(
-						map,
-						result = new SourceMapWithDocuments(
-							this.get(sourceScript.id, sourceScript.languageId, sourceScript.snapshot),
+			*getMaps(virtualCode) {
+				for (const [uri, [snapshot, map]] of context.language.maps.forEach(virtualCode)) {
+					if (!map2DocMap.has(map)) {
+						const embeddedUri = context.encodeEmbeddedDocumentUri(uri, virtualCode.id);
+						map2DocMap.set(map, new SourceMapWithDocuments(
+							this.get(uri, context.language.scripts.get(uri)!.languageId, snapshot),
 							this.get(embeddedUri, virtualCode.languageId, virtualCode.snapshot),
 							map,
 							virtualCode,
-						)
-					);
+						));
+					}
+					yield map2DocMap.get(map)!;
 				}
-				return result;
 			},
 			getLinkedCodeMap(virtualCode, documentUri) {
 				const map = context.language.linkedCodeMaps.get(virtualCode);
diff --git a/packages/language-service/lib/types.ts b/packages/language-service/lib/types.ts
index 746876ac..3f3bdc08 100644
--- a/packages/language-service/lib/types.ts
+++ b/packages/language-service/lib/types.ts
@@ -59,7 +59,7 @@ export interface LanguageServiceContext {
 	};
 	documents: {
 		get(uri: URI, languageId: string, snapshot: ts.IScriptSnapshot): TextDocument;
-		getSourceMap(virtualCode: VirtualCode): SourceMapWithDocuments;
+		getMaps(virtualCode: VirtualCode): Generator<SourceMapWithDocuments>;
 		getLinkedCodeMap(virtualCode: VirtualCode, documentUri: URI): LinkedCodeMapWithDocument | undefined;
 	};
 	plugins: [LanguageServicePlugin, LanguageServicePluginInstance][];
diff --git a/packages/language-service/lib/utils/featureWorkers.ts b/packages/language-service/lib/utils/featureWorkers.ts
index aa0bc27a..9cae6a9d 100644
--- a/packages/language-service/lib/utils/featureWorkers.ts
+++ b/packages/language-service/lib/utils/featureWorkers.ts
@@ -144,9 +144,13 @@ export function* forEachEmbeddedDocument(
 		}
 	}
 
-	const map = context.documents.getSourceMap(current);
-	if (!context.disabledEmbeddedDocumentUris.get(context.encodeEmbeddedDocumentUri(sourceScriptId, current.id))) {
-		yield map;
+	for (const map of context.documents.getMaps(current)) {
+		if (
+			sourceScriptId.toString() === map.sourceDocument.uri
+			&& !context.disabledEmbeddedDocumentUris.get(context.encodeEmbeddedDocumentUri(sourceScriptId, current.id))
+		) {
+			yield map;
+		}
 	}
 }
 
diff --git a/packages/language-service/lib/utils/transform.ts b/packages/language-service/lib/utils/transform.ts
index 4c89f043..cb8244ad 100644
--- a/packages/language-service/lib/utils/transform.ts
+++ b/packages/language-service/lib/utils/transform.ts
@@ -19,8 +19,11 @@ export function transformDocumentLinkTarget(_target: string, context: LanguageSe
 	target = decoded[0];
 
 	if (embeddedRange && virtualCode) {
-		const map = context.documents.getSourceMap(virtualCode);
-		if (map.map.mappings.some(mapping => isDocumentLinkEnabled(mapping.data))) {
+		for (const map of context.documents.getMaps(virtualCode)) {
+			if (!map.map.mappings.some(mapping => isDocumentLinkEnabled(mapping.data))) {
+				continue;
+			}
+
 			const startLine = Number(embeddedRange[1]) - 1;
 			const startCharacter = Number(embeddedRange[3] ?? 1) - 1;
 			if (embeddedRange[5] !== undefined) {
@@ -35,6 +38,7 @@ export function transformDocumentLinkTarget(_target: string, context: LanguageSe
 						fragment: 'L' + (sourceRange.start.line + 1) + ',' + (sourceRange.start.character + 1)
 							+ '-L' + (sourceRange.end.line + 1) + ',' + (sourceRange.end.character + 1),
 					});
+					break;
 				}
 			}
 			else {
@@ -43,6 +47,7 @@ export function transformDocumentLinkTarget(_target: string, context: LanguageSe
 					target = target.with({
 						fragment: 'L' + (sourcePos.line + 1) + ',' + (sourcePos.character + 1),
 					});
+					break;
 				}
 			}
 		}
@@ -324,10 +329,11 @@ export function transformWorkspaceEdit(
 		const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 		if (virtualCode) {
-			const map = context.documents.getSourceMap(virtualCode);
-			// TODO: check capability?
-			const uri = map.sourceDocument.uri;
-			sourceResult.changeAnnotations[uri] = tsAnno;
+			for (const map of context.documents.getMaps(virtualCode)) {
+				// TODO: check capability?
+				const uri = map.sourceDocument.uri;
+				sourceResult.changeAnnotations[uri] = tsAnno;
+			}
 		}
 		else {
 			sourceResult.changeAnnotations[tsUri] = tsAnno;
@@ -342,33 +348,34 @@ export function transformWorkspaceEdit(
 		const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 		if (virtualCode) {
-			const map = context.documents.getSourceMap(virtualCode);
-			const tsEdits = edit.changes[tsUri];
-			for (const tsEdit of tsEdits) {
-				if (mode === 'rename' || mode === 'fileName' || mode === 'codeAction') {
+			for (const map of context.documents.getMaps(virtualCode)) {
+				const tsEdits = edit.changes[tsUri];
+				for (const tsEdit of tsEdits) {
+					if (mode === 'rename' || mode === 'fileName' || mode === 'codeAction') {
 
-					let _data!: CodeInformation;
-
-					const range = map.getSourceRange(tsEdit.range, data => {
-						_data = data;
-						return isRenameEnabled(data);
-					});
+						let _data!: CodeInformation;
 
-					if (range) {
-						sourceResult.changes[map.sourceDocument.uri] ??= [];
-						sourceResult.changes[map.sourceDocument.uri].push({
-							newText: resolveRenameEditText(tsEdit.newText, _data),
-							range,
+						const range = map.getSourceRange(tsEdit.range, data => {
+							_data = data;
+							return isRenameEnabled(data);
 						});
-						hasResult = true;
+
+						if (range) {
+							sourceResult.changes[map.sourceDocument.uri] ??= [];
+							sourceResult.changes[map.sourceDocument.uri].push({
+								newText: resolveRenameEditText(tsEdit.newText, _data),
+								range,
+							});
+							hasResult = true;
+						}
 					}
-				}
-				else {
-					const range = map.getSourceRange(tsEdit.range);
-					if (range) {
-						sourceResult.changes[map.sourceDocument.uri] ??= [];
-						sourceResult.changes[map.sourceDocument.uri].push({ newText: tsEdit.newText, range });
-						hasResult = true;
+					else {
+						const range = map.getSourceRange(tsEdit.range);
+						if (range) {
+							sourceResult.changes[map.sourceDocument.uri] ??= [];
+							sourceResult.changes[map.sourceDocument.uri].push({ newText: tsEdit.newText, range });
+							hasResult = true;
+						}
 					}
 				}
 			}
@@ -391,43 +398,44 @@ export function transformWorkspaceEdit(
 				const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 				if (virtualCode) {
-					const map = context.documents.getSourceMap(virtualCode);
-					sourceEdit = {
-						textDocument: {
-							uri: map.sourceDocument.uri,
-							version: versions[map.sourceDocument.uri] ?? null,
-						},
-						edits: [],
-					} satisfies vscode.TextDocumentEdit;
-					for (const tsEdit of tsDocEdit.edits) {
-						if (mode === 'rename' || mode === 'fileName' || mode === 'codeAction') {
-							let _data!: CodeInformation;
-							const range = map.getSourceRange(tsEdit.range, data => {
-								_data = data;
-								// fix https://github.com/johnsoncodehk/volar/issues/1091
-								return isRenameEnabled(data);
-							});
-							if (range) {
-								sourceEdit.edits.push({
-									annotationId: 'annotationId' in tsEdit ? tsEdit.annotationId : undefined,
-									newText: resolveRenameEditText(tsEdit.newText, _data),
-									range,
+					for (const map of context.documents.getMaps(virtualCode)) {
+						sourceEdit = {
+							textDocument: {
+								uri: map.sourceDocument.uri,
+								version: versions[map.sourceDocument.uri] ?? null,
+							},
+							edits: [],
+						} satisfies vscode.TextDocumentEdit;
+						for (const tsEdit of tsDocEdit.edits) {
+							if (mode === 'rename' || mode === 'fileName' || mode === 'codeAction') {
+								let _data!: CodeInformation;
+								const range = map.getSourceRange(tsEdit.range, data => {
+									_data = data;
+									// fix https://github.com/johnsoncodehk/volar/issues/1091
+									return isRenameEnabled(data);
 								});
+								if (range) {
+									sourceEdit.edits.push({
+										annotationId: 'annotationId' in tsEdit ? tsEdit.annotationId : undefined,
+										newText: resolveRenameEditText(tsEdit.newText, _data),
+										range,
+									});
+								}
 							}
-						}
-						else {
-							const range = map.getSourceRange(tsEdit.range);
-							if (range) {
-								sourceEdit.edits.push({
-									annotationId: 'annotationId' in tsEdit ? tsEdit.annotationId : undefined,
-									newText: tsEdit.newText,
-									range,
-								});
+							else {
+								const range = map.getSourceRange(tsEdit.range);
+								if (range) {
+									sourceEdit.edits.push({
+										annotationId: 'annotationId' in tsEdit ? tsEdit.annotationId : undefined,
+										newText: tsEdit.newText,
+										range,
+									});
+								}
 							}
 						}
-					}
-					if (!sourceEdit.edits.length) {
-						sourceEdit = undefined;
+						if (!sourceEdit.edits.length) {
+							sourceEdit = undefined;
+						}
 					}
 				}
 				else {
@@ -444,15 +452,16 @@ export function transformWorkspaceEdit(
 				const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 				if (virtualCode) {
-					const map = context.documents.getSourceMap(virtualCode);
-					// TODO: check capability?
-					sourceEdit = {
-						kind: 'rename',
-						oldUri: map.sourceDocument.uri,
-						newUri: tsDocEdit.newUri /* TODO: remove .ts? */,
-						options: tsDocEdit.options,
-						annotationId: tsDocEdit.annotationId,
-					} satisfies vscode.RenameFile;
+					for (const map of context.documents.getMaps(virtualCode)) {
+						// TODO: check capability?
+						sourceEdit = {
+							kind: 'rename',
+							oldUri: map.sourceDocument.uri,
+							newUri: tsDocEdit.newUri /* TODO: remove .ts? */,
+							options: tsDocEdit.options,
+							annotationId: tsDocEdit.annotationId,
+						} satisfies vscode.RenameFile;
+					}
 				}
 				else {
 					sourceEdit = tsDocEdit;
@@ -465,14 +474,15 @@ export function transformWorkspaceEdit(
 				const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
 
 				if (virtualCode) {
-					const map = context.documents.getSourceMap(virtualCode);
-					// TODO: check capability?
-					sourceEdit = {
-						kind: 'delete',
-						uri: map.sourceDocument.uri,
-						options: tsDocEdit.options,
-						annotationId: tsDocEdit.annotationId,
-					} satisfies vscode.DeleteFile;
+					for (const map of context.documents.getMaps(virtualCode)) {
+						// TODO: check capability?
+						sourceEdit = {
+							kind: 'delete',
+							uri: map.sourceDocument.uri,
+							options: tsDocEdit.options,
+							annotationId: tsDocEdit.annotationId,
+						} satisfies vscode.DeleteFile;
+					}
 				}
 				else {
 					sourceEdit = tsDocEdit;
diff --git a/packages/typescript/lib/node/utils.ts b/packages/typescript/lib/node/utils.ts
index 101385d5..188bf2ef 100644
--- a/packages/typescript/lib/node/utils.ts
+++ b/packages/typescript/lib/node/utils.ts
@@ -9,7 +9,7 @@ export function getServiceScript(language: Language<string>, fileName: string) {
 	if (sourceScript?.generated) {
 		const serviceScript = sourceScript.generated.languagePlugin.typescript?.getServiceScript(sourceScript.generated.root);
 		if (serviceScript) {
-			const map = language.maps.get(serviceScript.code);
+			const map = language.maps.get(serviceScript.code, sourceScript.id);
 			if (map) {
 				return [serviceScript, sourceScript, map] as const;
 			}
diff --git a/packages/typescript/lib/protocol/createProject.ts b/packages/typescript/lib/protocol/createProject.ts
index fa82f7a8..387063e5 100644
--- a/packages/typescript/lib/protocol/createProject.ts
+++ b/packages/typescript/lib/protocol/createProject.ts
@@ -1,4 +1,4 @@
-import { FileMap, Language, TypeScriptExtraServiceScript, forEachEmbeddedCode } from '@volar/language-core';
+import { FileMap, Language, SourceScript, TypeScriptExtraServiceScript, forEachEmbeddedCode } from '@volar/language-core';
 import * as path from 'path-browserify';
 import type * as ts from 'typescript';
 import { createResolveModuleName } from '../resolveModuleName';
@@ -27,7 +27,7 @@ export function createLanguageServiceHost<T>(
 	let lastProjectVersion: number | string | undefined;
 	let tsProjectVersion = 0;
 	let tsFileRegistry = new FileMap<boolean>(sys.useCaseSensitiveFileNames);
-	let extraScriptRegistry = new FileMap<TypeScriptExtraServiceScript>(sys.useCaseSensitiveFileNames);
+	let extraScriptRegistry = new FileMap<[SourceScript<T>, TypeScriptExtraServiceScript]>(sys.useCaseSensitiveFileNames);
 	let lastTsVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 	let lastOtherVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 	let languageServiceHost: ts.LanguageServiceHost = {
@@ -99,7 +99,7 @@ export function createLanguageServiceHost<T>(
 			sync();
 
 			if (extraScriptRegistry.has(fileName)) {
-				return extraScriptRegistry.get(fileName)!.scriptKind;
+				return extraScriptRegistry.get(fileName)![1].scriptKind;
 			}
 
 			const sourceScript = language.scripts.get(asScrpitId(fileName));
@@ -218,7 +218,7 @@ export function createLanguageServiceHost<T>(
 				for (const extraServiceScript of sourceScript.generated.languagePlugin.typescript?.getExtraServiceScripts?.(fileName, sourceScript.generated.root) ?? []) {
 					newTsVirtualFileSnapshots.add(extraServiceScript.code.snapshot);
 					tsFileNamesSet.add(extraServiceScript.fileName);
-					extraScriptRegistry.set(extraServiceScript.fileName, extraServiceScript);
+					extraScriptRegistry.set(extraServiceScript.fileName, [sourceScript, extraServiceScript]);
 				}
 				for (const code of forEachEmbeddedCode(sourceScript.generated.root)) {
 					newOtherVirtualFileSnapshots.add(code.snapshot);
@@ -251,7 +251,7 @@ export function createLanguageServiceHost<T>(
 		sync();
 
 		if (extraScriptRegistry.has(fileName)) {
-			return extraScriptRegistry.get(fileName)!.code.snapshot;
+			return extraScriptRegistry.get(fileName)![1].code.snapshot;
 		}
 
 		const sourceScript = language.scripts.get(asScrpitId(fileName));
@@ -278,7 +278,7 @@ export function createLanguageServiceHost<T>(
 		const version = scriptVersions.get(fileName)!;
 
 		if (extraScriptRegistry.has(fileName)) {
-			const snapshot = extraScriptRegistry.get(fileName)!.code.snapshot;
+			const snapshot = extraScriptRegistry.get(fileName)![1].code.snapshot;
 			if (!version.map.has(snapshot)) {
 				version.map.set(snapshot, version.lastVersion++);
 			}

From 328eb6c664cf4b3076a30c58e47bf7843737b6b8 Mon Sep 17 00:00:00 2001
From: Johnson Chu <johnsoncodehk@gmail.com>
Date: Tue, 4 Jun 2024 14:39:01 +0800
Subject: [PATCH 2/5] wip

---
 packages/eslint/index.ts                      |   2 +-
 packages/language-core/index.ts               | 107 ++++++++++++++----
 packages/language-core/lib/linkedCodeMap.ts   |   2 +-
 packages/language-core/lib/types.ts           |  47 ++++----
 .../provideDocumentFormattingEdits.ts         |  47 ++++----
 .../language-service/lib/languageService.ts   |   2 +-
 packages/language-service/lib/types.ts        |   4 +-
 packages/language-service/lib/utils/common.ts |   6 +-
 .../lib/utils/featureWorkers.ts               |   2 +-
 packages/source-map/lib/buildMappings.ts      |   1 -
 packages/source-map/lib/sourceMap.ts          |  23 ++--
 packages/test-utils/index.ts                  |   6 +-
 packages/typescript/lib/node/transform.ts     |   4 +-
 packages/typescript/lib/node/utils.ts         |   2 +-
 .../typescript/lib/protocol/createProject.ts  |  12 +-
 15 files changed, 161 insertions(+), 106 deletions(-)

diff --git a/packages/eslint/index.ts b/packages/eslint/index.ts
index a08d56ad..dcdbb632 100644
--- a/packages/eslint/index.ts
+++ b/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,
diff --git a/packages/language-core/index.ts b/packages/language-core/index.ts
index 91ba8f36..abb5946b 100644
--- a/packages/language-core/index.ts
+++ b/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.isRelatedDirty || 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,
+						relateds: new Set(),
+						targets: 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.relatedMappings) {
+					for (const [relatedScriptId, relatedMappings] of virtualCode.relatedMappings) {
+						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.targets.forEach(id => {
+			const sourceScript = scriptRegistry.get(id);
+			if (sourceScript) {
+				sourceScript.isRelatedDirty = true;
+			}
+		});
+	}
+
+	function prepareCreateVirtualCode(sourceScript: SourceScript<T>): CodegenContext<T> {
+		for (const id of sourceScript.relateds) {
+			scriptRegistry.get(id)?.targets.delete(sourceScript.id);
+		}
+		sourceScript.relateds.clear();
+		sourceScript.isRelatedDirty = false;
+		return {
+			getRelatedSourceScript(id) {
+				sync(id);
+				const relatedSourceScript = scriptRegistry.get(id);
+				if (relatedSourceScript) {
+					relatedSourceScript.targets.add(sourceScript.id);
+					sourceScript.relateds.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) {
diff --git a/packages/language-core/lib/linkedCodeMap.ts b/packages/language-core/lib/linkedCodeMap.ts
index b0e85d6b..eac2bf49 100644
--- a/packages/language-core/lib/linkedCodeMap.ts
+++ b/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];
diff --git a/packages/language-core/lib/types.ts b/packages/language-core/lib/types.ts
index c849755d..e2ddf6de 100644
--- a/packages/language-core/lib/types.ts
+++ b/packages/language-core/lib/types.ts
@@ -8,13 +8,14 @@ export interface Language<T> {
 		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,7 +24,7 @@ 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;
 	};
@@ -33,21 +34,25 @@ export interface SourceScript<T> {
 	id: T;
 	languageId: string;
 	snapshot: ts.IScriptSnapshot;
+	relateds: Set<T>;
+	targets: Set<T>;
+	isRelatedDirty?: 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[];
+	relatedMappings?: Map<T, CodeMapping[]>;
+	embeddedCodes?: VirtualCode<T>[];
 	linkedCodeMappings?: Mapping[];
 }
 
@@ -77,19 +82,19 @@ 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> extends TypeScriptServiceScript<T> {
 	fileName: string;
 }
 
-export interface LanguagePlugin<T, K extends VirtualCode = VirtualCode> {
+export interface LanguagePlugin<T, 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.
 	 */
@@ -97,31 +102,35 @@ export interface LanguagePlugin<T, K extends VirtualCode = VirtualCode> {
 	/**
 	 * 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> {
+	getRelatedSourceScript(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;
 }
diff --git a/packages/language-service/lib/features/provideDocumentFormattingEdits.ts b/packages/language-service/lib/features/provideDocumentFormattingEdits.ts
index c602283d..dcf86d35 100644
--- a/packages/language-service/lib/features/provideDocumentFormattingEdits.ts
+++ b/packages/language-service/lib/features/provideDocumentFormattingEdits.ts
@@ -1,4 +1,4 @@
-import { CodeInformation, SourceMap, SourceScript, VirtualCode, forEachEmbeddedCode, isFormattingEnabled, updateVirtualCodeMapOfMap } from '@volar/language-core';
+import { SourceMap, SourceScript, VirtualCode, forEachEmbeddedCode, isFormattingEnabled } from '@volar/language-core';
 import type * as ts from 'typescript';
 import type * as vscode from 'vscode-languageserver-protocol';
 import { TextDocument } from 'vscode-languageserver-textdocument';
@@ -8,7 +8,6 @@ import type { EmbeddedCodeFormattingOptions, LanguageServiceContext } from '../t
 import { NoneCancellationToken } from '../utils/cancellation';
 import { findOverlapCodeRange, stringToSnapshot } from '../utils/common';
 import { getEmbeddedFilesByLevel as getEmbeddedCodesByLevel } from '../utils/featureWorkers';
-import { createUriMap } from '../utils/uriMap';
 
 export function register(context: LanguageServiceContext) {
 
@@ -253,32 +252,24 @@ export function register(context: LanguageServiceContext) {
 	};
 
 	function createDocMap(virtualCode: VirtualCode, documentUri: URI, sourceLanguageId: string, _sourceSnapshot: ts.IScriptSnapshot) {
-		const mapOfMap = createUriMap<[ts.IScriptSnapshot, SourceMap<CodeInformation>]>();
-		updateVirtualCodeMapOfMap(virtualCode, mapOfMap, sourceFileUri2 => {
-			if (!sourceFileUri2) {
-				return [documentUri, _sourceSnapshot];
-			}
-		});
-		if (mapOfMap.has(documentUri) && mapOfMap.get(documentUri)![0] === _sourceSnapshot) {
-			const map = mapOfMap.get(documentUri)!;
-			const version = fakeVersion++;
-			return new SourceMapWithDocuments(
-				TextDocument.create(
-					documentUri.toString(),
-					sourceLanguageId,
-					version,
-					_sourceSnapshot.getText(0, _sourceSnapshot.getLength())
-				),
-				TextDocument.create(
-					context.encodeEmbeddedDocumentUri(documentUri, virtualCode.id).toString(),
-					virtualCode.languageId,
-					version,
-					virtualCode.snapshot.getText(0, virtualCode.snapshot.getLength())
-				),
-				map[1],
-				virtualCode,
-			);
-		}
+		const map = new SourceMap(virtualCode.mappings);
+		const version = fakeVersion++;
+		return new SourceMapWithDocuments(
+			TextDocument.create(
+				documentUri.toString(),
+				sourceLanguageId,
+				version,
+				_sourceSnapshot.getText(0, _sourceSnapshot.getLength())
+			),
+			TextDocument.create(
+				context.encodeEmbeddedDocumentUri(documentUri, virtualCode.id).toString(),
+				virtualCode.languageId,
+				version,
+				virtualCode.snapshot.getText(0, virtualCode.snapshot.getLength())
+			),
+			map,
+			virtualCode,
+		);
 	}
 }
 
diff --git a/packages/language-service/lib/languageService.ts b/packages/language-service/lib/languageService.ts
index 90767bfa..3875a2b4 100644
--- a/packages/language-service/lib/languageService.ts
+++ b/packages/language-service/lib/languageService.ts
@@ -73,7 +73,7 @@ export function createLanguageService(
 				return map.get(uri)!;
 			},
 			*getMaps(virtualCode) {
-				for (const [uri, [snapshot, map]] of context.language.maps.forEach(virtualCode)) {
+				for (const [uri, snapshot, map] of context.language.maps.forEach(virtualCode)) {
 					if (!map2DocMap.has(map)) {
 						const embeddedUri = context.encodeEmbeddedDocumentUri(uri, virtualCode.id);
 						map2DocMap.set(map, new SourceMapWithDocuments(
diff --git a/packages/language-service/lib/types.ts b/packages/language-service/lib/types.ts
index 3f3bdc08..995ebf17 100644
--- a/packages/language-service/lib/types.ts
+++ b/packages/language-service/lib/types.ts
@@ -59,8 +59,8 @@ export interface LanguageServiceContext {
 	};
 	documents: {
 		get(uri: URI, languageId: string, snapshot: ts.IScriptSnapshot): TextDocument;
-		getMaps(virtualCode: VirtualCode): Generator<SourceMapWithDocuments>;
-		getLinkedCodeMap(virtualCode: VirtualCode, documentUri: URI): LinkedCodeMapWithDocument | undefined;
+		getMaps(virtualCode: VirtualCode<URI>): Generator<SourceMapWithDocuments>;
+		getLinkedCodeMap(virtualCode: VirtualCode<URI>, documentUri: URI): LinkedCodeMapWithDocument | undefined;
 	};
 	plugins: [LanguageServicePlugin, LanguageServicePluginInstance][];
 	disabledEmbeddedDocumentUris: UriMap<boolean>;
diff --git a/packages/language-service/lib/utils/common.ts b/packages/language-service/lib/utils/common.ts
index 083ccd8f..ce8e7365 100644
--- a/packages/language-service/lib/utils/common.ts
+++ b/packages/language-service/lib/utils/common.ts
@@ -31,14 +31,14 @@ export function findOverlapCodeRange(
 				const mappingEnd = mapping.sourceOffsets[mapping.sourceOffsets.length - 1] + mapping.lengths[mapping.lengths.length - 1];
 				const overlap = getOverlapRange(start, end, mappingStart, mappingEnd);
 				if (overlap) {
-					const curMappedStart = (overlap.start - mappingStart) + mapping.generatedOffsets[0]
+					const curMappedStart = (overlap.start - mappingStart) + mapping.generatedOffsets[0];
 
 					mappedStart = mappedStart === undefined ? curMappedStart : Math.min(mappedStart, curMappedStart);
 
 					const lastGeneratedLength = (mapping.generatedLengths ?? mapping.lengths)[mapping.generatedOffsets.length - 1];
-					const curMappedEndOffset = Math.min(overlap.end - mapping.sourceOffsets[mapping.sourceOffsets.length - 1], lastGeneratedLength)
+					const curMappedEndOffset = Math.min(overlap.end - mapping.sourceOffsets[mapping.sourceOffsets.length - 1], lastGeneratedLength);
 
-					const curMappedEnd = mapping.generatedOffsets[mapping.generatedOffsets.length - 1] + curMappedEndOffset
+					const curMappedEnd = mapping.generatedOffsets[mapping.generatedOffsets.length - 1] + curMappedEndOffset;
 
 					mappedEnd = mappedEnd === undefined ? curMappedEnd : Math.max(mappedEnd, curMappedEnd);
 				}
diff --git a/packages/language-service/lib/utils/featureWorkers.ts b/packages/language-service/lib/utils/featureWorkers.ts
index 9cae6a9d..306ad732 100644
--- a/packages/language-service/lib/utils/featureWorkers.ts
+++ b/packages/language-service/lib/utils/featureWorkers.ts
@@ -135,7 +135,7 @@ export async function safeCall<T>(cb: () => Thenable<T> | T, errorMsg?: string)
 export function* forEachEmbeddedDocument(
 	context: LanguageServiceContext,
 	sourceScriptId: URI,
-	current: VirtualCode
+	current: VirtualCode<URI>
 ): Generator<SourceMapWithDocuments> {
 
 	if (current.embeddedCodes) {
diff --git a/packages/source-map/lib/buildMappings.ts b/packages/source-map/lib/buildMappings.ts
index 57b64bf7..224b0251 100644
--- a/packages/source-map/lib/buildMappings.ts
+++ b/packages/source-map/lib/buildMappings.ts
@@ -10,7 +10,6 @@ export function buildMappings<T>(chunks: Segment<T>[]) {
 		}
 		else {
 			mappings.push({
-				source: segment[1],
 				sourceOffsets: [segment[2]],
 				generatedOffsets: [length],
 				lengths: [segment[0].length],
diff --git a/packages/source-map/lib/sourceMap.ts b/packages/source-map/lib/sourceMap.ts
index 199e80a1..ef4d6d31 100644
--- a/packages/source-map/lib/sourceMap.ts
+++ b/packages/source-map/lib/sourceMap.ts
@@ -3,24 +3,23 @@ import { translateOffset } from './translateOffset';
 
 export type CodeRangeKey = 'sourceOffsets' | 'generatedOffsets';
 
-export interface Mapping<T = any> {
-	source?: string;
+export interface Mapping<Data = unknown> {
 	sourceOffsets: number[];
 	generatedOffsets: number[];
 	lengths: number[];
 	generatedLengths?: number[];
-	data: T;
+	data: Data;
 }
 
-interface MappingMemo<Data> {
+interface MappingMemo {
 	offsets: number[];
-	mappings: Set<Mapping<Data>>[];
+	mappings: Set<Mapping>[];
 }
 
-export class SourceMap<Data = any> {
+export class SourceMap<Data = unknown> {
 
-	private sourceCodeOffsetsMemo: MappingMemo<Data> | undefined;
-	private generatedCodeOffsetsMemo: MappingMemo<Data> | undefined;
+	private sourceCodeOffsetsMemo: MappingMemo | undefined;
+	private generatedCodeOffsetsMemo: MappingMemo | undefined;
 
 	constructor(public readonly mappings: Mapping<Data>[]) { }
 
@@ -39,7 +38,7 @@ export class SourceMap<Data = any> {
 		}
 
 		const { low: start, high: end } = binarySearch(memo.offsets, offset);
-		const skip = new Set<Mapping<Data>>();
+		const skip = new Set<Mapping>();
 
 		for (let i = start; i <= end; i++) {
 			for (const mapping of memo.mappings[i]) {
@@ -50,7 +49,7 @@ export class SourceMap<Data = any> {
 
 				const mapped = translateOffset(offset, mapping[fromRange], mapping[toRange], getLengths(mapping, fromRange), getLengths(mapping, toRange));
 				if (mapped !== undefined) {
-					yield [mapped, mapping] as const;
+					yield [mapped, mapping as Mapping<Data>] as const;
 				}
 			}
 		}
@@ -62,7 +61,7 @@ export class SourceMap<Data = any> {
 			: this.generatedCodeOffsetsMemo ??= this.createMemo('generatedOffsets');
 	}
 
-	private createMemo(key: CodeRangeKey): MappingMemo<Data> {
+	private createMemo(key: CodeRangeKey): MappingMemo {
 		const offsetsSet = new Set<number>();
 		for (const mapping of this.mappings) {
 			for (let i = 0; i < mapping[key].length; i++) {
@@ -88,6 +87,6 @@ export class SourceMap<Data = any> {
 	}
 }
 
-function getLengths<Data>(mapping: Mapping<Data>, key: CodeRangeKey) {
+function getLengths(mapping: Mapping, key: CodeRangeKey) {
 	return key == "sourceOffsets" ? mapping.lengths : mapping.generatedLengths ?? mapping.lengths;
 }
diff --git a/packages/test-utils/index.ts b/packages/test-utils/index.ts
index bab69489..abccba27 100644
--- a/packages/test-utils/index.ts
+++ b/packages/test-utils/index.ts
@@ -500,8 +500,7 @@ export function* printSnapshot(
 					' '.repeat(log.lineOffset),
 					sourceLineHead,
 					'(exact match)',
-					`(${log.mapping.source
-					+ ':' + (sourcePosition.line + 1)
+					`(${':' + (sourcePosition.line + 1)
 					+ ':' + (sourcePosition.character + 1)
 					})`,
 				].join(' ');
@@ -512,8 +511,7 @@ export function* printSnapshot(
 					' '.repeat(log.lineOffset),
 					sourceLineHead,
 					normalizeLogText(sourceLine),
-					`(${log.mapping.source
-					+ ':' + (sourcePosition.line + 1)
+					`(${':' + (sourcePosition.line + 1)
 					+ ':' + (sourcePosition.character + 1)})`,
 				].join(' ');
 				yield [
diff --git a/packages/typescript/lib/node/transform.ts b/packages/typescript/lib/node/transform.ts
index 5347410b..33badb0d 100644
--- a/packages/typescript/lib/node/transform.ts
+++ b/packages/typescript/lib/node/transform.ts
@@ -184,7 +184,7 @@ export function transformTextSpan(
 export function toSourceOffset(
 	serviceScript: TypeScriptServiceScript,
 	sourceScript: SourceScript<string>,
-	map: SourceMap,
+	map: SourceMap<CodeInformation>,
 	position: number,
 	filter: (data: CodeInformation) => boolean
 ) {
@@ -198,7 +198,7 @@ export function toSourceOffset(
 export function toGeneratedOffset(
 	serviceScript: TypeScriptServiceScript,
 	sourceScript: SourceScript<string>,
-	map: SourceMap,
+	map: SourceMap<CodeInformation>,
 	position: number,
 	filter: (data: CodeInformation) => boolean
 ) {
diff --git a/packages/typescript/lib/node/utils.ts b/packages/typescript/lib/node/utils.ts
index 188bf2ef..101385d5 100644
--- a/packages/typescript/lib/node/utils.ts
+++ b/packages/typescript/lib/node/utils.ts
@@ -9,7 +9,7 @@ export function getServiceScript(language: Language<string>, fileName: string) {
 	if (sourceScript?.generated) {
 		const serviceScript = sourceScript.generated.languagePlugin.typescript?.getServiceScript(sourceScript.generated.root);
 		if (serviceScript) {
-			const map = language.maps.get(serviceScript.code, sourceScript.id);
+			const map = language.maps.get(serviceScript.code);
 			if (map) {
 				return [serviceScript, sourceScript, map] as const;
 			}
diff --git a/packages/typescript/lib/protocol/createProject.ts b/packages/typescript/lib/protocol/createProject.ts
index 387063e5..e46c5ff9 100644
--- a/packages/typescript/lib/protocol/createProject.ts
+++ b/packages/typescript/lib/protocol/createProject.ts
@@ -1,4 +1,4 @@
-import { FileMap, Language, SourceScript, TypeScriptExtraServiceScript, forEachEmbeddedCode } from '@volar/language-core';
+import { FileMap, Language, TypeScriptExtraServiceScript, forEachEmbeddedCode } from '@volar/language-core';
 import * as path from 'path-browserify';
 import type * as ts from 'typescript';
 import { createResolveModuleName } from '../resolveModuleName';
@@ -27,7 +27,7 @@ export function createLanguageServiceHost<T>(
 	let lastProjectVersion: number | string | undefined;
 	let tsProjectVersion = 0;
 	let tsFileRegistry = new FileMap<boolean>(sys.useCaseSensitiveFileNames);
-	let extraScriptRegistry = new FileMap<[SourceScript<T>, TypeScriptExtraServiceScript]>(sys.useCaseSensitiveFileNames);
+	let extraScriptRegistry = new FileMap<TypeScriptExtraServiceScript<T>>(sys.useCaseSensitiveFileNames);
 	let lastTsVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 	let lastOtherVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 	let languageServiceHost: ts.LanguageServiceHost = {
@@ -99,7 +99,7 @@ export function createLanguageServiceHost<T>(
 			sync();
 
 			if (extraScriptRegistry.has(fileName)) {
-				return extraScriptRegistry.get(fileName)![1].scriptKind;
+				return extraScriptRegistry.get(fileName)!.scriptKind;
 			}
 
 			const sourceScript = language.scripts.get(asScrpitId(fileName));
@@ -218,7 +218,7 @@ export function createLanguageServiceHost<T>(
 				for (const extraServiceScript of sourceScript.generated.languagePlugin.typescript?.getExtraServiceScripts?.(fileName, sourceScript.generated.root) ?? []) {
 					newTsVirtualFileSnapshots.add(extraServiceScript.code.snapshot);
 					tsFileNamesSet.add(extraServiceScript.fileName);
-					extraScriptRegistry.set(extraServiceScript.fileName, [sourceScript, extraServiceScript]);
+					extraScriptRegistry.set(extraServiceScript.fileName, extraServiceScript);
 				}
 				for (const code of forEachEmbeddedCode(sourceScript.generated.root)) {
 					newOtherVirtualFileSnapshots.add(code.snapshot);
@@ -251,7 +251,7 @@ export function createLanguageServiceHost<T>(
 		sync();
 
 		if (extraScriptRegistry.has(fileName)) {
-			return extraScriptRegistry.get(fileName)![1].code.snapshot;
+			return extraScriptRegistry.get(fileName)!.code.snapshot;
 		}
 
 		const sourceScript = language.scripts.get(asScrpitId(fileName));
@@ -278,7 +278,7 @@ export function createLanguageServiceHost<T>(
 		const version = scriptVersions.get(fileName)!;
 
 		if (extraScriptRegistry.has(fileName)) {
-			const snapshot = extraScriptRegistry.get(fileName)![1].code.snapshot;
+			const snapshot = extraScriptRegistry.get(fileName)!.code.snapshot;
 			if (!version.map.has(snapshot)) {
 				version.map.set(snapshot, version.lastVersion++);
 			}

From c9597b709e1ee02d609fe05d8618e8aeb48a6348 Mon Sep 17 00:00:00 2001
From: Johnson Chu <johnsoncodehk@gmail.com>
Date: Tue, 4 Jun 2024 19:21:34 +0800
Subject: [PATCH 3/5] wip

---
 packages/language-core/index.ts     | 6 +++---
 packages/language-core/lib/types.ts | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/language-core/index.ts b/packages/language-core/index.ts
index abb5946b..4c4428b0 100644
--- a/packages/language-core/index.ts
+++ b/packages/language-core/index.ts
@@ -142,8 +142,8 @@ export function createLanguage<T>(
 				}
 				yield [sourceScript.id, sourceScript.snapshot, mapCache.get(sourceScript.snapshot)!];
 
-				if (virtualCode.relatedMappings) {
-					for (const [relatedScriptId, relatedMappings] of virtualCode.relatedMappings) {
+				if (virtualCode.associatedScriptMappings) {
+					for (const [relatedScriptId, relatedMappings] of virtualCode.associatedScriptMappings) {
 						const relatedSourceScript = scriptRegistry.get(relatedScriptId);
 						if (relatedSourceScript) {
 							if (!mapCache.has(relatedSourceScript.snapshot)) {
@@ -194,7 +194,7 @@ export function createLanguage<T>(
 		sourceScript.relateds.clear();
 		sourceScript.isRelatedDirty = false;
 		return {
-			getRelatedSourceScript(id) {
+			getAssociatedScript(id) {
 				sync(id);
 				const relatedSourceScript = scriptRegistry.get(id);
 				if (relatedSourceScript) {
diff --git a/packages/language-core/lib/types.ts b/packages/language-core/lib/types.ts
index e2ddf6de..3ccf45a3 100644
--- a/packages/language-core/lib/types.ts
+++ b/packages/language-core/lib/types.ts
@@ -51,7 +51,7 @@ export interface VirtualCode<T = unknown> {
 	languageId: string;
 	snapshot: ts.IScriptSnapshot;
 	mappings: CodeMapping[];
-	relatedMappings?: Map<T, CodeMapping[]>;
+	associatedScriptMappings?: Map<T, CodeMapping[]>;
 	embeddedCodes?: VirtualCode<T>[];
 	linkedCodeMappings?: Mapping[];
 }
@@ -115,7 +115,7 @@ export interface LanguagePlugin<T, K extends VirtualCode<T> = VirtualCode<T>> {
 }
 
 export interface CodegenContext<T> {
-	getRelatedSourceScript(scriptId: T): SourceScript<T> | undefined;
+	getAssociatedScript(scriptId: T): SourceScript<T> | undefined;
 }
 
 /**

From 83be045b1220cfeb74ddbd0b8b44bb7380f28973 Mon Sep 17 00:00:00 2001
From: Johnson Chu <johnsoncodehk@gmail.com>
Date: Tue, 4 Jun 2024 19:33:44 +0800
Subject: [PATCH 4/5] asdf

---
 packages/language-core/index.ts     | 22 +++++++++++-----------
 packages/language-core/lib/types.ts |  6 +++---
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/packages/language-core/index.ts b/packages/language-core/index.ts
index 4c4428b0..58e4661a 100644
--- a/packages/language-core/index.ts
+++ b/packages/language-core/index.ts
@@ -48,7 +48,7 @@ export function createLanguage<T>(
 						this.delete(id);
 						return this.set(id, snapshot, languageId);
 					}
-					else if (sourceScript.isRelatedDirty || sourceScript.snapshot !== snapshot) {
+					else if (sourceScript.isAssociationDirty || sourceScript.snapshot !== snapshot) {
 						// snapshot updated
 						sourceScript.snapshot = snapshot;
 						const codegenCtx = prepareCreateVirtualCode(sourceScript);
@@ -84,8 +84,8 @@ export function createLanguage<T>(
 						id: id,
 						languageId,
 						snapshot,
-						relateds: new Set(),
-						targets: new Set(),
+						associatedIds: new Set(),
+						targetIds: new Set(),
 					};
 					scriptRegistry.set(id, sourceScript);
 
@@ -179,27 +179,27 @@ export function createLanguage<T>(
 	};
 
 	function triggerTargetsDirty(sourceScript: SourceScript<T>) {
-		sourceScript.targets.forEach(id => {
+		sourceScript.targetIds.forEach(id => {
 			const sourceScript = scriptRegistry.get(id);
 			if (sourceScript) {
-				sourceScript.isRelatedDirty = true;
+				sourceScript.isAssociationDirty = true;
 			}
 		});
 	}
 
 	function prepareCreateVirtualCode(sourceScript: SourceScript<T>): CodegenContext<T> {
-		for (const id of sourceScript.relateds) {
-			scriptRegistry.get(id)?.targets.delete(sourceScript.id);
+		for (const id of sourceScript.associatedIds) {
+			scriptRegistry.get(id)?.targetIds.delete(sourceScript.id);
 		}
-		sourceScript.relateds.clear();
-		sourceScript.isRelatedDirty = false;
+		sourceScript.associatedIds.clear();
+		sourceScript.isAssociationDirty = false;
 		return {
 			getAssociatedScript(id) {
 				sync(id);
 				const relatedSourceScript = scriptRegistry.get(id);
 				if (relatedSourceScript) {
-					relatedSourceScript.targets.add(sourceScript.id);
-					sourceScript.relateds.add(relatedSourceScript.id);
+					relatedSourceScript.targetIds.add(sourceScript.id);
+					sourceScript.associatedIds.add(relatedSourceScript.id);
 				}
 				return relatedSourceScript;
 			},
diff --git a/packages/language-core/lib/types.ts b/packages/language-core/lib/types.ts
index 3ccf45a3..b600d0f7 100644
--- a/packages/language-core/lib/types.ts
+++ b/packages/language-core/lib/types.ts
@@ -34,9 +34,9 @@ export interface SourceScript<T> {
 	id: T;
 	languageId: string;
 	snapshot: ts.IScriptSnapshot;
-	relateds: Set<T>;
-	targets: Set<T>;
-	isRelatedDirty?: boolean;
+	targetIds: Set<T>;
+	associatedIds: Set<T>;
+	isAssociationDirty?: boolean;
 	generated?: {
 		root: VirtualCode<T>;
 		languagePlugin: LanguagePlugin<T>;

From ba09c3a6c15b2bdcc1aa60bc62e9167a3fbf36a0 Mon Sep 17 00:00:00 2001
From: Johnson Chu <johnsoncodehk@gmail.com>
Date: Wed, 5 Jun 2024 08:33:21 +0800
Subject: [PATCH 5/5] Update types.ts

---
 packages/language-core/lib/types.ts | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/language-core/lib/types.ts b/packages/language-core/lib/types.ts
index b600d0f7..4fb2157d 100644
--- a/packages/language-core/lib/types.ts
+++ b/packages/language-core/lib/types.ts
@@ -2,7 +2,7 @@ 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;
@@ -30,7 +30,7 @@ export interface Language<T> {
 	};
 }
 
-export interface SourceScript<T> {
+export interface SourceScript<T = unknown> {
 	id: T;
 	languageId: string;
 	snapshot: ts.IScriptSnapshot;
@@ -90,11 +90,11 @@ export interface TypeScriptServiceScript<T = unknown> {
 	preventLeadingOffset?: boolean;
 }
 
-export interface TypeScriptExtraServiceScript<T> extends TypeScriptServiceScript<T> {
+export interface TypeScriptExtraServiceScript<T = unknown> extends TypeScriptServiceScript<T> {
 	fileName: string;
 }
 
-export interface LanguagePlugin<T, K extends VirtualCode<T> = VirtualCode<T>> {
+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.
 	 */
@@ -114,7 +114,7 @@ export interface LanguagePlugin<T, K extends VirtualCode<T> = VirtualCode<T>> {
 	typescript?: TypeScriptGenericOptions<T, K> & TypeScriptNonTSPluginOptions<T, K>;
 }
 
-export interface CodegenContext<T> {
+export interface CodegenContext<T = unknown> {
 	getAssociatedScript(scriptId: T): SourceScript<T> | undefined;
 }