From c0cb8bfd96b9096c603b12333fbd51bf84f38b8c Mon Sep 17 00:00:00 2001
From: jens <jens@doellmann.com>
Date: Thu, 25 Nov 2021 17:23:35 +0100
Subject: [PATCH] fix(types): should unwrap refs in collections when reading
 objects

---
 packages/reactivity/src/index.ts   |  3 +
 packages/reactivity/src/ref.ts     | 34 +++++++++--
 packages/runtime-core/src/index.ts |  3 +
 test-dts/ref.test-d.ts             | 95 ++++++++++++++++++++++++++++++
 test-dts/tsconfig.build.json       |  1 +
 test-dts/tsconfig.json             |  3 +-
 6 files changed, 134 insertions(+), 5 deletions(-)

diff --git a/packages/reactivity/src/index.ts b/packages/reactivity/src/index.ts
index 676f8598fb3..63e067247ca 100644
--- a/packages/reactivity/src/index.ts
+++ b/packages/reactivity/src/index.ts
@@ -13,6 +13,9 @@ export {
   ToRefs,
   UnwrapRef,
   ShallowRef,
+  UnwrappedMap,
+  UnwrappedWeakMap,
+  UnwrappedSet,
   ShallowUnwrapRef,
   RefUnwrapBailTypes,
   CustomRefFactory
diff --git a/packages/reactivity/src/ref.ts b/packages/reactivity/src/ref.ts
index 55059485d6c..69ca2be0402 100644
--- a/packages/reactivity/src/ref.ts
+++ b/packages/reactivity/src/ref.ts
@@ -66,9 +66,6 @@ export function isRef(r: any): r is Ref {
   return Boolean(r && r.__v_isRef === true)
 }
 
-export function ref<T extends object>(
-  value: T
-): [T] extends [Ref] ? T : Ref<UnwrapRef<T>>
 export function ref<T>(value: T): Ref<UnwrapRef<T>>
 export function ref<T = any>(): Ref<T | undefined>
 export function ref(value?: unknown) {
@@ -287,13 +284,42 @@ export type UnwrapRef<T> = T extends ShallowRef<infer V>
   ? UnwrapRefSimple<V>
   : UnwrapRefSimple<T>
 
+export interface UnwrappedMap<K, V>
+  extends Omit<Map<K, UnwrapRefSimple<V>>, 'set'> {
+  set(key: K, value: V): this
+}
+
+export interface UnwrappedWeakMap<K extends object, V>
+  extends Omit<WeakMap<K, UnwrapRefSimple<V>>, 'set'> {
+  set(key: K, value: V): this
+}
+
+export interface UnwrappedSet<T>
+  extends Omit<Set<UnwrapRefSimple<T>>, 'add' | 'has' | 'delete'> {
+  add(value: T): this
+  has(value: T): boolean
+  delete(value: T): boolean
+}
+
+type UnwrapRefCollectionTypes<T extends CollectionTypes> = T extends Map<
+  infer TMapKey,
+  infer TMap
+>
+  ? UnwrappedMap<TMapKey, TMap>
+  : T extends WeakMap<infer TWeakMapKey, infer TWeakMap>
+  ? UnwrappedWeakMap<TWeakMapKey, TWeakMap>
+  : T extends Set<infer TSet>
+  ? UnwrappedSet<TSet>
+  : T
+
 export type UnwrapRefSimple<T> = T extends
   | Function
-  | CollectionTypes
   | BaseTypes
   | Ref
   | RefUnwrapBailTypes[keyof RefUnwrapBailTypes]
   ? T
+  : T extends CollectionTypes
+  ? UnwrapRefCollectionTypes<T>
   : T extends Array<any>
   ? { [K in keyof T]: UnwrapRefSimple<T[K]> }
   : T extends object & { [ShallowReactiveMarker]?: never }
diff --git a/packages/runtime-core/src/index.ts b/packages/runtime-core/src/index.ts
index 0d5e320e83f..0c4d6335088 100644
--- a/packages/runtime-core/src/index.ts
+++ b/packages/runtime-core/src/index.ts
@@ -153,6 +153,9 @@ export {
   ToRefs,
   UnwrapRef,
   ShallowRef,
+  UnwrappedMap,
+  UnwrappedWeakMap,
+  UnwrappedSet,
   ShallowUnwrapRef,
   CustomRefFactory,
   ReactiveFlags,
diff --git a/test-dts/ref.test-d.ts b/test-dts/ref.test-d.ts
index f91d6415f70..7713b6f1470 100644
--- a/test-dts/ref.test-d.ts
+++ b/test-dts/ref.test-d.ts
@@ -70,6 +70,101 @@ function plainType(arg: number | Ref<number>) {
   // should still unwrap in objects nested in arrays
   const arr2 = ref([{ a: ref(1) }]).value
   expectType<number>(arr2[0].a)
+
+  // unwrapped collections should all extend the native ones
+  // and be assignable if no unwrapping has happened
+  {
+    const someSet: Set<Ref<number>> = new Set()
+    const someWeakSet: WeakSet<Ref<number>> = new WeakSet()
+    const someMap: Map<string, Ref<number>> = new Map()
+    const someWeakMap: WeakMap<object, Ref<number>> = new WeakMap()
+    const setRef = ref(someSet)
+    const weakSetRef = ref(someWeakSet)
+    const mapRef = ref(someMap)
+    const weakMapRef = ref(someWeakMap)
+
+    expectType<Set<Ref<number>>>(setRef.value)
+    expectType<WeakSet<Ref<number>>>(weakSetRef.value)
+    expectType<Map<string, Ref<number>>>(mapRef.value)
+    expectType<WeakMap<object, Ref<number>>>(weakMapRef.value)
+  }
+
+  // should still unwrap in objects nested in collections
+  {
+    let someSet: Set<{ a: Ref<number> }> = new Set()
+    const setRef = ref(someSet)
+
+    const obj = {
+      a: ref(1)
+    }
+
+    // do not require unwrapped object
+    // on the `add`, `has`, and `delete` method
+    setRef.value.add(obj)
+    setRef.value.has(obj)
+    setRef.value.delete(obj)
+
+    // make sure to unwrap the object when reading
+    for (const { a } of setRef.value.values()) {
+      expectType<number>(a)
+    }
+
+    // @ts-expect-error should not be assignable back
+    // because `a` was unwrapped
+    someSet = setRef.value
+  }
+
+  {
+    let someWeakSet: WeakSet<{ a: Ref<number> }> = new WeakSet()
+    const weakSetRef = ref(someWeakSet)
+
+    const obj = {
+      a: ref(1)
+    }
+
+    weakSetRef.value.add(obj)
+    weakSetRef.value.has(obj)
+    weakSetRef.value.delete(obj)
+
+    // should be assignable back because WeakSets are not unwrapped
+    // (their values can never be read and only be used for lookup)
+    someWeakSet = weakSetRef.value
+  }
+
+  {
+    let someMap: Map<string, { a: Ref<number> }> = new Map()
+    const mapRef = ref(someMap)
+
+    const obj = {
+      a: ref(1)
+    }
+
+    // do not require unwrapped object on the `set` method
+    mapRef.value.set('key', obj)
+
+    // make sure to unwrap the object when reading
+    expectType<{ a: number } | undefined>(mapRef.value.get('key'))
+
+    // @ts-expect-error
+    someMap = mapRef.value
+  }
+
+  {
+    let someWeakMap: WeakMap<{}, { a: Ref<number> }> = new Map()
+    const weakMapRef = ref(someWeakMap)
+
+    const key = {}
+    const obj = {
+      a: ref(1)
+    }
+
+    weakMapRef.value.set(key, obj)
+
+    expectType<{ a: number } | undefined>(weakMapRef.value.get(key))
+
+    // @ts-expect-error
+    someWeakMap = weakMapRef.value
+  }
 }
 
 plainType(1)
diff --git a/test-dts/tsconfig.build.json b/test-dts/tsconfig.build.json
index ea058b1d826..87628b962c3 100644
--- a/test-dts/tsconfig.build.json
+++ b/test-dts/tsconfig.build.json
@@ -5,6 +5,7 @@
     "module": "esnext",
     "strict": true,
     "moduleResolution": "node",
+    "downlevelIteration": true,
     "lib": ["esnext", "dom"]
   },
   "include": ["./*"]
diff --git a/test-dts/tsconfig.json b/test-dts/tsconfig.json
index 433de219a8e..0d15c9fbeb2 100644
--- a/test-dts/tsconfig.json
+++ b/test-dts/tsconfig.json
@@ -2,7 +2,8 @@
   "extends": "../tsconfig.json",
   "compilerOptions": {
     "noEmit": true,
-    "declaration": true
+    "declaration": true,
+    "downlevelIteration": true
   },
   "exclude": ["../packages/*/__tests__", "../packages/template-explorer"]
 }