feat(typescript): re-support extra virtual scripts for LSP and Kit (#132

)
volarjs · Feb 4, 2024 · 2fb4bb4 · 2fb4bb4
1 parent 5130348
commit 2fb4bb4
Show file tree

Hide file tree

Showing 5 changed files with 95 additions and 20 deletions.
diff --git a/packages/language-core/lib/types.ts b/packages/language-core/lib/types.ts
@@ -53,18 +53,37 @@ export interface CodeInformation {
 	format: boolean;
 }
 
+export interface ServiceScript {
+	code: VirtualCode;
+	extension: '.ts' | '.js' | '.mts' | '.mjs' | '.cjs' | '.cts' | '.d.ts' | string;
+	scriptKind: ts.ScriptKind;
+}
+
+export interface ExtraServiceScript extends ServiceScript {
+	fileName: string;
+}
+
 export interface LanguagePlugin<T extends VirtualCode = VirtualCode> {
 	createVirtualCode(fileId: string, languageId: string, snapshot: ts.IScriptSnapshot, files?: FileRegistry): T | undefined;
 	updateVirtualCode(fileId: string, virtualCode: T, newSnapshot: ts.IScriptSnapshot, files?: FileRegistry): T;
 	disposeVirtualCode?(fileId: string, virtualCode: T, files?: FileRegistry): void;
 	typescript?: {
+		/**
+		 * LSP + TS Plugin
+		 */
 		extraFileExtensions: ts.FileExtensionInfo[];
+		/**
+		 * LSP + TS Plugin
+		 */
+		getScript(rootVirtualCode: T): ServiceScript | undefined;
+		/**
+		 * LSP only
+		 */
+		getExtraScripts?(fileName: string, rootVirtualCode: T): ExtraServiceScript[];
+		/**
+		 * LSP only
+		 */
 		resolveLanguageServiceHost?(host: ts.LanguageServiceHost): ts.LanguageServiceHost;
-		getScript(rootVirtualCode: T): {
-			code: VirtualCode;
-			extension: '.ts' | '.js' | '.mts' | '.mjs' | '.cjs' | '.cts' | '.d.ts' | string;
-			scriptKind: ts.ScriptKind;
-		} | undefined;
 	};
 }
 
@@ -75,6 +94,7 @@ export interface LanguageContext {
 		sys: ts.System & { sync?(): Promise<number>; };
 		projectHost: TypeScriptProjectHost;
 		languageServiceHost: ts.LanguageServiceHost;
+		getExtraScript(fileName: string): ExtraServiceScript | undefined;
 	};
 }
 

diff --git a/packages/language-server/lib/register/registerEditorFeatures.ts b/packages/language-server/lib/register/registerEditorFeatures.ts
@@ -132,10 +132,16 @@ export function registerEditorFeatures(
 					if (uri.startsWith(rootUri)) {
 						const sourceFile = languageService.context.language.files.get(uri);
 						if (sourceFile?.generated) {
-							const virtualFile = sourceFile.generated.languagePlugin.typescript?.getScript(sourceFile.generated.code);
-							if (virtualFile) {
-								const { snapshot } = virtualFile.code;
-								fs.writeFile(uri + virtualFile.extension, snapshot.getText(0, snapshot.getLength()), () => { });
+							const mainScript = sourceFile.generated.languagePlugin.typescript?.getScript(sourceFile.generated.code);
+							if (mainScript) {
+								const { snapshot } = mainScript.code;
+								fs.writeFile(fileName + mainScript.extension, snapshot.getText(0, snapshot.getLength()), () => { });
+							}
+							if (sourceFile.generated.languagePlugin.typescript?.getExtraScripts) {
+								for (const extraScript of sourceFile.generated.languagePlugin.typescript.getExtraScripts(uri, sourceFile.generated.code)) {
+									const { snapshot } = extraScript.code;
+									fs.writeFile(fileName, snapshot.getText(0, snapshot.getLength()), () => { });
+								}
 							}
 						}
 					}

diff --git a/packages/typescript/lib/node/decorateLanguageServiceHost.ts b/packages/typescript/lib/node/decorateLanguageServiceHost.ts
@@ -180,6 +180,9 @@ export function decorateLanguageServiceHost(
 						patchedText += script.code.snapshot.getText(0, script.code.snapshot.getLength());
 					}
 					snapshotSnapshot = ts.ScriptSnapshot.fromString(patchedText);
+					if (sourceFile.generated.languagePlugin.typescript?.getExtraScripts) {
+						console.warn('getExtraScripts() is not available in this use case.');
+					}
 				}
 			}
 			else if (virtualFiles.get(fileName)) {

diff --git a/packages/typescript/lib/node/proxyCreateProgram.ts b/packages/typescript/lib/node/proxyCreateProgram.ts
@@ -83,12 +83,16 @@ export function proxyCreateProgram(
 						assert(!!sourceFile, '!!sourceFile');
 						let patchedText = originalSourceFile.text.split('\n').map(line => ' '.repeat(line.length)).join('\n');
 						let scriptKind = ts.ScriptKind.TS;
-						if (sourceFile.generated) {
-							const script = sourceFile.generated.languagePlugin.typescript?.getScript(sourceFile.generated.code);
+						if (sourceFile.generated?.languagePlugin.typescript) {
+							const { getScript, getExtraScripts } = sourceFile.generated.languagePlugin.typescript;
+							const script = getScript(sourceFile.generated.code);
 							if (script) {
 								scriptKind = script.scriptKind;
 								patchedText += script.code.snapshot.getText(0, script.code.snapshot.getLength());
 							}
+							if (getExtraScripts) {
+								console.warn('getExtraScripts() is not available in this use case.');
+							}
 						}
 						sourceFile2 = ts.createSourceFile(
 							fileName,

diff --git a/packages/typescript/lib/protocol/createProject.ts b/packages/typescript/lib/protocol/createProject.ts
@@ -1,4 +1,4 @@
-import { createFileRegistry, FileMap, LanguagePlugin, LanguageContext, TypeScriptProjectHost } from '@volar/language-core';
+import { createFileRegistry, FileMap, LanguagePlugin, LanguageContext, TypeScriptProjectHost, ExtraServiceScript } from '@volar/language-core';
 import type * as ts from 'typescript';
 import { forEachEmbeddedCode } from '@volar/language-core';
 import * as path from 'path-browserify';
@@ -51,7 +51,7 @@ export function createLanguage(
 		}
 	});
 
-	let languageServiceHost = createLanguageServiceHost();
+	let { languageServiceHost, getExtraScript } = createLanguageServiceHost();
 
 	for (const language of languagePlugins) {
 		if (language.typescript?.resolveLanguageServiceHost) {
@@ -134,6 +134,7 @@ export function createLanguage(
 			sys,
 			projectHost,
 			languageServiceHost,
+			getExtraScript,
 		},
 	};
 
@@ -142,6 +143,7 @@ export function createLanguage(
 		let lastProjectVersion: number | string | undefined;
 		let tsProjectVersion = 0;
 		let tsFileRegistry = new FileMap<boolean>(sys.useCaseSensitiveFileNames);
+		let extraScriptRegistry = new FileMap<ExtraServiceScript>(sys.useCaseSensitiveFileNames);
 		let lastTsVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 		let lastOtherVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 
@@ -201,14 +203,21 @@ export function createLanguage(
 				return getScriptVersion(fileName) !== '';
 			},
 			getProjectVersion() {
-				syncProject();
+				sync();
 				return tsProjectVersion + ('version' in sys ? `:${sys.version}` : '');
 			},
 			getScriptFileNames() {
-				syncProject();
+				sync();
 				return [...tsFileRegistry.keys()];
 			},
 			getScriptKind(fileName) {
+
+				sync();
+
+				if (extraScriptRegistry.has(fileName)) {
+					return extraScriptRegistry.get(fileName)!.scriptKind;
+				}
+
 				const sourceFile = files.get(fileNameToFileId(fileName));
 				if (sourceFile?.generated) {
 					const tsCode = sourceFile.generated.languagePlugin.typescript?.getScript(sourceFile.generated.code);
@@ -239,9 +248,17 @@ export function createLanguage(
 			getScriptSnapshot,
 		};
 
-		return languageServiceHost;
+		return {
+			languageServiceHost,
+			getExtraScript,
+		};
+
+		function getExtraScript(fileName: string) {
+			sync();
+			return extraScriptRegistry.get(fileName);
+		}
 
-		function syncProject() {
+		function sync() {
 
 			const newProjectVersion = projectHost.getProjectVersion?.();
 			const shouldUpdate = newProjectVersion === undefined || newProjectVersion !== lastProjectVersion;
@@ -250,6 +267,7 @@ export function createLanguage(
 			}
 
 			lastProjectVersion = newProjectVersion;
+			extraScriptRegistry.clear();
 
 			const newTsVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
 			const newOtherVirtualFileSnapshots = new Set<ts.IScriptSnapshot>();
@@ -261,10 +279,15 @@ export function createLanguage(
 					const script = sourceFile.generated.languagePlugin.typescript?.getScript(sourceFile.generated.code);
 					if (script) {
 						newTsVirtualFileSnapshots.add(script.code.snapshot);
-						tsFileNamesSet.add(fileName); // virtual .ts
+						tsFileNamesSet.add(fileName);
 					}
-					for (const file of forEachEmbeddedCode(sourceFile.generated.code)) {
-						newOtherVirtualFileSnapshots.add(file.snapshot);
+					for (const extraScript of sourceFile.generated.languagePlugin.typescript?.getExtraScripts?.(fileName, sourceFile.generated.code) ?? []) {
+						newTsVirtualFileSnapshots.add(extraScript.code.snapshot);
+						tsFileNamesSet.add(extraScript.fileName);
+						extraScriptRegistry.set(extraScript.fileName, extraScript);
+					}
+					for (const code of forEachEmbeddedCode(sourceFile.generated.code)) {
+						newOtherVirtualFileSnapshots.add(code.snapshot);
 					}
 				}
 				else {
@@ -291,6 +314,12 @@ export function createLanguage(
 
 		function getScriptSnapshot(fileName: string) {
 
+			sync();
+
+			if (extraScriptRegistry.has(fileName)) {
+				return extraScriptRegistry.get(fileName)!.code.snapshot;
+			}
+
 			const sourceFile = files.get(fileNameToFileId(fileName));
 
 			if (sourceFile?.generated) {
@@ -306,12 +335,24 @@ export function createLanguage(
 
 		function getScriptVersion(fileName: string): string {
 
+			sync();
+
 			if (!scriptVersions.has(fileName)) {
 				scriptVersions.set(fileName, { lastVersion: 0, map: new WeakMap() });
 			}
 
 			const version = scriptVersions.get(fileName)!;
+
+			if (extraScriptRegistry.has(fileName)) {
+				const snapshot = extraScriptRegistry.get(fileName)!.code.snapshot;
+				if (!version.map.has(snapshot)) {
+					version.map.set(snapshot, version.lastVersion++);
+				}
+				return version.map.get(snapshot)!.toString();
+			}
+
 			const sourceFile = files.get(fileNameToFileId(fileName));
+
 			if (sourceFile?.generated) {
 				const script = sourceFile.generated.languagePlugin.typescript?.getScript(sourceFile.generated.code);
 				if (script) {
@@ -323,6 +364,7 @@ export function createLanguage(
 			}
 
 			const isOpenedFile = !!projectHost.getScriptSnapshot(fileName);
+
 			if (isOpenedFile) {
 				const sourceFile = files.get(fileNameToFileId(fileName));
 				if (sourceFile && !sourceFile.generated) {