fix(types): reactive collection type is not as expected (#5954)

packages/compiler-ssr/src/transforms/ssrTransformComponent.ts

@@ -334,20 +334,19 @@ function subTransform(
   // traverse
   traverseNode(childRoot, childContext)
   // merge helpers/components/directives into parent context
-  ;(['helpers', 'components', 'directives'] as const).forEach(key => {
-    childContext[key].forEach((value: any, helperKey: any) => {
-      if (key === 'helpers') {
-        const parentCount = parentContext.helpers.get(helperKey)
-        if (parentCount === undefined) {
-          parentContext.helpers.set(helperKey, value)
-        } else {
-          parentContext.helpers.set(helperKey, value + parentCount)
-        }
-      } else {
-        ;(parentContext[key] as any).add(value)
-      }
+  ;(['components', 'directives'] as const).forEach(key => {
+    childContext[key].forEach((value: any) => {
+      ;(parentContext[key] as any).add(value)
     })
   })
+  childContext['helpers'].forEach((value: any, helperKey: any) => {
+    const parentCount = parentContext.helpers.get(helperKey)
+    if (parentCount === undefined) {
+      parentContext.helpers.set(helperKey, value)
+    } else {
+      parentContext.helpers.set(helperKey, value + parentCount)
+    }
+  })
   // imports/hoists are not merged because:
   // - imports are only used for asset urls and should be consistent between
   //   node/client branches

packages/reactivity/__tests__/collections/shallowReadonly.spec.ts

@@ -2,7 +2,7 @@ import { isReactive, isReadonly, shallowReadonly } from '../../src'
 
 describe('reactivity/collections', () => {
   describe('shallowReadonly/Map', () => {
-    ;[Map, WeakMap].forEach(Collection => {
+    ;[Map, WeakMap].forEach((Collection: any) => {
       test('should make the map/weak-map readonly', () => {
         const key = {}
         const val = { foo: 1 }
@@ -81,7 +81,7 @@ describe('reactivity/collections', () => {
 
   describe('shallowReadonly/Set', () => {
     test('should make the set/weak-set readonly', () => {
-      ;[Set, WeakSet].forEach(Collection => {
+      ;[Set, WeakSet].forEach((Collection: any) => {
         const obj = { foo: 1 }
         const original = new Collection([obj])
         const sroSet = shallowReadonly(original)

packages/reactivity/__tests__/shallowReadonly.spec.ts

@@ -38,7 +38,7 @@ describe('reactivity/shallowReadonly', () => {
   })
 
   describe('collection/Map', () => {
-    ;[Map, WeakMap].forEach(Collection => {
+    ;[Map, WeakMap].forEach((Collection: any) => {
       test('should make the map/weak-map readonly', () => {
         const key = {}
         const val = { foo: 1 }
@@ -117,7 +117,7 @@ describe('reactivity/shallowReadonly', () => {
 
   describe('collection/Set', () => {
     test('should make the set/weak-set readonly', () => {
-      ;[Set, WeakSet].forEach(Collection => {
+      ;[Set, WeakSet].forEach((Collection: any) => {
         const obj = { foo: 1 }
         const original = new Collection([obj])
         const sroSet = shallowReadonly(original)

packages/reactivity/index.d.ts

@@ -0,0 +1,20 @@
+import { CollectionUnwrapRefs } from './dist/reactivity'
+
+export * from './dist/reactivity'
+
+declare global {
+  interface Map<K, V> {
+    get<T extends ThisType<Map<K, V>>>(this: T, key: K): CollectionUnwrapRefs<T, V> | undefined
+    forEach<T extends ThisType<Map<K, V>>>(this: T, callbackfn: (value: CollectionUnwrapRefs<T, V>, key: CollectionUnwrapRefs<T, K>, map: T) => void, thisArg?: any): void;
+  }
+  interface WeakMap<K, V> {
+    get<T extends ThisType<WeakMap<K, V>>>(this: T, key: K): CollectionUnwrapRefs<T, V> | undefined
+    forEach<T extends ThisType<WeakMap<K, V>>>(this: T, callbackfn: (value: CollectionUnwrapRefs<T, V>, key: CollectionUnwrapRefs<T, K>, map: T) => void, thisArg?: any): void;
+  }
+  interface Set<T> {
+    forEach<TT extends ThisType<Set<T>>>(this: TT, callbackfn: (value1: CollectionUnwrapRefs<TT, T>, value2: CollectionUnwrapRefs<TT, T>, set: TT) => void, thisArg?: any): void;
+  }
+  interface WeakSet<T> {
+    forEach<TT extends ThisType<WeakSet<T>>>(this: TT, callbackfn: (value1: CollectionUnwrapRefs<TT, T>, value2: CollectionUnwrapRefs<TT, T>, set: TT) => void, thisArg?: any): void;
+  }
+}

packages/reactivity/package.json

@@ -4,11 +4,12 @@
   "description": "@vue/reactivity",
   "main": "index.js",
   "module": "dist/reactivity.esm-bundler.js",
-  "types": "dist/reactivity.d.ts",
+  "types": "index.d.ts",
   "unpkg": "dist/reactivity.global.js",
   "jsdelivr": "dist/reactivity.global.js",
   "files": [
     "index.js",
+    "index.d.ts",
     "dist"
   ],
   "sideEffects": false,

packages/reactivity/src/collectionHandlers.ts

@@ -1,8 +1,41 @@
-import { toRaw, ReactiveFlags, toReactive, toReadonly } from './reactive'
+import { 
+  toRaw,
+  ReactiveFlags,
+  toReactive,
+  toReadonly,
+  ReactiveObject,
+  IsCollectionReactive,
+  IsCollectionReadonly,
+  ReadonlyObject
+} from './reactive'
 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 CollectionUnwrapRefs<T extends object, V> = 
+  IsCollectionReactive<T> extends true
+  ? ReactiveObject<V>
+  : IsCollectionReadonly<T> extends true
+  ? ReadonlyObject<T>
+  : V
+
+declare global {
+  interface Map<K, V> {
+    get<T extends ThisType<Map<K, V>>>(this: T, key: K): CollectionUnwrapRefs<T, V> | undefined
+    forEach<T extends ThisType<Map<K, V>>>(this: T, callbackfn: (value: CollectionUnwrapRefs<T, V>, key: CollectionUnwrapRefs<T, K>, map: T) => void, thisArg?: any): void;
+  }
+  interface WeakMap<K, V> {
+    get<T extends ThisType<WeakMap<K, V>>>(this: T, key: K): CollectionUnwrapRefs<T, V> | undefined
+    forEach<T extends ThisType<WeakMap<K, V>>>(this: T, callbackfn: (value: CollectionUnwrapRefs<T, V>, key: CollectionUnwrapRefs<T, K>, map: T) => void, thisArg?: any): void;
+  }
+  interface Set<T> {
+    forEach<TT extends ThisType<Set<T>>>(this: TT, callbackfn: (value1: CollectionUnwrapRefs<TT, T>, value2: CollectionUnwrapRefs<TT, T>, set: TT) => void, thisArg?: any): void;
+  }
+  interface WeakSet<T> {
+    forEach<TT extends ThisType<WeakSet<T>>>(this: TT, callbackfn: (value1: CollectionUnwrapRefs<TT, T>, value2: CollectionUnwrapRefs<TT, T>, set: TT) => void, thisArg?: any): void;
+  }
+}
+
 export type CollectionTypes = IterableCollections | WeakCollections
 
 type IterableCollections = Map<any, any> | Set<any>

packages/reactivity/src/index.ts

@@ -66,3 +66,4 @@ export {
   onScopeDispose
 } from './effectScope'
 export { TrackOpTypes, TriggerOpTypes } from './operations'
+export { CollectionUnwrapRefs } from './collectionHandlers'

packages/reactivity/src/reactive.ts

@@ -6,13 +6,27 @@ import {
   shallowReadonlyHandlers
 } from './baseHandlers'
 import {
+  CollectionTypes,
   mutableCollectionHandlers,
   readonlyCollectionHandlers,
   shallowCollectionHandlers,
   shallowReadonlyCollectionHandlers
 } from './collectionHandlers'
 import type { UnwrapRefSimple, Ref, RawSymbol } from './ref'
 
+export declare const CollectionReactiveMarker: unique symbol
+export declare const CollectionReadonlyMarker: unique symbol
+export declare const ShallowReactiveMarker: unique symbol
+
+export type IsCollectionReadonly<T extends object> = T extends {
+  [CollectionReadonlyMarker]?: true
+} ? true : false
+
+export type IsCollectionReactive<T extends object> = T extends {
+  [CollectionReactiveMarker]?: true
+} ? true : false
+
+
 export const enum ReactiveFlags {
   SKIP = '__v_skip',
   IS_REACTIVE = '__v_isReactive',
@@ -64,6 +78,11 @@ function getTargetType(value: Target) {
 // only unwrap nested ref
 export type UnwrapNestedRefs<T> = T extends Ref ? T : UnwrapRefSimple<T>
 
+export type ReactiveObject<T> = 
+  T extends CollectionTypes 
+    ? UnwrapNestedRefs<T> & { [CollectionReactiveMarker]?: true }
+    : UnwrapNestedRefs<T>
+
 /**
  * Creates a reactive copy of the original object.
  *
@@ -86,7 +105,7 @@ export type UnwrapNestedRefs<T> = T extends Ref ? T : UnwrapRefSimple<T>
  * count.value // -> 1
  * ```
  */
-export function reactive<T extends object>(target: T): UnwrapNestedRefs<T>
+export function reactive<T extends object>(target: T): ReactiveObject<T>
 export function reactive(target: object) {
   // if trying to observe a readonly proxy, return the readonly version.
   if (isReadonly(target)) {
@@ -101,8 +120,6 @@ export function reactive(target: object) {
   )
 }
 
-export declare const ShallowReactiveMarker: unique symbol
-
 export type ShallowReactive<T> = T & { [ShallowReactiveMarker]?: true }
 
 /**
@@ -146,13 +163,17 @@ export type DeepReadonly<T> = T extends Builtin
   ? { readonly [K in keyof T]: DeepReadonly<T[K]> }
   : Readonly<T>
 
+export type ReadonlyObject<T> = 
+  T extends CollectionTypes 
+    ? DeepReadonly<UnwrapNestedRefs<T>> & { [CollectionReadonlyMarker]?: true }
+    : DeepReadonly<UnwrapNestedRefs<T>>
 /**
  * Creates a readonly copy of the original object. Note the returned copy is not
  * made reactive, but `readonly` can be called on an already reactive object.
  */
 export function readonly<T extends object>(
   target: T
-): DeepReadonly<UnwrapNestedRefs<T>> {
+): ReadonlyObject<T> {
   return createReactiveObject(
     target,
     true,

test-dts/reactivity.test-d.ts

@@ -60,3 +60,162 @@ describe('shallowReadonly ref unwrap', () => {
   expectType<Ref>(r.count.n)
   r.count.n.value = 123
 })
+
+describe('collection-type', () => {
+  it('primitive type as value', () => {
+    const m = reactive(new Map<string, number>())
+    const s = reactive(new Set<string>())
+    m.set('a', 1)
+    s.add('b')
+
+    expectType<Map<string, number>>(m)
+    expectType<Set<string>>(s)
+    expectType<number>(m.get('a')!)
+    m.forEach((v1, v2, m) => {
+      expectType<number>(v1)
+      expectType<string>(v2)
+      expectType<Map<string, number>>(m)
+    })
+    s.forEach((v1, v2, s) => {
+      expectType<string>(v1)
+      expectType<string>(v2)
+      expectType<Set<string>>(s)
+    })
+  })
+
+  it('composite types as values', () => {
+    const m = reactive(new Map<string, {a: number}>())
+    const s = reactive(new Set<{a: number}>())
+    m.set('a', {a: 1})
+    s.add({a: 1})
+
+    expectType<{a: number}>(m.get('a')!)
+    m.forEach((v1, v2, m) => {
+      expectType<{a: number}>(v1)
+      expectType<string>(v2)
+    })
+    s.forEach((v1, v2, s) => {
+      expectType<{a: number}>(v1)
+      expectType<{a: number}>(v2)
+    })
+  })
+
+  it('composite type as key', () => {
+    const m = reactive(new Map<{a: number}, string>())
+    m.set({a: 1}, 'a')
+    m.forEach((v1, v2, m) => {
+      expectType<string>(v1)
+      expectType<{a: number}>(v2)
+    })
+  })
+
+  it('ref as value', () => {
+    const m = reactive(new Map<string, Ref<number>>())
+    const s = reactive(new Set<Ref<string>>())
+    m.set('a', ref(1))
+    s.add(ref('b'))
+
+    expectType<Map<string, Ref<number>>>(m)
+    expectType<Set<Ref<string>>>(s)
+    expectType<Ref<number>>(m.get('a')!)
+
+    m.forEach((v1, v2, m) => {
+      expectType<Ref<number>>(v1)
+      expectType<string>(v2)
+      expectType<Map<string, Ref<number>>>(m)
+    })
+
+    s.forEach((v1, v2, s) => {
+      expectType<Ref<string>>(v1)
+      expectType<Ref<string>>(v2)
+      expectType<Set<Ref<string>>>(s)
+    })
+  })
+
+  it('when value is an object and ref is used as an object property', () => {
+    const m = reactive(new Map<string, { foo: Ref<number> }>())
+    const s = reactive(new Set<{ foo:Ref<number> }>())
+    m.set('a', {
+      foo: ref(1)
+    })
+    s.add({
+      foo: ref(1)
+    })
+
+
+    expectType<Map<string, { foo: Ref<number> }>>(m)
+    expectType<Set<{ foo:Ref<number> }>>(s)
+    expectType<{foo: number}>(m.get('a')!)
+
+    m.forEach((v1, v2, m) => {
+      expectType<{ foo: number }>(v1)
+      expectType<string>(v2)
+      expectType<Map<string, { foo: Ref<number> }>>(m)
+    })
+
+    s.forEach((v1, v2, s) => {
+      expectType<{foo: number}>(v1)
+      expectType<{foo: number}>(v2)
+      expectType<Set<{foo: Ref<number>}>>(s)
+    })
+  })
+  it('When the key is an object and the responsive data is an object property', () => {
+    const m = reactive(new Map<{ foo: Ref<number> }, string>())
+    const obj = {
+      foo: ref(1)
+    }
+    m.set(obj, 'a')
+
+    expectType<Map<{ foo:Ref<number> }, string>>(m)
+    m.forEach((v1, v2, m) => {
+      expectType<string>(v1)
+      expectType<{foo: number}>(v2)
+      expectType<Map<{ foo: Ref<number> }, string>>(m)
+    })
+  })
+  it('normal use under non-reactive', () => {
+    const m = new Map<string, number>()
+    const s = new Set<string>()
+    const mr = new Map<{a: Ref<number>}, {b: Ref<number>}>()
+    const sr = new Set<{a: Ref<number>}>()
+    const objA = {
+      a: ref(1)
+    }
+    const objB = {
+      b: ref(1)
+    }
+
+    m.set('a', 1)
+    s.add('b')
+    mr.set(objA, objB)
+    sr.add(objA)
+
+    expectType<Map<string, number>>(m)
+    expectType<Set<string>>(s)
+    expectType<Map<{a: Ref<number>}, {b: Ref<number>}>>(mr)
+    expectType<Set<{a: Ref<number>}>>(sr)
+    expectType<number>(m.get('a')!)
+    expectType<{b: Ref<number>}>(mr.get(objA)!)
+
+    m.forEach((v1, v2, m) => {
+      expectType<number>(v1)
+      expectType<string>(v2)
+      expectType<Map<string, number>>(m)
+    })
+    s.forEach((v1, v2, s) => {
+      expectType<string>(v1)
+      expectType<string>(v2)
+      expectType<Set<string>>(s)
+    })
+    mr.forEach((v1, v2, m) => {
+      expectType<{b: Ref<number>}>(v1)
+      expectType<{a: Ref<number>}>(v2)
+      expectType<Map<{a: Ref<number>}, {b: Ref<number>}>>(m)
+    })
+    sr.forEach((v1, v2, s) => {
+      expectType<{a: Ref<number>}>(v1)
+      expectType<{a: Ref<number>}>(v2)
+      expectType<Set<{a: Ref<number>}>>(s)
+    })
+  })
+})