From 1cfab4c695b0c28f549f8c97faee5099581792a7 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Thu, 13 Apr 2023 10:28:22 +0800
Subject: [PATCH] feat(compiler-sfc): support limited built-in utility types in
 macros

---
 .../compileScript/resolveType.spec.ts         | 25 ++++++
 .../compiler-sfc/src/script/resolveType.ts    | 79 +++++++++++++++++--
 2 files changed, 97 insertions(+), 7 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/compileScript/resolveType.spec.ts b/packages/compiler-sfc/__tests__/compileScript/resolveType.spec.ts
index 3d5e3750798..3e61e1b231a 100644
--- a/packages/compiler-sfc/__tests__/compileScript/resolveType.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript/resolveType.spec.ts
@@ -190,6 +190,31 @@ describe('resolveType', () => {
     })
   })
 
+  test('utility type: Pick', () => {
+    expect(
+      resolve(`
+    type T = { foo: number, bar: string, baz: boolean }
+    type K = 'foo' | 'bar'
+    type Target = Pick<T, K>
+    `).elements
+    ).toStrictEqual({
+      foo: ['Number'],
+      bar: ['String']
+    })
+  })
+
+  test('utility type: Omit', () => {
+    expect(
+      resolve(`
+    type T = { foo: number, bar: string, baz: boolean }
+    type K = 'foo' | 'bar'
+    type Target = Omit<T, K>
+    `).elements
+    ).toStrictEqual({
+      baz: ['Boolean']
+    })
+  })
+
   describe('errors', () => {
     test('error on computed keys', () => {
       expect(() => resolve(`type Target = { [Foo]: string }`)).toThrow(
diff --git a/packages/compiler-sfc/src/script/resolveType.ts b/packages/compiler-sfc/src/script/resolveType.ts
index 101f3b4f0a8..0f6a4eb5230 100644
--- a/packages/compiler-sfc/src/script/resolveType.ts
+++ b/packages/compiler-sfc/src/script/resolveType.ts
@@ -72,8 +72,18 @@ function innerResolveTypeElements(
       if (resolved) {
         return resolveTypeElements(ctx, resolved)
       } else {
-        // TODO Pick / Omit
-        ctx.error(`Failed to resolved type reference`, node)
+        const typeName = getReferenceName(node)
+        if (
+          typeof typeName === 'string' &&
+          // @ts-ignore
+          SupportedBuiltinsSet.has(typeName)
+        ) {
+          return resolveBuiltin(ctx, node, typeName as any)
+        }
+        ctx.error(
+          `Failed to resolved type reference, or unsupported built-in utlility type.`,
+          node
+        )
       }
     }
     case 'TSUnionType':
@@ -290,17 +300,60 @@ function resolveTemplateKeys(
   return res
 }
 
+const SupportedBuiltinsSet = new Set([
+  'Partial',
+  'Required',
+  'Readonly',
+  'Pick',
+  'Omit'
+] as const)
+
+type GetSetType<T> = T extends Set<infer V> ? V : never
+
+function resolveBuiltin(
+  ctx: ScriptCompileContext,
+  node: TSTypeReference | TSExpressionWithTypeArguments,
+  name: GetSetType<typeof SupportedBuiltinsSet>
+): ResolvedElements {
+  const t = resolveTypeElements(ctx, node.typeParameters!.params[0])
+  switch (name) {
+    case 'Partial':
+    case 'Required':
+    case 'Readonly':
+      return t
+    case 'Pick': {
+      const picked = resolveStringType(ctx, node.typeParameters!.params[1])
+      const res: ResolvedElements = {}
+      if (t.__callSignatures) addCallSignature(res, t.__callSignatures)
+      for (const key of picked) {
+        res[key] = t[key]
+      }
+      return res
+    }
+    case 'Omit':
+      const omitted = resolveStringType(ctx, node.typeParameters!.params[1])
+      const res: ResolvedElements = {}
+      if (t.__callSignatures) addCallSignature(res, t.__callSignatures)
+      for (const key in t) {
+        if (!omitted.includes(key)) {
+          res[key] = t[key]
+        }
+      }
+      return res
+  }
+}
+
 function resolveTypeReference(
   ctx: ScriptCompileContext,
   node: TSTypeReference | TSExpressionWithTypeArguments,
   scope = getRootScope(ctx)
 ): Node | undefined {
-  const ref = node.type === 'TSTypeReference' ? node.typeName : node.expression
-  if (ref.type === 'Identifier') {
-    if (scope.imports[ref.name]) {
+  const name = getReferenceName(node)
+  if (typeof name === 'string') {
+    if (scope.imports[name]) {
       // TODO external import
-    } else if (scope.types[ref.name]) {
-      return scope.types[ref.name]
+    } else if (scope.types[name]) {
+      return scope.types[name]
     }
   } else {
     // TODO qualified name, e.g. Foo.Bar
@@ -308,6 +361,18 @@ function resolveTypeReference(
   }
 }
 
+function getReferenceName(
+  node: TSTypeReference | TSExpressionWithTypeArguments
+): string | string[] {
+  const ref = node.type === 'TSTypeReference' ? node.typeName : node.expression
+  if (ref.type === 'Identifier') {
+    return ref.name
+  } else {
+    // TODO qualified name, e.g. Foo.Bar
+    return []
+  }
+}
+
 function getRootScope(ctx: ScriptCompileContext): TypeScope {
   if (ctx.scope) {
     return ctx.scope