fix(language-core): use a global WeakMap variable to cache inline ts asts

Author: KazariEX

packages/language-core/lib/codegen/template/context.ts

@@ -1,5 +1,6 @@
 import * as CompilerDOM from '@vue/compiler-dom';
 import type { Code, VueCodeInformation } from '../../types';
+import { templateInlineTsAsts } from '../../virtualFile/computedSfc';
 import { codeFeatures } from '../codeFeatures';
 import type { InlayHintInfo } from '../inlayHints';
 import { endOfLine, newLine } from '../utils';
@@ -107,7 +108,10 @@ const commentDirectiveRegex = /^<!--\s*@vue-(?<name>[-\w]+)\b(?<content>[\s\S]*)
  * an error/diagnostic was encountered for a region of code covered by a `@vue-expect-error` directive,
  * and additionally how we use that to determine whether to propagate diagnostics back upward.
  */
-export function createTemplateCodegenContext(options: Pick<TemplateCodegenOptions, 'scriptSetupBindingNames'>) {
+export function createTemplateCodegenContext(
+	options: Pick<TemplateCodegenOptions, 'scriptSetupBindingNames'>,
+	templateAst?: CompilerDOM.RootNode,
+) {
 	let variableId = 0;
 
 	function resolveCodeFeatures(features: VueCodeInformation) {
@@ -194,6 +198,7 @@ export function createTemplateCodegenContext(options: Pick<TemplateCodegenOption
 			},
 		}),
 		resolveCodeFeatures,
+		inlineTsAsts: templateAst && templateInlineTsAsts.get(templateAst),
 		inVFor: false,
 		slots,
 		dynamicSlots,

packages/language-core/lib/codegen/template/elementEvents.ts

@@ -117,7 +117,7 @@ export function* generateEventExpression(
 		let suffix = `)`;
 		let isFirstMapping = true;
 
-		const ast = createTsAst(options.ts, options.template.ast, prop.exp.content);
+		const ast = createTsAst(options.ts, ctx.inlineTsAsts, prop.exp.content);
 		const isCompound = isCompoundExpression(options.ts, ast);
 
 		if (isCompound) {

packages/language-core/lib/codegen/template/index.ts

@@ -25,7 +25,7 @@ export interface TemplateCodegenOptions {
 }
 
 export function* generateTemplate(options: TemplateCodegenOptions): Generator<Code, TemplateCodegenContext> {
-	const ctx = createTemplateCodegenContext(options);
+	const ctx = createTemplateCodegenContext(options, options.template.ast);
 
 	if (options.slotsAssignName) {
 		ctx.addLocalVariable(options.slotsAssignName);

packages/language-core/lib/codegen/template/interpolation.ts

@@ -1,6 +1,6 @@
 import { isGloballyAllowed, makeMap } from '@vue/shared';
 import type * as ts from 'typescript';
-import type { Code, Sfc, VueCodeInformation } from '../../types';
+import type { Code, VueCodeInformation } from '../../types';
 import { getNodeText, getStartEnd } from '../../utils/shared';
 import type { ScriptCodegenOptions } from '../script';
 import { collectVars, createTsAst, identifierRegex } from '../utils';
@@ -25,12 +25,11 @@ export function* generateInterpolation(
 		destructuredPropNames,
 		templateRefNames,
 	} = options;
-	const template = 'template' in options ? options.template : options.sfc.template;
 
 	for (
 		let [section, offset, type] of forEachInterpolationSegment(
 			ts,
-			template,
+			ctx.inlineTsAsts,
 			destructuredPropNames,
 			templateRefNames,
 			ctx,
@@ -94,7 +93,7 @@ type Segment = [
 
 function* forEachInterpolationSegment(
 	ts: typeof import('typescript'),
-	template: Sfc['template'],
+	inlineTsAsts: Map<string, ts.SourceFile> | undefined,
 	destructuredPropNames: Set<string> | undefined,
 	templateRefNames: Set<string> | undefined,
 	ctx: TemplateCodegenContext,
@@ -113,7 +112,7 @@ function* forEachInterpolationSegment(
 			offset: prefix.length,
 		});
 	} else {
-		const ast = createTsAst(ts, template?.ast, code);
+		const ast = createTsAst(ts, inlineTsAsts, code);
 		const varCb = (id: ts.Identifier, isShorthand: boolean) => {
 			const text = getNodeText(ts, id, ast);
 			if (!shouldIdentifierSkipped(ctx, text, destructuredPropNames)) {

packages/language-core/lib/codegen/template/vFor.ts

@@ -17,7 +17,7 @@ export function* generateVFor(
 
 	yield `for (const [`;
 	if (leftExpressionRange && leftExpressionText) {
-		const collectAst = createTsAst(options.ts, options.template.ast, `const [${leftExpressionText}]`);
+		const collectAst = createTsAst(options.ts, ctx.inlineTsAsts, `const [${leftExpressionText}]`);
 		collectVars(options.ts, collectAst, collectAst, forBlockVars);
 		yield [
 			leftExpressionText,

packages/language-core/lib/codegen/template/vSlot.ts

@@ -62,7 +62,7 @@ export function* generateVSlot(
 	}
 
 	if (slotDir?.exp?.type === CompilerDOM.NodeTypes.SIMPLE_EXPRESSION) {
-		const slotAst = createTsAst(options.ts, options.template.ast, `(${slotDir.exp.content}) => {}`);
+		const slotAst = createTsAst(options.ts, ctx.inlineTsAsts, `(${slotDir.exp.content}) => {}`);
 		collectVars(options.ts, slotAst, slotAst, slotBlockVars);
 		yield* generateSlotParameters(options, ctx, slotAst, slotDir.exp, slotVar);
 	}

packages/language-core/lib/codegen/utils/index.ts

@@ -65,17 +65,16 @@ export function normalizeAttributeValue(node: CompilerDOM.TextNode): [string, nu
 
 export function createTsAst(
 	ts: typeof import('typescript'),
-	templateAst: CompilerDOM.RootNode | undefined,
+	inlineTsAsts: Map<string, ts.SourceFile> | undefined,
 	text: string,
 ) {
-	const inlineTsAsts = (templateAst as any)?.__volar_inlineTsAsts;
 	let ast = inlineTsAsts?.get(text);
 	if (!ast) {
 		ast = ts.createSourceFile('/a.ts', text, 99 satisfies ts.ScriptTarget.ESNext);
 		inlineTsAsts?.set(text, ast);
 	}
-	ast.__volar_used = true;
-	return ast as ts.SourceFile;
+	(ast as any).__volar_used = true;
+	return ast;
 }
 
 export function generateSfcBlockSection(

packages/language-core/lib/plugins/vue-template-inline-ts.ts

@@ -5,6 +5,7 @@ import { parseInterpolationNode } from '../codegen/template/templateChild';
 import { parseVForNode } from '../codegen/template/vFor';
 import { createTsAst } from '../codegen/utils';
 import type { Code, Sfc, VueLanguagePlugin } from '../types';
+import { templateInlineTsAsts } from '../virtualFile/computedSfc';
 
 const codeFeatures: CodeInformation = {
 	format: true,
@@ -68,6 +69,7 @@ const plugin: VueLanguagePlugin = ctx => {
 			return data;
 		}
 		const templateContent = sfc.template.content;
+		const inlineTsAsts = sfc.template.ast && templateInlineTsAsts.get(sfc.template.ast);
 		let i = 0;
 		sfc.template.ast.children.forEach(visit);
 		return data;
@@ -104,7 +106,7 @@ const plugin: VueLanguagePlugin = ctx => {
 						&& prop.exp.constType !== CompilerDOM.ConstantTypes.CAN_STRINGIFY // style='z-index: 2' will compile to {'z-index':'2'}
 					) {
 						if (prop.name === 'on' && prop.arg?.type === CompilerDOM.NodeTypes.SIMPLE_EXPRESSION) {
-							const ast = createTsAst(ctx.modules.typescript, sfc.template!.ast, prop.exp.content);
+							const ast = createTsAst(ctx.modules.typescript, inlineTsAsts, prop.exp.content);
 							if (isCompoundExpression(ctx.modules.typescript, ast)) {
 								addFormatCodes(
 									prop.exp.loc.source,

packages/language-core/lib/virtualFile/computedSfc.ts

@@ -8,6 +8,8 @@ import { parseCssImports } from '../utils/parseCssImports';
 import { parseCssVars } from '../utils/parseCssVars';
 import { computedArray, computedItems } from '../utils/signals';
 
+export const templateInlineTsAsts = new WeakMap<CompilerDOM.RootNode, Map<string, ts.SourceFile>>();
+
 export function computedSfc(
 	ts: typeof import('typescript'),
 	plugins: VueLanguagePluginReturn[],
@@ -215,8 +217,11 @@ export function computedSfc(
 		let inlineTsAsts: Map<string, any> | undefined;
 
 		function updateInlineTsAsts(newAst: CompilerDOM.RootNode, oldAst?: CompilerDOM.RootNode) {
-			const newTsAsts: Map<string, any> = (newAst as any).__volar_inlineTsAsts ??= new Map();
-			const oldTsAsts: Map<string, any> = (oldAst as any)?.__volar_inlineTsAsts ?? inlineTsAsts;
+			let newTsAsts = templateInlineTsAsts.get(newAst);
+			if (!newTsAsts) {
+				templateInlineTsAsts.set(newAst, newTsAsts = new Map());
+			}
+			const oldTsAsts = oldAst && templateInlineTsAsts.get(oldAst) || inlineTsAsts;
 
 			if (oldTsAsts) {
 				for (const [text, ast] of oldTsAsts) {