From 2de2ea27239829c8fb7e04032000999d74069c64 Mon Sep 17 00:00:00 2001
From: so1ve <i@mk1.io>
Date: Tue, 7 Nov 2023 11:09:38 +0800
Subject: [PATCH] feat: [wip] `defineEmits` support

---
 src/core/language.ts                          | 57 ++++++++++----
 src/core/printer.ts                           |  6 +-
 src/core/transform.ts                         | 76 +++++++++++++++++--
 test/__fixtures__/defineEmits/basic-ts.ts     |  4 +
 test/__fixtures__/defineEmits/basic.vue       |  5 ++
 .../{ => defineProps}/basic-ts.ts             |  0
 test/__fixtures__/{ => defineProps}/basic.vue |  0
 .../{ => defineProps}/core#8286.vue           |  0
 .../{ => defineProps}/core#8468-props.ts      |  0
 .../{ => defineProps}/core#8468.vue           |  0
 .../import-from-vue-vue.exclude.vue           |  0
 .../{ => defineProps}/import-from-vue.vue     |  0
 .../fixtures-compiled.test.ts.snap            | 27 ++++++-
 test/__snapshots__/fixtures.test.ts.snap      | 17 ++++-
 test/fixtures-compiled.test.ts                |  2 +-
 test/fixtures.test.ts                         |  2 +-
 16 files changed, 162 insertions(+), 34 deletions(-)
 create mode 100644 test/__fixtures__/defineEmits/basic-ts.ts
 create mode 100644 test/__fixtures__/defineEmits/basic.vue
 rename test/__fixtures__/{ => defineProps}/basic-ts.ts (100%)
 rename test/__fixtures__/{ => defineProps}/basic.vue (100%)
 rename test/__fixtures__/{ => defineProps}/core#8286.vue (100%)
 rename test/__fixtures__/{ => defineProps}/core#8468-props.ts (100%)
 rename test/__fixtures__/{ => defineProps}/core#8468.vue (100%)
 rename test/__fixtures__/{ => defineProps}/import-from-vue-vue.exclude.vue (100%)
 rename test/__fixtures__/{ => defineProps}/import-from-vue.vue (100%)

diff --git a/src/core/language.ts b/src/core/language.ts
index b5f3ebc..d2febdb 100644
--- a/src/core/language.ts
+++ b/src/core/language.ts
@@ -114,10 +114,19 @@ function baseCreateLanguageWorker(
 	const program = tsLs.getProgram()!;
 	const typeChecker = program.getTypeChecker();
 
-	function findScriptRangesNode(
+	function findScriptRangesAndNode(
 		filepath: string,
 		filter: (ranges: vue.ScriptSetupRanges) => vue.TextRange | undefined,
-	): { node: ts.Node; range: vue.TextRange } | undefined {
+	):
+		| {
+				virtualFileNode: ts.Node;
+				scriptNode: ts.Node;
+				setupRange: vue.TextRange;
+				scriptSetupAst: ts.SourceFile;
+				virtualFileAst: ts.SourceFile;
+				offset: number;
+		  }
+		| undefined {
 		filepath = normalizePath(filepath);
 		const sourceFile = core.virtualFiles.getSource(filepath)?.root;
 		if (!(sourceFile instanceof vue.VueFile)) {
@@ -128,7 +137,7 @@ function baseCreateLanguageWorker(
 		}
 		const {
 			ast: scriptSetupAst,
-			startTagEnd,
+			startTagEnd: offset,
 			lang,
 		} = sourceFile.sfc.scriptSetup;
 		if (lang !== "ts" && lang !== "tsx") {
@@ -137,8 +146,8 @@ function baseCreateLanguageWorker(
 		// We get volar's generated virtual ts file because `scriptSetupAst`
 		// is the ast of `<script setup>` block, which is not in the program
 		// and cannot be used by type checker.
-		const virtualTsFile = program.getSourceFile(`${filepath}.${lang}`);
-		if (!virtualTsFile) {
+		const virtualFileAst = program.getSourceFile(`${filepath}.${lang}`);
+		if (!virtualFileAst) {
 			return;
 		}
 		const scriptSetupRanges = vue.parseScriptSetupRanges(
@@ -148,7 +157,7 @@ function baseCreateLanguageWorker(
 		);
 		const virtualTsFileScriptSetupRanges = vue.parseScriptSetupRanges(
 			ts,
-			virtualTsFile,
+			virtualFileAst,
 			vueCompilerOptions,
 		);
 		const sourceRange = filter(scriptSetupRanges);
@@ -156,28 +165,46 @@ function baseCreateLanguageWorker(
 		if (!sourceRange || !virtualTsRange) {
 			return;
 		}
-		let foundNode!: ts.Node;
-		virtualTsFile.forEachChild(function traverse(node: ts.Node) {
+		let foundVirtualFileNode!: ts.Node;
+		virtualFileAst.forEachChild(function traverse(node: ts.Node) {
 			if (
-				node.getStart(virtualTsFile) === virtualTsRange.start &&
+				node.getStart(virtualFileAst) === virtualTsRange.start &&
 				node.getEnd() === virtualTsRange.end
 			) {
-				foundNode = node;
+				foundVirtualFileNode = node;
+			} else {
+				node.forEachChild(traverse);
+			}
+		});
+		let foundSetupNode!: ts.Node;
+		scriptSetupAst.forEachChild(function traverse(node: ts.Node) {
+			if (
+				node.getStart(scriptSetupAst) === sourceRange.start &&
+				node.getEnd() === sourceRange.end
+			) {
+				foundSetupNode = node;
 			} else {
 				node.forEachChild(traverse);
 			}
 		});
 
-		const range: vue.TextRange = {
-			start: startTagEnd + sourceRange.start,
-			end: startTagEnd + sourceRange.end,
+		const setupRange: vue.TextRange = {
+			start: offset + sourceRange.start,
+			end: offset + sourceRange.end,
 		};
 
-		return { node: foundNode, range };
+		return {
+			virtualFileNode: foundVirtualFileNode,
+			scriptNode: foundSetupNode,
+			setupRange,
+			scriptSetupAst,
+			virtualFileAst,
+			offset,
+		};
 	}
 
 	return {
-		findScriptRangesNode,
+		findScriptRangesNode: findScriptRangesAndNode,
 		__internal__: {
 			tsLs,
 			program,
diff --git a/src/core/printer.ts b/src/core/printer.ts
index 8152e97..e37f023 100644
--- a/src/core/printer.ts
+++ b/src/core/printer.ts
@@ -16,11 +16,11 @@ export class Printer {
 	}
 
 	private printIntersectionTypeNode(node: ts.IntersectionTypeNode): string {
-		return node.types.map((t) => this.print(t)).join(" & ");
+		return node.types.map((t) => this.printTypeArg(t)).join(" & ");
 	}
 
 	private printUnionTypeNode(node: ts.UnionTypeNode): string {
-		return node.types.map((t) => this.print(t)).join(" | ");
+		return node.types.map((t) => this.printTypeArg(t)).join(" | ");
 	}
 
 	private printTypeLiteralNode(node: ts.TypeLiteralNode): string {
@@ -75,7 +75,7 @@ export class Printer {
 		return parts.join("\n");
 	}
 
-	public print(node: ts.Node): string {
+	public printTypeArg(node: ts.Node): string {
 		// Intersection and Union
 		if (ts.isIntersectionTypeNode(node)) {
 			return this.printIntersectionTypeNode(node);
diff --git a/src/core/transform.ts b/src/core/transform.ts
index 4361458..bc45cb8 100644
--- a/src/core/transform.ts
+++ b/src/core/transform.ts
@@ -1,14 +1,12 @@
 import MagicString from "magic-string";
+import ts from "typescript";
 import type { TransformResult } from "unplugin";
 
 import { getLanguage } from "./language";
 import { Printer } from "./printer";
 
-export function transform(code: string, id: string): TransformResult {
+function transformDefineProps(printer: Printer, s: MagicString, id: string) {
 	const language = getLanguage();
-	const s = new MagicString(code);
-	const typeChecker = language.__internal__.typeChecker;
-	const printer = new Printer(typeChecker);
 	const definePropsTypeArg = language.findScriptRangesNode(
 		id,
 		(scriptSetupRanges) => scriptSetupRanges.props.define?.typeArg,
@@ -18,11 +16,77 @@ export function transform(code: string, id: string): TransformResult {
 		return;
 	}
 
-	const { node: typeArgNode, range: typeArgRange } = definePropsTypeArg;
+	const { virtualFileNode: typeArgNode, setupRange: typeArgRange } =
+		definePropsTypeArg;
 
-	const printedType = printer.print(typeArgNode);
+	const printedType = printer.printTypeArg(typeArgNode);
 
 	s.overwrite(typeArgRange.start, typeArgRange.end, printedType);
+}
+
+function transformDefineEmits(printer: Printer, s: MagicString, id: string) {
+	const language = getLanguage();
+	const defineEmits = language.findScriptRangesNode(
+		id,
+		(scriptSetupRanges) => scriptSetupRanges.emits.define,
+	);
+	if (!defineEmits) {
+		return;
+	}
+
+	const {
+		scriptNode: defineEmitsNode,
+		virtualFileNode: virtualFileDefineEmitsNode,
+		scriptSetupAst,
+		offset,
+	} = defineEmits;
+
+	// TODO: refactor when https://github.com/vuejs/language-tools/pull/3710 is merged
+	const defineEmitsTypeArg =
+		virtualFileDefineEmitsNode &&
+		ts.isCallExpression(virtualFileDefineEmitsNode) &&
+		virtualFileDefineEmitsNode.typeArguments?.[0];
+
+	if (!defineEmitsTypeArg) {
+		return;
+	}
+
+	const tokens = defineEmitsNode.getChildren(scriptSetupAst);
+
+	// defineEmits<     Arg   >()
+	//            ^           ^
+	const lessThanToken = tokens.find(
+		(t) => t.kind === ts.SyntaxKind.LessThanToken,
+	);
+	const greaterThanToken = tokens.find(
+		(t) => t.kind === ts.SyntaxKind.GreaterThanToken,
+	);
+	// defineEmits<     Arg   >()
+	//                         ^
+	const openParenToken = tokens.find(
+		(t) => t.kind === ts.SyntaxKind.OpenParenToken,
+	);
+
+	if (!lessThanToken || !greaterThanToken || !openParenToken) {
+		return;
+	}
+
+	const defineEmitsTypeArgRange = [
+		offset + lessThanToken.pos,
+		offset + greaterThanToken.end,
+	] as const;
+
+	// Remove the type argument
+	s.remove(...defineEmitsTypeArgRange);
+}
+
+export function transform(code: string, id: string): TransformResult {
+	const s = new MagicString(code);
+	const language = getLanguage();
+	const typeChecker = language.__internal__.typeChecker;
+	const printer = new Printer(typeChecker);
+	transformDefineProps(printer, s, id);
+	transformDefineEmits(printer, s, id);
 
 	return {
 		code: s.toString(),
diff --git a/test/__fixtures__/defineEmits/basic-ts.ts b/test/__fixtures__/defineEmits/basic-ts.ts
new file mode 100644
index 0000000..f930356
--- /dev/null
+++ b/test/__fixtures__/defineEmits/basic-ts.ts
@@ -0,0 +1,4 @@
+export interface Events {
+	(event: "foo"): void;
+	(event: 1 extends 2 ? "bar" : "baz"): void;
+}
diff --git a/test/__fixtures__/defineEmits/basic.vue b/test/__fixtures__/defineEmits/basic.vue
new file mode 100644
index 0000000..e1be167
--- /dev/null
+++ b/test/__fixtures__/defineEmits/basic.vue
@@ -0,0 +1,5 @@
+<script setup lang="ts">
+import type { Events } from "./basic-ts";
+
+defineEmits<Events>();
+</script>
diff --git a/test/__fixtures__/basic-ts.ts b/test/__fixtures__/defineProps/basic-ts.ts
similarity index 100%
rename from test/__fixtures__/basic-ts.ts
rename to test/__fixtures__/defineProps/basic-ts.ts
diff --git a/test/__fixtures__/basic.vue b/test/__fixtures__/defineProps/basic.vue
similarity index 100%
rename from test/__fixtures__/basic.vue
rename to test/__fixtures__/defineProps/basic.vue
diff --git a/test/__fixtures__/core#8286.vue b/test/__fixtures__/defineProps/core#8286.vue
similarity index 100%
rename from test/__fixtures__/core#8286.vue
rename to test/__fixtures__/defineProps/core#8286.vue
diff --git a/test/__fixtures__/core#8468-props.ts b/test/__fixtures__/defineProps/core#8468-props.ts
similarity index 100%
rename from test/__fixtures__/core#8468-props.ts
rename to test/__fixtures__/defineProps/core#8468-props.ts
diff --git a/test/__fixtures__/core#8468.vue b/test/__fixtures__/defineProps/core#8468.vue
similarity index 100%
rename from test/__fixtures__/core#8468.vue
rename to test/__fixtures__/defineProps/core#8468.vue
diff --git a/test/__fixtures__/import-from-vue-vue.exclude.vue b/test/__fixtures__/defineProps/import-from-vue-vue.exclude.vue
similarity index 100%
rename from test/__fixtures__/import-from-vue-vue.exclude.vue
rename to test/__fixtures__/defineProps/import-from-vue-vue.exclude.vue
diff --git a/test/__fixtures__/import-from-vue.vue b/test/__fixtures__/defineProps/import-from-vue.vue
similarity index 100%
rename from test/__fixtures__/import-from-vue.vue
rename to test/__fixtures__/defineProps/import-from-vue.vue
diff --git a/test/__snapshots__/fixtures-compiled.test.ts.snap b/test/__snapshots__/fixtures-compiled.test.ts.snap
index a1581b9..fc02985 100644
--- a/test/__snapshots__/fixtures-compiled.test.ts.snap
+++ b/test/__snapshots__/fixtures-compiled.test.ts.snap
@@ -1,6 +1,25 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
-exports[`fixtures compiled > __fixtures__/basic.vue 1`] = `
+exports[`fixtures compiled > __fixtures__/defineEmits/basic.vue 1`] = `
+"import { defineComponent } from 'vue';
+import _export_sfc from '[NULL]/plugin-vue/export-helper';
+
+\\"use strict\\";
+var _sfc_main = /* @__PURE__ */ defineComponent({
+  __name: \\"basic\\",
+  setup(__props) {
+    return () => {
+    };
+  }
+});
+
+var basic = /* @__PURE__ */ _export_sfc(_sfc_main, [__FILE__]);
+
+export { basic as default };
+"
+`;
+
+exports[`fixtures compiled > __fixtures__/defineProps/basic.vue 1`] = `
 "import { defineComponent } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
 
@@ -36,7 +55,7 @@ export { basic as default };
 "
 `;
 
-exports[`fixtures compiled > __fixtures__/core#8286.vue 1`] = `
+exports[`fixtures compiled > __fixtures__/defineProps/core#8286.vue 1`] = `
 "import { defineComponent } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
 
@@ -239,7 +258,7 @@ export { core_8286 as default };
 "
 `;
 
-exports[`fixtures compiled > __fixtures__/core#8468.vue 1`] = `
+exports[`fixtures compiled > __fixtures__/defineProps/core#8468.vue 1`] = `
 "import { defineComponent } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
 
@@ -262,7 +281,7 @@ export { core_8468 as default };
 "
 `;
 
-exports[`fixtures compiled > __fixtures__/import-from-vue.vue 1`] = `
+exports[`fixtures compiled > __fixtures__/defineProps/import-from-vue.vue 1`] = `
 "import { defineComponent } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
 
diff --git a/test/__snapshots__/fixtures.test.ts.snap b/test/__snapshots__/fixtures.test.ts.snap
index 7db9bbe..b429b42 100644
--- a/test/__snapshots__/fixtures.test.ts.snap
+++ b/test/__snapshots__/fixtures.test.ts.snap
@@ -1,6 +1,15 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
-exports[`fixtures > __fixtures__/basic.vue 1`] = `
+exports[`fixtures > __fixtures__/defineEmits/basic.vue 1`] = `
+"<script setup lang=\\"ts\\">
+import type { Events } from \\"./basic-ts\\";
+
+defineEmits();
+</script>
+"
+`;
+
+exports[`fixtures > __fixtures__/defineProps/basic.vue 1`] = `
 "<script
 	setup
 	lang=\\"ts\\"
@@ -36,7 +45,7 @@ typeLiteral: TypeLiteral
 "
 `;
 
-exports[`fixtures > __fixtures__/core#8286.vue 1`] = `
+exports[`fixtures > __fixtures__/defineProps/core#8286.vue 1`] = `
 "<script setup lang=\\"ts\\">
 import { ButtonHTMLAttributes } from \\"vue\\";
 
@@ -230,7 +239,7 @@ onTransitionstart: any
 "
 `;
 
-exports[`fixtures > __fixtures__/core#8468.vue 1`] = `
+exports[`fixtures > __fixtures__/defineProps/core#8468.vue 1`] = `
 "<script setup lang=\\"ts\\" generic=\\"P extends Inputs\\">
 import type { Inputs } from \\"./core#8468-props.ts\\";
 
@@ -242,7 +251,7 @@ value: string | number
 "
 `;
 
-exports[`fixtures > __fixtures__/import-from-vue.vue 1`] = `
+exports[`fixtures > __fixtures__/defineProps/import-from-vue.vue 1`] = `
 "<script setup lang=\\"ts\\" generic=\\"T extends string\\">
 import type { Foo } from \\"./import-from-vue-vue.exclude.vue\\";
 
diff --git a/test/fixtures-compiled.test.ts b/test/fixtures-compiled.test.ts
index 4c7fb6d..c1ff2ed 100644
--- a/test/fixtures-compiled.test.ts
+++ b/test/fixtures-compiled.test.ts
@@ -16,7 +16,7 @@ import VueComplexTypes from "../src/rollup";
 
 describe("fixtures compiled", async () => {
 	await testFixtures(
-		["__fixtures__/*.vue", "!__fixtures__/*.exclude.vue"],
+		["__fixtures__/**/*.vue", "!__fixtures__/**/*.exclude.vue"],
 		(_args, id) =>
 			rollupBuild(id, [
 				VueComplexTypes({
diff --git a/test/fixtures.test.ts b/test/fixtures.test.ts
index 4633dae..d4e1b17 100644
--- a/test/fixtures.test.ts
+++ b/test/fixtures.test.ts
@@ -13,7 +13,7 @@ beforeAll(() => {
 
 describe("fixtures", async () => {
 	await testFixtures(
-		["__fixtures__/*.vue", "!__fixtures__/*.exclude.vue"],
+		["__fixtures__/**/*.vue", "!__fixtures__/**/*.exclude.vue"],
 		async (_args, id) => {
 			const text = await readFile(id, { encoding: "utf-8" });