From 94a60f3999d505c370cdae96bca6246c1b856a43 Mon Sep 17 00:00:00 2001
From: rudyxu1102 <xuhaodong66@foxmail.com>
Date: Tue, 17 Oct 2023 22:44:58 +0800
Subject: [PATCH] feat(compiler-sfc): support defineAttrs macro

---
 .../__snapshots__/defineAttrs.spec.ts.snap    | 46 +++++++++++++++++++
 .../compileScript/defineAttrs.spec.ts         | 40 ++++++++++++++++
 packages/compiler-sfc/src/compileScript.ts    |  7 ++-
 packages/compiler-sfc/src/script/context.ts   |  1 +
 .../compiler-sfc/src/script/defineAttrs.ts    | 33 +++++++++++++
 packages/runtime-core/src/apiSetupHelpers.ts  | 12 ++++-
 packages/runtime-core/src/componentOptions.ts | 11 +++++
 packages/runtime-core/src/index.ts            |  1 +
 .../types/scriptSetupHelpers.d.ts             |  2 +
 9 files changed, 150 insertions(+), 3 deletions(-)
 create mode 100644 packages/compiler-sfc/__tests__/compileScript/__snapshots__/defineAttrs.spec.ts.snap
 create mode 100644 packages/compiler-sfc/__tests__/compileScript/defineAttrs.spec.ts
 create mode 100644 packages/compiler-sfc/src/script/defineAttrs.ts

diff --git a/packages/compiler-sfc/__tests__/compileScript/__snapshots__/defineAttrs.spec.ts.snap b/packages/compiler-sfc/__tests__/compileScript/__snapshots__/defineAttrs.spec.ts.snap
new file mode 100644
index 00000000000..80d066e8a67
--- /dev/null
+++ b/packages/compiler-sfc/__tests__/compileScript/__snapshots__/defineAttrs.spec.ts.snap
@@ -0,0 +1,46 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`defineAttrs() > basic usage 1`] = `
+"import { useAttrs as _useAttrs, defineComponent as _defineComponent } from 'vue'
+
+export default /*#__PURE__*/_defineComponent({
+  setup(__props, { expose: __expose }) {
+  __expose();
+
+      const attrs = _useAttrs()
+      
+return { attrs }
+}
+
+})"
+`;
+
+exports[`defineAttrs() > w/o generic params 1`] = `
+"import { useAttrs as _useAttrs } from 'vue'
+
+export default {
+  setup(__props, { expose: __expose }) {
+  __expose();
+
+      const attrs = _useAttrs()
+      
+return { attrs }
+}
+
+}"
+`;
+
+exports[`defineAttrs() > w/o return value 1`] = `
+"import { defineComponent as _defineComponent } from 'vue'
+
+export default /*#__PURE__*/_defineComponent({
+  setup(__props, { expose: __expose }) {
+  __expose();
+
+      
+      
+return {  }
+}
+
+})"
+`;
diff --git a/packages/compiler-sfc/__tests__/compileScript/defineAttrs.spec.ts b/packages/compiler-sfc/__tests__/compileScript/defineAttrs.spec.ts
new file mode 100644
index 00000000000..e155e3c6bfd
--- /dev/null
+++ b/packages/compiler-sfc/__tests__/compileScript/defineAttrs.spec.ts
@@ -0,0 +1,40 @@
+import { compileSFCScript as compile, assertCode } from '../utils'
+
+describe('defineAttrs()', () => {
+  test('basic usage', () => {
+    const { content } = compile(`
+      <script setup lang="ts">
+      const attrs = defineAttrs<{
+        bar?: number
+      }>()
+      </script>
+    `)
+    assertCode(content)
+    expect(content).toMatch(`const attrs = _useAttrs()`)
+    expect(content).not.toMatch('defineAttrs')
+  })
+
+  test('w/o return value', () => {
+    const { content } = compile(`
+      <script setup lang="ts">
+      defineAttrs<{
+        bar?: number
+      }>()
+      </script>
+    `)
+    assertCode(content)
+    expect(content).not.toMatch('defineAttrs')
+    expect(content).not.toMatch(`_useAttrs`)
+  })
+
+  test('w/o generic params', () => {
+    const { content } = compile(`
+      <script setup>
+      const attrs = defineAttrs()
+      </script>
+    `)
+    assertCode(content)
+    expect(content).toMatch(`const attrs = _useAttrs()`)
+    expect(content).not.toMatch('defineAttrs')
+  })
+})
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index cfcc607c72d..fe51563dfd2 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -53,6 +53,7 @@ import {
 import { analyzeScriptBindings } from './script/analyzeScriptBindings'
 import { isImportUsed } from './script/importUsageCheck'
 import { processAwait } from './script/topLevelAwait'
+import { processDefineAttrs } from './script/defineAttrs'
 
 export interface SFCScriptCompileOptions {
   /**
@@ -512,7 +513,8 @@ export function compileScript(
         processDefineProps(ctx, expr) ||
         processDefineEmits(ctx, expr) ||
         processDefineOptions(ctx, expr) ||
-        processDefineSlots(ctx, expr)
+        processDefineSlots(ctx, expr) ||
+        processDefineAttrs(ctx, expr)
       ) {
         ctx.s.remove(node.start! + startOffset, node.end! + startOffset)
       } else if (processDefineExpose(ctx, expr)) {
@@ -550,7 +552,8 @@ export function compileScript(
             !isDefineProps && processDefineEmits(ctx, init, decl.id)
           !isDefineEmits &&
             (processDefineSlots(ctx, init, decl.id) ||
-              processDefineModel(ctx, init, decl.id))
+              processDefineModel(ctx, init, decl.id) ||
+              processDefineAttrs(ctx, init, decl.id))
 
           if (
             isDefineProps &&
diff --git a/packages/compiler-sfc/src/script/context.ts b/packages/compiler-sfc/src/script/context.ts
index 5fe09d28a42..030370ebbad 100644
--- a/packages/compiler-sfc/src/script/context.ts
+++ b/packages/compiler-sfc/src/script/context.ts
@@ -36,6 +36,7 @@ export class ScriptCompileContext {
   hasDefineOptionsCall = false
   hasDefineSlotsCall = false
   hasDefineModelCall = false
+  hasDefineAttrsCall = false
 
   // defineProps
   propsCall: CallExpression | undefined
diff --git a/packages/compiler-sfc/src/script/defineAttrs.ts b/packages/compiler-sfc/src/script/defineAttrs.ts
new file mode 100644
index 00000000000..d2eda8d9962
--- /dev/null
+++ b/packages/compiler-sfc/src/script/defineAttrs.ts
@@ -0,0 +1,33 @@
+import { LVal, Node } from '@babel/types'
+import { isCallOf } from './utils'
+import { ScriptCompileContext } from './context'
+
+export const DEFINE_ATTRS = 'defineAttrs'
+
+export function processDefineAttrs(
+  ctx: ScriptCompileContext,
+  node: Node,
+  declId?: LVal
+): boolean {
+  if (!isCallOf(node, DEFINE_ATTRS)) {
+    return false
+  }
+  if (ctx.hasDefineAttrsCall) {
+    ctx.error(`duplicate ${DEFINE_ATTRS}() call`, node)
+  }
+  ctx.hasDefineAttrsCall = true
+
+  if (node.arguments.length > 0) {
+    ctx.error(`${DEFINE_ATTRS}() cannot accept arguments`, node)
+  }
+
+  if (declId) {
+    ctx.s.overwrite(
+      ctx.startOffset! + node.start!,
+      ctx.startOffset! + node.end!,
+      `${ctx.helper('useAttrs')}()`
+    )
+  }
+
+  return true
+}
diff --git a/packages/runtime-core/src/apiSetupHelpers.ts b/packages/runtime-core/src/apiSetupHelpers.ts
index 93200667081..b673d20e7f8 100644
--- a/packages/runtime-core/src/apiSetupHelpers.ts
+++ b/packages/runtime-core/src/apiSetupHelpers.ts
@@ -18,7 +18,8 @@ import {
   ComponentOptionsMixin,
   ComponentOptionsWithoutProps,
   ComputedOptions,
-  MethodOptions
+  MethodOptions,
+  StrictUnwrapAttrsType
 } from './componentOptions'
 import {
   ComponentPropsOptions,
@@ -215,6 +216,15 @@ export function defineSlots<
   return null as any
 }
 
+export function defineAttrs<
+  Attrs extends Record<string, any> = Record<string, any>
+>(): StrictUnwrapAttrsType<Attrs> {
+  if (__DEV__) {
+    warnRuntimeUsage(`defineAttrs`)
+  }
+  return null as any
+}
+
 /**
  * (**Experimental**) Vue `<script setup>` compiler macro for declaring a
  * two-way binding prop that can be consumed via `v-model` from the parent
diff --git a/packages/runtime-core/src/componentOptions.ts b/packages/runtime-core/src/componentOptions.ts
index de4d304448a..3b3516ed5a2 100644
--- a/packages/runtime-core/src/componentOptions.ts
+++ b/packages/runtime-core/src/componentOptions.ts
@@ -423,6 +423,17 @@ export type ExtractComputedReturns<T extends any> = {
     : never
 }
 
+export type StrictUnwrapAttrsType<
+  Attrs extends Record<string, unknown>,
+  T = NonNullable<Attrs>
+> = [keyof T] extends [never]
+  ? Record<string, unknown>
+  : T extends new () => { $props: infer P }
+  ? Readonly<NonNullable<P>>
+  : T extends (props: infer P, ...args: any) => any
+  ? Readonly<NonNullable<P>>
+  : Readonly<T>
+
 export type ObjectWatchOptionItem = {
   handler: WatchCallback | string
 } & WatchOptions
diff --git a/packages/runtime-core/src/index.ts b/packages/runtime-core/src/index.ts
index 0d5fc39d5f6..27a2f106365 100644
--- a/packages/runtime-core/src/index.ts
+++ b/packages/runtime-core/src/index.ts
@@ -71,6 +71,7 @@ export {
   defineExpose,
   defineOptions,
   defineSlots,
+  defineAttrs,
   defineModel,
   withDefaults,
   useModel,
diff --git a/packages/runtime-core/types/scriptSetupHelpers.d.ts b/packages/runtime-core/types/scriptSetupHelpers.d.ts
index a9ba734a53d..9224ebb9a5b 100644
--- a/packages/runtime-core/types/scriptSetupHelpers.d.ts
+++ b/packages/runtime-core/types/scriptSetupHelpers.d.ts
@@ -7,6 +7,7 @@ type _defineOptions = typeof defineOptions
 type _defineSlots = typeof defineSlots
 type _defineModel = typeof defineModel
 type _withDefaults = typeof withDefaults
+type _defineAttrs = typeof defineAttrs
 
 declare global {
   const defineProps: _defineProps
@@ -16,4 +17,5 @@ declare global {
   const defineSlots: _defineSlots
   const defineModel: _defineModel
   const withDefaults: _withDefaults
+  const defineAttrs: _defineAttrs
 }