From 815bfcffae7a9f04ee996367a731e6e072af6bd2 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Sat, 4 Sep 2021 16:42:39 -0400
Subject: [PATCH] fix(types): fix ref macro types

fix #4499
---
 packages/reactivity/src/computed.ts   |  3 ++
 packages/reactivity/src/ref.ts        |  3 +-
 packages/vue/ref-macros.d.ts          | 52 ++++++++++++++------
 test-dts/refTransformMacros.test-d.ts | 69 ++++++++++++++++++++++-----
 4 files changed, 99 insertions(+), 28 deletions(-)

diff --git a/packages/reactivity/src/computed.ts b/packages/reactivity/src/computed.ts
index 3b3734388a1..b3292588435 100644
--- a/packages/reactivity/src/computed.ts
+++ b/packages/reactivity/src/computed.ts
@@ -4,8 +4,11 @@ import { isFunction, NOOP } from '@vue/shared'
 import { ReactiveFlags, toRaw } from './reactive'
 import { Dep } from './dep'
 
+declare const ComoutedRefSymbol: unique symbol
+
 export interface ComputedRef<T = any> extends WritableComputedRef<T> {
   readonly value: T
+  [ComoutedRefSymbol]: true
 }
 
 export interface WritableComputedRef<T> extends Ref<T> {
diff --git a/packages/reactivity/src/ref.ts b/packages/reactivity/src/ref.ts
index 382afaa8016..32426aaf3e5 100644
--- a/packages/reactivity/src/ref.ts
+++ b/packages/reactivity/src/ref.ts
@@ -251,7 +251,8 @@ export interface RefUnwrapBailTypes {}
 export type ShallowUnwrapRef<T> = {
   [K in keyof T]: T[K] extends Ref<infer V>
     ? V
-    : T[K] extends Ref<infer V> | undefined // if `V` is `unknown` that means it does not extend `Ref` and is undefined
+    : // if `V` is `unknown` that means it does not extend `Ref` and is undefined
+    T[K] extends Ref<infer V> | undefined
     ? unknown extends V
       ? undefined
       : V | undefined
diff --git a/packages/vue/ref-macros.d.ts b/packages/vue/ref-macros.d.ts
index ee8ddfed9b9..c38b4bad719 100644
--- a/packages/vue/ref-macros.d.ts
+++ b/packages/vue/ref-macros.d.ts
@@ -4,18 +4,30 @@ import {
   ComputedRef,
   WritableComputedOptions,
   DebuggerOptions,
-  WritableComputedRef,
-  ShallowUnwrapRef
+  WritableComputedRef
 } from '@vue/runtime-dom'
 
-declare const RefMarker: unique symbol
-type RefValue<T> = T & { [RefMarker]?: any }
+declare const RefType: unique symbol
 
-declare const ComputedRefMarker: unique symbol
-type ComputedRefValue<T> = T & { [ComputedRefMarker]?: any }
+declare const enum RefTypes {
+  Ref = 1,
+  ComputedRef = 2,
+  WritableComputedRef = 3
+}
+
+type RefValue<T> = T extends null | undefined
+  ? T
+  : T & { [RefType]?: RefTypes.Ref }
+
+type ComputedRefValue<T> = T extends null | undefined
+  ? T
+  : T & { [RefType]?: RefTypes.ComputedRef }
+
+type WritableComputedRefValue<T> = T extends null | undefined
+  ? T
+  : T & { [RefType]?: RefTypes.WritableComputedRef }
 
-declare const WritableComputedRefMarker: unique symbol
-type WritableComputedRefValue<T> = T & { [WritableComputedRefMarker]?: any }
+type NormalObject<T extends object> = T & { [RefType]?: never }
 
 /**
  * Vue ref transform macro for binding refs as reactive variables.
@@ -23,25 +35,35 @@ type WritableComputedRefValue<T> = T & { [WritableComputedRefMarker]?: any }
 declare function _$<T>(arg: ComputedRef<T>): ComputedRefValue<T>
 declare function _$<T>(arg: WritableComputedRef<T>): WritableComputedRefValue<T>
 declare function _$<T>(arg: Ref<T>): RefValue<T>
-declare function _$<T extends object>(arg?: T): ShallowUnwrapRef<T>
+declare function _$<T extends object>(arg?: T): DestructureRefs<T>
+
+type DestructureRefs<T extends object> = {
+  [K in keyof T]: T[K] extends ComputedRef<infer V>
+    ? ComputedRefValue<V>
+    : T[K] extends WritableComputedRef<infer V>
+    ? WritableComputedRefValue<V>
+    : T[K] extends Ref<infer V>
+    ? RefValue<V>
+    : T[K]
+}
 
 /**
  * Vue ref transform macro for accessing underlying refs of reactive varaibles.
  */
+declare function _$$<T extends object>(arg: NormalObject<T>): ToRawRefs<T>
+declare function _$$<T>(value: RefValue<T>): Ref<T>
 declare function _$$<T>(value: ComputedRefValue<T>): ComputedRef<T>
 declare function _$$<T>(
   value: WritableComputedRefValue<T>
 ): WritableComputedRef<T>
-declare function _$$<T>(value: RefValue<T>): Ref<T>
-declare function _$$<T extends object>(arg: T): ToRawRefs<T>
 
 type ToRawRefs<T extends object> = {
-  [K in keyof T]: T[K] extends ComputedRefValue<infer V>
-    ? ComputedRefValue<V>
+  [K in keyof T]: T[K] extends RefValue<infer V>
+    ? Ref<V>
+    : T[K] extends ComputedRefValue<infer V>
+    ? ComputedRef<V>
     : T[K] extends WritableComputedRefValue<infer V>
     ? WritableComputedRef<V>
-    : T[K] extends RefValue<infer V>
-    ? Ref<V>
     : T[K] extends object
     ? T[K] extends
         | Function
diff --git a/test-dts/refTransformMacros.test-d.ts b/test-dts/refTransformMacros.test-d.ts
index dfabb55fb5e..f0abb30424f 100644
--- a/test-dts/refTransformMacros.test-d.ts
+++ b/test-dts/refTransformMacros.test-d.ts
@@ -1,14 +1,55 @@
 import { WritableComputedRef } from '@vue/reactivity'
-import { expectType, ref, Ref, ComputedRef } from './index'
+import { expectType, ref, computed, Ref, ComputedRef } from './index'
 import 'vue/ref-macros'
+import { RefType, RefTypes } from 'vue/ref-macros'
 
 // wrapping refs
+
 // normal
+let n = $(ref(1))
+n = 2
+// @ts-expect-error
+n = 'foo'
+
+// #4499 nullable
+let msg = $(ref<string | null>(null))
+msg = 'hello world'
+msg = null
+expectType<RefTypes.Ref | undefined>(msg![RefType])
+
 // computed
+let m = $(computed(() => n + 1))
+m * 1
+// @ts-expect-error
+m.slice()
+expectType<RefTypes.ComputedRef | undefined>(m[RefType])
+
 // writable computed
+let wc = $(
+  computed({
+    get: () => n + 1,
+    set: v => (n = v - 1)
+  })
+)
+wc = 2
+// @ts-expect-error
+wc = 'foo'
+expectType<RefTypes.WritableComputedRef | undefined>(wc[RefType])
 
 // destructure
-const { x, y, z } = $(useFoo())
+function useFoo() {
+  let x = $ref(1)
+  let y = $computed(() => 'hi')
+
+  return $$({
+    x,
+    y,
+    z: 123
+  })
+}
+
+const fooRes = useFoo()
+const { x, y, z } = $(fooRes)
 expectType<number>(x)
 expectType<string>(y)
 expectType<number>(z)
@@ -39,17 +80,12 @@ expectType<number>(
   })
 )
 
-function useFoo() {
-  return {
-    x: ref(1),
-    y: ref('hi'),
-    z: 123
-  }
-}
-
 // $$
-expectType<Ref<number>>($$(x))
-expectType<Ref<string>>($$(y))
+const xRef = $$(x)
+expectType<Ref<number>>(xRef)
+
+const yRef = $$(y)
+expectType<ComputedRef<string>>(yRef)
 
 const c = $computed(() => 1)
 const cRef = $$(c)
@@ -63,3 +99,12 @@ const c2Ref = $$(c2)
 expectType<WritableComputedRef<number>>(c2Ref)
 
 // $$ on object
+const obj = $$({
+  n,
+  m,
+  wc
+})
+
+expectType<Ref<number>>(obj.n)
+expectType<ComputedRef<number>>(obj.m)
+expectType<WritableComputedRef<number>>(obj.wc)