From ad274737015c36906d76f3189203093fa3a2e4e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=89=E5=92=B2=E6=99=BA=E5=AD=90=20Kevin=20Deng?=
 <sxzz@sxzz.moe>
Date: Fri, 1 Dec 2023 02:33:09 +0800
Subject: [PATCH] fix(types): fix reactive collection types (#8960)

close #8904
---
 packages/dts-test/reactivity.test-d.ts        | 28 +++++++++++++++++++
 packages/reactivity/src/collectionHandlers.ts |  2 +-
 packages/reactivity/src/ref.ts                | 24 ++++++++++------
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/packages/dts-test/reactivity.test-d.ts b/packages/dts-test/reactivity.test-d.ts
index 4504f9e352d..431bbafbdf4 100644
--- a/packages/dts-test/reactivity.test-d.ts
+++ b/packages/dts-test/reactivity.test-d.ts
@@ -62,3 +62,31 @@ describe('should unwrap tuple correctly', () => {
   const reactiveTuple = reactive(tuple)
   expectType<Ref<number>>(reactiveTuple[0])
 })
+
+describe('should unwrap Map correctly', () => {
+  const map = reactive(new Map<string, Ref<number>>())
+  expectType<Ref<number>>(map.get('a')!)
+
+  const map2 = reactive(new Map<string, { wrap: Ref<number> }>())
+  expectType<number>(map2.get('a')!.wrap)
+
+  const wm = reactive(new WeakMap<object, Ref<number>>())
+  expectType<Ref<number>>(wm.get({})!)
+
+  const wm2 = reactive(new WeakMap<object, { wrap: Ref<number> }>())
+  expectType<number>(wm2.get({})!.wrap)
+})
+
+describe('should unwrap Set correctly', () => {
+  const set = reactive(new Set<Ref<number>>())
+  expectType<Set<Ref<number>>>(set)
+
+  const set2 = reactive(new Set<{ wrap: Ref<number> }>())
+  expectType<Set<{ wrap: number }>>(set2)
+
+  const ws = reactive(new WeakSet<Ref<number>>())
+  expectType<WeakSet<Ref<number>>>(ws)
+
+  const ws2 = reactive(new WeakSet<{ wrap: Ref<number> }>())
+  expectType<WeakSet<{ wrap: number }>>(ws2)
+})
diff --git a/packages/reactivity/src/collectionHandlers.ts b/packages/reactivity/src/collectionHandlers.ts
index ecdcae00b7c..abf623f551c 100644
--- a/packages/reactivity/src/collectionHandlers.ts
+++ b/packages/reactivity/src/collectionHandlers.ts
@@ -3,7 +3,7 @@ import { track, trigger, ITERATE_KEY, MAP_KEY_ITERATE_KEY } from './effect'
 import { TrackOpTypes, TriggerOpTypes } from './operations'
 import { capitalize, hasOwn, hasChanged, toRawType, isMap } from '@vue/shared'
 
-export type CollectionTypes = IterableCollections | WeakCollections
+type CollectionTypes = IterableCollections | WeakCollections
 
 type IterableCollections = Map<any, any> | Set<any>
 type WeakCollections = WeakMap<any, any> | WeakSet<any>
diff --git a/packages/reactivity/src/ref.ts b/packages/reactivity/src/ref.ts
index 201766158ed..9f7ebdad822 100644
--- a/packages/reactivity/src/ref.ts
+++ b/packages/reactivity/src/ref.ts
@@ -17,7 +17,6 @@ import {
   isShallow
 } from './reactive'
 import type { ShallowReactiveMarker } from './reactive'
-import { CollectionTypes } from './collectionHandlers'
 import { createDep, Dep } from './dep'
 
 declare const RefSymbol: unique symbol
@@ -492,16 +491,23 @@ export type UnwrapRef<T> = T extends ShallowRef<infer V>
 
 export type UnwrapRefSimple<T> = T extends
   | Function
-  | CollectionTypes
   | BaseTypes
   | Ref
   | RefUnwrapBailTypes[keyof RefUnwrapBailTypes]
   | { [RawSymbol]?: true }
   ? T
-  : T extends ReadonlyArray<any>
-    ? { [K in keyof T]: UnwrapRefSimple<T[K]> }
-    : T extends object & { [ShallowReactiveMarker]?: never }
-      ? {
-          [P in keyof T]: P extends symbol ? T[P] : UnwrapRef<T[P]>
-        }
-      : T
+  : T extends Map<infer K, infer V>
+    ? Map<K, UnwrapRefSimple<V>>
+    : T extends WeakMap<infer K, infer V>
+      ? WeakMap<K, UnwrapRefSimple<V>>
+      : T extends Set<infer V>
+        ? Set<UnwrapRefSimple<V>>
+        : T extends WeakSet<infer V>
+          ? WeakSet<UnwrapRefSimple<V>>
+          : T extends ReadonlyArray<any>
+            ? { [K in keyof T]: UnwrapRefSimple<T[K]> }
+            : T extends object & { [ShallowReactiveMarker]?: never }
+              ? {
+                  [P in keyof T]: P extends symbol ? T[P] : UnwrapRef<T[P]>
+                }
+              : T