fix(language-service): read ast from codegen instead of parsing it repeatedly (#5086)

Author: KazariEX

packages/language-core/index.ts

@@ -8,7 +8,6 @@ export * from './lib/utils/parseSfc';
 export * from './lib/utils/ts';
 export * from './lib/virtualFile/vueFile';
 
-export * as scriptRanges from './lib/parsers/scriptRanges';
 export { tsCodegen } from './lib/plugins/vue-tsx';
 export * from './lib/utils/shared';
 

packages/language-service/lib/plugins/utils.ts

@@ -0,0 +1,12 @@
+import type { TextDocument } from 'vscode-languageserver-textdocument';
+
+export function sleep(ms: number) {
+	return new Promise(resolve => setTimeout(resolve, ms));
+}
+
+export function isTsDocument(document: TextDocument) {
+	return document.languageId === 'javascript' ||
+		document.languageId === 'typescript' ||
+		document.languageId === 'javascriptreact' ||
+		document.languageId === 'typescriptreact';
+}

packages/language-service/lib/plugins/vue-autoinsert-dotvalue.ts

@@ -1,19 +1,9 @@
 import type { LanguageServiceContext, LanguageServicePlugin } from '@volar/language-service';
-import { hyphenateAttr } from '@vue/language-core';
+import { hyphenateAttr, VueVirtualCode } from '@vue/language-core';
 import type * as ts from 'typescript';
 import type { TextDocument } from 'vscode-languageserver-textdocument';
 import { URI } from 'vscode-uri';
-
-const asts = new WeakMap<ts.IScriptSnapshot, ts.SourceFile>();
-
-function getAst(ts: typeof import('typescript'), fileName: string, snapshot: ts.IScriptSnapshot, scriptKind?: ts.ScriptKind) {
-	let ast = asts.get(snapshot);
-	if (!ast) {
-		ast = ts.createSourceFile(fileName, snapshot.getText(0, snapshot.getLength()), ts.ScriptTarget.Latest, undefined, scriptKind);
-		asts.set(snapshot, ast);
-	}
-	return ast;
-}
+import { isTsDocument, sleep } from './utils';
 
 export function create(
 	ts: typeof import('typescript'),
@@ -61,46 +51,49 @@ export function create(
 					const decoded = context.decodeEmbeddedDocumentUri(uri);
 					const sourceScript = decoded && context.language.scripts.get(decoded[0]);
 					const virtualCode = decoded && sourceScript?.generated?.embeddedCodes.get(decoded[1]);
-					if (!sourceScript) {
+					if (!sourceScript?.generated || !virtualCode) {
 						return;
 					}
 
-					let ast: ts.SourceFile | undefined;
-					let sourceCodeOffset = document.offsetAt(selection);
+					const root = sourceScript.generated.root;
+					if (!(root instanceof VueVirtualCode)) {
+						return;
+					}
 
-					const fileName = context.project.typescript?.uriConverter.asFileName(sourceScript.id)
-						?? sourceScript.id.fsPath.replace(/\\/g, '/');
+					const blocks = [
+						root._sfc.script,
+						root._sfc.scriptSetup,
+					].filter(block => !!block);
+					if (!blocks.length) {
+						return;
+					}
 
-					if (sourceScript.generated) {
-						const serviceScript = sourceScript.generated.languagePlugin.typescript?.getServiceScript(sourceScript.generated.root);
-						if (!serviceScript || serviceScript?.code !== virtualCode) {
-							return;
-						}
-						ast = getAst(ts, fileName, virtualCode.snapshot, serviceScript.scriptKind);
-						let mapped = false;
-						for (const [_sourceScript, map] of context.language.maps.forEach(virtualCode)) {
-							for (const [sourceOffset] of map.toSourceLocation(document.offsetAt(selection))) {
-								sourceCodeOffset = sourceOffset;
-								mapped = true;
-								break;
-							}
-							if (mapped) {
-								break;
-							}
+					let sourceCodeOffset = document.offsetAt(selection);
+					let mapped = false;
+					for (const [, map] of context.language.maps.forEach(virtualCode)) {
+						for (const [sourceOffset] of map.toSourceLocation(sourceCodeOffset)) {
+							sourceCodeOffset = sourceOffset;
+							mapped = true;
+							break;
 						}
-						if (!mapped) {
-							return;
+						if (mapped) {
+							break;
 						}
 					}
-					else {
-						ast = getAst(ts, fileName, sourceScript.snapshot);
+					if (!mapped) {
+						return;
 					}
 
-					if (isBlacklistNode(ts, ast, document.offsetAt(selection), false)) {
-						return;
+					for (const { ast, startTagEnd, endTagStart } of blocks) {
+						if (sourceCodeOffset < startTagEnd || sourceCodeOffset > endTagStart) {
+							continue;
+						}
+						if (isBlacklistNode(ts, ast, sourceCodeOffset - startTagEnd, false)) {
+							return;
+						}
 					}
 
-					const props = await tsPluginClient?.getPropertiesAtLocation(fileName, sourceCodeOffset) ?? [];
+					const props = await tsPluginClient?.getPropertiesAtLocation(root.fileName, sourceCodeOffset) ?? [];
 					if (props.some(prop => prop === 'value')) {
 						return '${1:.value}';
 					}
@@ -110,20 +103,9 @@ export function create(
 	};
 }
 
-function sleep(ms: number) {
-	return new Promise(resolve => setTimeout(resolve, ms));
-}
-
-export function isTsDocument(document: TextDocument) {
-	return document.languageId === 'javascript' ||
-		document.languageId === 'typescript' ||
-		document.languageId === 'javascriptreact' ||
-		document.languageId === 'typescriptreact';
-}
-
 const charReg = /\w/;
 
-export function isCharacterTyping(document: TextDocument, change: { text: string; rangeOffset: number; rangeLength: number; }) {
+function isCharacterTyping(document: TextDocument, change: { text: string; rangeOffset: number; rangeLength: number; }) {
 	const lastCharacter = change.text[change.text.length - 1];
 	const nextCharacter = document.getText().slice(
 		change.rangeOffset + change.text.length,
@@ -138,7 +120,7 @@ export function isCharacterTyping(document: TextDocument, change: { text: string
 	return charReg.test(lastCharacter) && !charReg.test(nextCharacter);
 }
 
-export function isBlacklistNode(ts: typeof import('typescript'), node: ts.Node, pos: number, allowAccessDotValue: boolean) {
+function isBlacklistNode(ts: typeof import('typescript'), node: ts.Node, pos: number, allowAccessDotValue: boolean) {
 	if (ts.isVariableDeclaration(node) && pos >= node.name.getFullStart() && pos <= node.name.getEnd()) {
 		return true;
 	}

packages/language-service/lib/plugins/vue-complete-define-assignment.ts

@@ -2,7 +2,7 @@ import type { LanguageServicePlugin } from '@volar/language-service';
 import { TextRange, tsCodegen, VueVirtualCode } from '@vue/language-core';
 import type * as vscode from 'vscode-languageserver-protocol';
 import { URI } from 'vscode-uri';
-import { isTsDocument } from './vue-autoinsert-dotvalue';
+import { isTsDocument } from './utils';
 
 export function create(): LanguageServicePlugin {
 	return {

packages/language-service/lib/plugins/vue-document-drop.ts

@@ -111,7 +111,7 @@ export function create(
 					});
 
 					if (sfc.script) {
-						const edit = createAddComponentToOptionEdit(ts, sfc.script.ast, newName);
+						const edit = createAddComponentToOptionEdit(ts, sfc, sfc.script.ast, newName);
 						if (edit) {
 							additionalEdit.changes[embeddedDocumentUriStr].push({
 								range: {

packages/language-service/lib/plugins/vue-extract-file.ts

@@ -1,6 +1,6 @@
 import type { CreateFile, LanguageServiceContext, LanguageServicePlugin, TextDocumentEdit, TextEdit } from '@volar/language-service';
 import type { ExpressionNode, TemplateChildNode } from '@vue/compiler-dom';
-import { Sfc, VueVirtualCode, scriptRanges } from '@vue/language-core';
+import { Sfc, VueVirtualCode, tsCodegen } from '@vue/language-core';
 import type * as ts from 'typescript';
 import type * as vscode from 'vscode-languageserver-protocol';
 import { URI } from 'vscode-uri';
@@ -158,7 +158,7 @@ export function create(
 					];
 
 					if (sfc.script) {
-						const edit = createAddComponentToOptionEdit(ts, sfc.script.ast, newName);
+						const edit = createAddComponentToOptionEdit(ts, sfc, sfc.script.ast, newName);
 						if (edit) {
 							sfcEdits.push({
 								range: {
@@ -304,12 +304,13 @@ export function getLastImportNode(ts: typeof import('typescript'), sourceFile: t
 	return lastImportNode;
 }
 
-export function createAddComponentToOptionEdit(ts: typeof import('typescript'), ast: ts.SourceFile, componentName: string) {
+export function createAddComponentToOptionEdit(ts: typeof import('typescript'), sfc: Sfc, ast: ts.SourceFile, componentName: string) {
 
-	const exportDefault = scriptRanges.parseScriptRanges(ts, ast, false, true).exportDefault;
-	if (!exportDefault) {
+	const scriptRanges = tsCodegen.get(sfc)?.scriptRanges.get();
+	if (!scriptRanges?.exportDefault) {
 		return;
 	}
+	const { exportDefault } = scriptRanges;
 
 	// https://github.com/microsoft/TypeScript/issues/36174
 	const printer = ts.createPrinter();

packages/language-service/lib/plugins/vue-template.ts

@@ -1,5 +1,5 @@
 import type { Disposable, LanguageServiceContext, LanguageServicePluginInstance } from '@volar/language-service';
-import { VueCompilerOptions, VueVirtualCode, hyphenateAttr, hyphenateTag, parseScriptSetupRanges } from '@vue/language-core';
+import { VueVirtualCode, hyphenateAttr, hyphenateTag, tsCodegen } from '@vue/language-core';
 import { camelize, capitalize } from '@vue/shared';
 import { getComponentSpans } from '@vue/typescript-plugin/lib/common';
 import { create as createHtmlService } from 'volar-service-html';
@@ -157,7 +157,6 @@ export function create(
 					if (!context.project.vue) {
 						return;
 					}
-					const vueCompilerOptions = context.project.vue.compilerOptions;
 
 					let sync: (() => Promise<number>) | undefined;
 					let currentVersion: number | undefined;
@@ -172,7 +171,7 @@ export function create(
 						// #4298: Precompute HTMLDocument before provideHtmlData to avoid parseHTMLDocument requesting component names from tsserver
 						baseServiceInstance.provideCompletionItems?.(document, position, completionContext, token);
 
-						sync = (await provideHtmlData(vueCompilerOptions, sourceScript!.id, root)).sync;
+						sync = (await provideHtmlData(sourceScript!.id, root)).sync;
 						currentVersion = await sync();
 					}
 
@@ -462,7 +461,7 @@ export function create(
 				},
 			};
 
-			async function provideHtmlData(vueCompilerOptions: VueCompilerOptions, sourceDocumentUri: URI, vueCode: VueVirtualCode) {
+			async function provideHtmlData(sourceDocumentUri: URI, vueCode: VueVirtualCode) {
 
 				await (initializing ??= initialize());
 
@@ -520,9 +519,7 @@ export function create(
 								})());
 								return [];
 							}
-							const scriptSetupRanges = vueCode._sfc.scriptSetup
-								? parseScriptSetupRanges(ts, vueCode._sfc.scriptSetup.ast, vueCompilerOptions)
-								: undefined;
+							const scriptSetupRanges = tsCodegen.get(vueCode._sfc)?.scriptSetupRanges.get();
 							const names = new Set<string>();
 							const tags: html.ITagData[] = [];