From 1462f5ed906dc2782bd9b2c9692638dcd9b81bf2 Mon Sep 17 00:00:00 2001
From: Johnson Chu <johnsoncodehk@gmail.com>
Date: Fri, 10 May 2024 13:44:57 +0800
Subject: [PATCH] feat(typescript-plugin): create script setup block when auto
 import if needed

---
 packages/language-core/lib/types.ts           | 13 --------
 .../lib/virtualFile/computedSfc.ts            |  3 --
 packages/language-service/index.ts            |  2 +-
 packages/typescript-plugin/index.ts           |  2 +-
 packages/typescript-plugin/lib/common.ts      | 30 ++++++++++++++++++-
 5 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/packages/language-core/lib/types.ts b/packages/language-core/lib/types.ts
index 21242209ff..ab814b279a 100644
--- a/packages/language-core/lib/types.ts
+++ b/packages/language-core/lib/types.ts
@@ -123,19 +123,6 @@ export interface Sfc {
 	customBlocks: readonly (SfcBlock & {
 		type: string;
 	})[];
-
-	/**
-	 * @deprecated use `template.ast` instead
-	 */
-	templateAst: CompilerDOM.RootNode | undefined;
-	/**
-	 * @deprecated use `script.ast` instead
-	 */
-	scriptAst: ts.SourceFile | undefined;
-	/**
-	 * @deprecated use `scriptSetup.ast` instead
-	 */
-	scriptSetupAst: ts.SourceFile | undefined;
 }
 
 export interface TextRange {
diff --git a/packages/language-core/lib/virtualFile/computedSfc.ts b/packages/language-core/lib/virtualFile/computedSfc.ts
index 71fc283c9f..7e2803089f 100644
--- a/packages/language-core/lib/virtualFile/computedSfc.ts
+++ b/packages/language-core/lib/virtualFile/computedSfc.ts
@@ -126,9 +126,6 @@ export function computedSfc(
 		get scriptSetup() { return scriptSetup(); },
 		get styles() { return styles; },
 		get customBlocks() { return customBlocks; },
-		get templateAst() { return template()?.ast; },
-		get scriptAst() { return script()?.ast; },
-		get scriptSetupAst() { return scriptSetup()?.ast; },
 	};
 
 	function computedTemplateAst(base: SfcBlock) {
diff --git a/packages/language-service/index.ts b/packages/language-service/index.ts
index e0c0a27c20..f1d89fd69f 100644
--- a/packages/language-service/index.ts
+++ b/packages/language-service/index.ts
@@ -55,7 +55,7 @@ export function getVueLanguageServicePlugins(
 						}
 						const languageService = (created.provide as import('volar-service-typescript').Provide)['typescript/languageService']();
 						const vueOptions = getVueOptions(context.env);
-						decorateLanguageServiceForVue(context.language, languageService, vueOptions, ts, false);
+						decorateLanguageServiceForVue(context.language, languageService, vueOptions, ts, false, fileName => context.env.typescript!.fileNameToUri(fileName));
 						return created;
 					},
 				};
diff --git a/packages/typescript-plugin/index.ts b/packages/typescript-plugin/index.ts
index 1bea049499..6bdf407992 100644
--- a/packages/typescript-plugin/index.ts
+++ b/packages/typescript-plugin/index.ts
@@ -64,7 +64,7 @@ function createLanguageServicePlugin(): ts.server.PluginModuleFactory {
 					projects.set(info.project, { info, language, vueOptions });
 
 					decorateLanguageService(language, info.languageService);
-					decorateLanguageServiceForVue(language, info.languageService, vueOptions, ts, true);
+					decorateLanguageServiceForVue(language, info.languageService, vueOptions, ts, true, fileName => fileName);
 					decorateLanguageServiceHost(ts, language, info.languageServiceHost);
 					startNamedPipeServer(ts, info.project.projectKind, info.project.getCurrentDirectory());
 
diff --git a/packages/typescript-plugin/lib/common.ts b/packages/typescript-plugin/lib/common.ts
index 8b2f83f0b5..39f73c77b2 100644
--- a/packages/typescript-plugin/lib/common.ts
+++ b/packages/typescript-plugin/lib/common.ts
@@ -9,6 +9,7 @@ export function decorateLanguageServiceForVue(
 	vueOptions: vue.VueCompilerOptions,
 	ts: typeof import('typescript'),
 	isTsPlugin: boolean,
+	getScriptId: (fileName: string) => string,
 ) {
 	const {
 		getCompletionsAtPosition,
@@ -49,6 +50,12 @@ export function decorateLanguageServiceForVue(
 							break;
 						}
 					}
+					if (item.data) {
+						// @ts-expect-error
+						item.data.__isAutoImport = {
+							fileName,
+						};
+					}
 				}
 			}
 		}
@@ -69,6 +76,27 @@ export function decorateLanguageServiceForVue(
 				}
 			}
 		}
+		// @ts-expect-error
+		if (args[6]?.__isAutoImport) {
+			// @ts-expect-error
+			const { fileName } = args[6]?.__isAutoImport;
+			const sourceScript = language.scripts.get(getScriptId(fileName));
+			if (sourceScript?.generated?.root instanceof vue.VueVirtualCode) {
+				const sfc = sourceScript.generated.root.getVueSfc();
+				if (!sfc?.descriptor.script && !sfc?.descriptor.scriptSetup) {
+					for (const codeAction of details?.codeActions ?? []) {
+						for (const change of codeAction.changes) {
+							for (const textChange of change.textChanges) {
+								textChange.newText = `<script setup lang="ts">${textChange.newText}</script>\n\n`;
+								break;
+							}
+							break;
+						}
+						break;
+					}
+				}
+			}
+		}
 		return details;
 	};
 	languageService.getCodeFixesAtPosition = (...args) => {
@@ -80,7 +108,7 @@ export function decorateLanguageServiceForVue(
 	if (isTsPlugin) {
 		languageService.getEncodedSemanticClassifications = (fileName, span, format) => {
 			const result = getEncodedSemanticClassifications(fileName, span, format);
-			const file = language.scripts.get(fileName);
+			const file = language.scripts.get(getScriptId(fileName));
 			if (file?.generated?.root instanceof vue.VueVirtualCode) {
 				const { template } = file.generated.root.sfc;
 				if (template) {