feat(language-core): auto infer type for $refs & useTemplateRef (#4644)

packages/language-core/lib/codegen/script/component.ts

@@ -36,6 +36,9 @@ export function* generateComponent(
 	if (options.sfc.script && options.scriptRanges) {
 		yield* generateScriptOptions(options.sfc.script, options.scriptRanges);
 	}
+	if (options.vueCompilerOptions.target >= 3.5 && scriptSetupRanges.templateRefs.length) {
+		yield `__typeRefs: {} as __VLS_Refs,${newLine}`;
+	}
 	yield `})`;
 }
 
@@ -83,10 +86,10 @@ export function* generatePropsOption(
 	if (options.vueCompilerOptions.target >= 3.5) {
 		const types = [];
 		if (inheritAttrs && options.templateCodegen?.inheritedAttrVars.size) {
-			types.push('typeof __VLS_template>[1]');
+			types.push(`ReturnType<typeof __VLS_template>['attrs']`);
 		}
 		if (ctx.generatedPropsType) {
-			types.push('{} as __VLS_PublicProps');
+			types.push(`{} as __VLS_PublicProps`);
 		}
 		if (types.length) {
 			yield `__typeProps: ${types.join(' & ')},${newLine}`;
@@ -97,7 +100,7 @@ export function* generatePropsOption(
 
 		if (inheritAttrs && options.templateCodegen?.inheritedAttrVars.size) {
 			codegens.push(function* () {
-				yield `{} as ${ctx.helperTypes.TypePropsToOption.name}<__VLS_PickNotAny<${ctx.helperTypes.OmitIndexSignature.name}<ReturnType<typeof __VLS_template>[1]>, {}>>`;
+				yield `{} as ${ctx.helperTypes.TypePropsToOption.name}<__VLS_PickNotAny<${ctx.helperTypes.OmitIndexSignature.name}<ReturnType<typeof __VLS_template>['attrs']>, {}>>`;
 			});
 		}
 

packages/language-core/lib/codegen/script/globalTypes.ts

@@ -129,6 +129,13 @@ declare global {
 		>
 	>;
 	type __VLS_PrettifyGlobal<T> = { [K in keyof T]: T[K]; } & {};
+	type __VLS_PickRefsExpose<T> = T extends object
+		? { [K in keyof T]: (T[K] extends any[]
+		? Parameters<T[K][0]['expose']>[0][]
+		: T[K] extends { expose?: (exposed: infer E) => void }
+		? E
+		: T[K]) | null }
+		: never;
 }
 export const __VLS_globalTypesEnd = {};
 `;

packages/language-core/lib/codegen/script/internalComponent.ts

@@ -48,6 +48,9 @@ export function* generateInternalComponent(
 		}
 		yield `}${endOfLine}`; // return {
 		yield `},${newLine}`; // setup() {
+		if (options.vueCompilerOptions.target >= 3.5) {
+			yield `__typeRefs: {} as __VLS_Refs,${newLine}`;
+		}
 		if (options.sfc.scriptSetup && options.scriptSetupRanges && !ctx.bypassDefineComponent) {
 			yield* generateScriptSetupOptions(options, ctx, options.sfc.scriptSetup, options.scriptSetupRanges, false);
 		}

packages/language-core/lib/codegen/script/scriptSetup.ts

@@ -67,7 +67,7 @@ export function* generateScriptSetup(
 			+ `			props: ${ctx.helperTypes.Prettify.name}<typeof __VLS_functionalComponentProps & __VLS_PublicProps> & __VLS_BuiltInPublicProps,${newLine}`
 			+ `			expose(exposed: import('${options.vueCompilerOptions.lib}').ShallowUnwrapRef<${scriptSetupRanges.expose.define ? 'typeof __VLS_exposed' : '{}'}>): void,${newLine}`
 			+ `			attrs: any,${newLine}`
-			+ `			slots: ReturnType<typeof __VLS_template>[0],${newLine}`
+			+ `			slots: __VLS_Slots,${newLine}`
 			+ `			emit: ${emitTypes.join(' & ')},${newLine}`
 			+ `		}${endOfLine}`;
 		yield `	})(),${newLine}`; // __VLS_setup = (async () => {
@@ -116,7 +116,7 @@ function* generateSetupFunction(
 	if (options.vueCompilerOptions.target >= 3.3) {
 		yield `const { `;
 		for (const macro of Object.keys(options.vueCompilerOptions.macros)) {
-			if (!ctx.bindingNames.has(macro)) {
+			if (!ctx.bindingNames.has(macro) && macro !== 'templateRef') {
 				yield macro + `, `;
 			}
 		}
@@ -211,6 +211,11 @@ function* generateSetupFunction(
 			]);
 		}
 	}
+	for (const { define } of scriptSetupRanges.templateRefs) {
+		if (define?.arg) {
+			setupCodeModifies.push([[`<__VLS_Refs[${scriptSetup.content.slice(define.arg.start, define.arg.end)}], keyof __VLS_Refs>`], define.arg.start - 1, define.arg.start - 1]);
+		}
+	}
 	setupCodeModifies = setupCodeModifies.sort((a, b) => a[1] - b[1]);
 
 	if (setupCodeModifies.length) {
@@ -243,14 +248,16 @@ function* generateSetupFunction(
 	yield* generateComponentProps(options, ctx, scriptSetup, scriptSetupRanges, definePropMirrors);
 	yield* generateModelEmits(options, scriptSetup, scriptSetupRanges);
 	yield* generateTemplate(options, ctx, false);
+	yield `type __VLS_Refs = ReturnType<typeof __VLS_template>['refs']${endOfLine}`;
+	yield `type __VLS_Slots = ReturnType<typeof __VLS_template>['slots']${endOfLine}`;
 
 	if (syntax) {
 		if (!options.vueCompilerOptions.skipTemplateCodegen && (options.templateCodegen?.hasSlot || scriptSetupRanges?.slots.define)) {
 			yield `const __VLS_component = `;
 			yield* generateComponent(options, ctx, scriptSetup, scriptSetupRanges);
 			yield endOfLine;
 			yield `${syntax} `;
-			yield `{} as ${ctx.helperTypes.WithTemplateSlots.name}<typeof __VLS_component, ReturnType<typeof __VLS_template>[0]>${endOfLine}`;
+			yield `{} as ${ctx.helperTypes.WithTemplateSlots.name}<typeof __VLS_component, __VLS_Slots>${endOfLine}`;
 		}
 		else {
 			yield `${syntax} `;

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

@@ -165,11 +165,16 @@ function* generateTemplateContext(
 		yield `// no template${newLine}`;
 		if (!options.scriptSetupRanges?.slots.define) {
 			yield `const __VLS_slots = {}${endOfLine}`;
+			yield `const __VLS_refs = {}${endOfLine}`;
 			yield `const __VLS_inheritedAttrs = {}${endOfLine}`;
 		}
 	}
 
-	yield `return [${options.scriptSetupRanges?.slots.name ?? '__VLS_slots'}, __VLS_inheritedAttrs] as const${endOfLine}`;
+	yield `return {${newLine}`;
+	yield `slots: ${options.scriptSetupRanges?.slots.name ?? '__VLS_slots'},${newLine}`;
+	yield `refs: __VLS_refs as __VLS_PickRefsExpose<typeof __VLS_refs>,${newLine}`;
+	yield `attrs: __VLS_inheritedAttrs,${newLine}`;
+	yield `}${endOfLine}`;
 }
 
 function* generateCssClassProperty(

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

@@ -259,13 +259,26 @@ export function* generateComponent(
 		yield endOfLine;
 	}
 
-	yield* generateVScope(options, ctx, node, props);
+	const refName = yield* generateVScope(options, ctx, node, props);
 
 	ctx.usedComponentCtxVars.add(componentCtxVar);
 	const usedComponentEventsVar = yield* generateElementEvents(options, ctx, node, var_functionalComponent, var_componentInstance, var_componentEmit, var_componentEvents);
 
 	if (var_defineComponentCtx && ctx.usedComponentCtxVars.has(var_defineComponentCtx)) {
 		yield `const ${componentCtxVar} = __VLS_nonNullable(__VLS_pickFunctionalComponentCtx(${var_originalComponent}, ${var_componentInstance}))${endOfLine}`;
+		if (refName) {
+			yield `// @ts-ignore${newLine}`;
+			if (node.codegenNode?.type === CompilerDOM.NodeTypes.VNODE_CALL
+				&& node.codegenNode.props?.type === CompilerDOM.NodeTypes.JS_OBJECT_EXPRESSION
+				&& node.codegenNode.props.properties.find(({ key }) => key.type === CompilerDOM.NodeTypes.SIMPLE_EXPRESSION && key.content === 'ref_for')
+			) {
+				yield `(${refName} ??= []).push(${var_defineComponentCtx})`;
+			} else {
+				yield `${refName} = ${var_defineComponentCtx}`;
+			}
+
+			yield endOfLine;
+		}
 	}
 	if (usedComponentEventsVar) {
 		yield `let ${var_componentEmit}!: typeof ${componentCtxVar}.emit${endOfLine}`;
@@ -350,7 +363,17 @@ export function* generateElement(
 		yield endOfLine;
 	}
 
-	yield* generateVScope(options, ctx, node, node.props);
+	const refName = yield* generateVScope(options, ctx, node, node.props);
+	if (refName) {
+		yield `// @ts-ignore${newLine}`;
+		yield `${refName} = __VLS_intrinsicElements`;
+		yield* generatePropertyAccess(
+			options,
+			ctx,
+			node.tag
+		);
+		yield endOfLine;
+	}
 
 	const slotDir = node.props.find(p => p.type === CompilerDOM.NodeTypes.DIRECTIVE && p.name === 'slot') as CompilerDOM.DirectiveNode;
 	if (slotDir && componentCtxVar) {
@@ -397,13 +420,14 @@ function* generateVScope(
 	}
 
 	yield* generateElementDirectives(options, ctx, node);
-	yield* generateReferencesForElements(options, ctx, node); // <el ref="foo" />
+	const refName = yield* generateReferencesForElements(options, ctx, node); // <el ref="foo" />
 	yield* generateReferencesForScopedCssClasses(options, ctx, node);
 
 	if (inScope) {
 		yield `}${newLine}`;
 		ctx.blockConditions.length = originalConditionsNum;
 	}
+	return refName;
 }
 
 export function getCanonicalComponentName(tagText: string) {
@@ -571,6 +595,10 @@ function* generateReferencesForElements(
 				')'
 			);
 			yield endOfLine;
+
+			const refName = CompilerDOM.toValidAssetId(prop.value.content, '_VLS_refs' as any);
+			options.templateRefNames.set(prop.value.content, refName);
+			return refName;
 		}
 	}
 }

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

@@ -15,6 +15,7 @@ export interface TemplateCodegenOptions {
 	template: NonNullable<Sfc['template']>;
 	scriptSetupBindingNames: Set<string>;
 	scriptSetupImportComponentNames: Set<string>;
+	templateRefNames: Map<string, string>;
 	hasDefineSlots?: boolean;
 	slotsAssignName?: string;
 	propsAssignName?: string;
@@ -49,8 +50,21 @@ export function* generateTemplate(options: TemplateCodegenOptions): Generator<Co
 
 	yield* ctx.generateAutoImportCompletion();
 
+	yield* generateRefs()
+
 	return ctx;
 
+	function* generateRefs(): Generator<Code> {
+		for (const [, validId] of options.templateRefNames) {
+			yield `let ${validId}${newLine}`;
+		}
+		yield `const __VLS_refs = {${newLine}`;
+		for (const [name, validId] of options.templateRefNames) {
+			yield `'${name}': ${validId}!,${newLine}`;
+		}
+		yield `}${endOfLine}`;
+	}
+
 	function* generateSlotsType(): Generator<Code> {
 		for (const { expVar, varName } of ctx.dynamicSlots) {
 			ctx.hasSlot = true;

packages/language-core/lib/parsers/scriptSetupRanges.ts

@@ -42,6 +42,10 @@ export function parseScriptSetupRanges(
 		name?: string;
 		inheritAttrs?: string;
 	} = {};
+	const templateRefs: {
+		name?: string;
+		define?: ReturnType<typeof parseDefineFunction>;
+	}[] = [];
 
 	const definePropProposalA = vueCompilerOptions.experimentalDefinePropProposal === 'kevinEdition' || ast.text.trimStart().startsWith('// @experimentalDefinePropProposal=kevinEdition');
 	const definePropProposalB = vueCompilerOptions.experimentalDefinePropProposal === 'johnsonEdition' || ast.text.trimStart().startsWith('// @experimentalDefinePropProposal=johnsonEdition');
@@ -104,6 +108,7 @@ export function parseScriptSetupRanges(
 		expose,
 		options,
 		defineProp,
+		templateRefs,
 	};
 
 	function _getStartEnd(node: ts.Node) {
@@ -298,6 +303,17 @@ export function parseScriptSetupRanges(
 						}
 					}
 				}
+			} else if (vueCompilerOptions.macros.templateRef.includes(callText) && node.arguments.length && !node.typeArguments?.length) {
+				const define = parseDefineFunction(node);
+				define.arg = _getStartEnd(node.arguments[0]);
+				let name;
+				if (ts.isVariableDeclaration(parent)) {
+					name = getNodeText(ts, parent.name, ast);
+				}
+				templateRefs.push({
+					name,
+					define
+				});
 			}
 		}
 		ts.forEachChild(node, child => {

packages/language-core/lib/plugins/vue-tsx.ts

@@ -94,6 +94,7 @@ function createTsx(
 			template: _sfc.template,
 			scriptSetupBindingNames: scriptSetupBindingNames(),
 			scriptSetupImportComponentNames: scriptSetupImportComponentNames(),
+			templateRefNames: new Map(),
 			hasDefineSlots: hasDefineSlots(),
 			slotsAssignName: slotsAssignName(),
 			propsAssignName: propsAssignName(),

packages/language-core/lib/types.ts

@@ -41,6 +41,7 @@ export interface VueCompilerOptions {
 		defineModel: string[];
 		defineOptions: string[];
 		withDefaults: string[];
+		templateRef: string[];
 	};
 	plugins: VueLanguagePlugin[];
 

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

@@ -230,6 +230,7 @@ export function resolveVueCompilerOptions(vueOptions: Partial<VueCompilerOptions
 			defineModel: ['defineModel'],
 			defineOptions: ['defineOptions'],
 			withDefaults: ['withDefaults'],
+			templateRef: ['templateRef', 'useTemplateRef'],
 			...vueOptions.macros,
 		},
 		plugins: vueOptions.plugins ?? [],

packages/language-core/schemas/vue-tsconfig.schema.json

@@ -79,7 +79,8 @@
 						"defineSlots": [ "defineSlots" ],
 						"defineEmits": [ "defineEmits" ],
 						"defineExpose": [ "defineExpose" ],
-						"withDefaults": [ "withDefaults" ]
+						"withDefaults": [ "withDefaults" ],
+						"templateRef": [ "templateRef", "useTemplateRef" ]
 					}
 				},
 				"experimentalResolveStyleCssClasses": {