From 6e46a574eddb5fa43c8a4ce10c620ecdf1caf3b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thorsten=20L=C3=BCnborg?= <t.luenborg@googlemail.com>
Date: Fri, 27 Nov 2020 20:16:00 +0100
Subject: [PATCH] fix(reactivity): ensure add/set on reactive collections
 return the proxy (#2534)

fix #2530
---
 packages/reactivity/__tests__/collections/Map.spec.ts     | 6 ++++++
 packages/reactivity/__tests__/collections/Set.spec.ts     | 6 ++++++
 packages/reactivity/__tests__/collections/WeakMap.spec.ts | 5 +++++
 packages/reactivity/__tests__/collections/WeakSet.spec.ts | 6 ++++++
 packages/reactivity/src/collectionHandlers.ts             | 8 ++++----
 5 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/packages/reactivity/__tests__/collections/Map.spec.ts b/packages/reactivity/__tests__/collections/Map.spec.ts
index 00fd6539d8f..27ee5b103fb 100644
--- a/packages/reactivity/__tests__/collections/Map.spec.ts
+++ b/packages/reactivity/__tests__/collections/Map.spec.ts
@@ -467,5 +467,11 @@ describe('reactivity/collections', () => {
       proxy.clear()
       expect(spy).toBeCalledTimes(1)
     })
+
+    it('should return proxy from Map.set call', () => {
+      const map = reactive(new Map())
+      const result = map.set('a', 'a')
+      expect(result).toBe(map)
+    })
   })
 })
diff --git a/packages/reactivity/__tests__/collections/Set.spec.ts b/packages/reactivity/__tests__/collections/Set.spec.ts
index 9e76bce8838..b5936c07101 100644
--- a/packages/reactivity/__tests__/collections/Set.spec.ts
+++ b/packages/reactivity/__tests__/collections/Set.spec.ts
@@ -454,5 +454,11 @@ describe('reactivity/collections', () => {
       proxy.clear()
       expect(spy).toBeCalledTimes(1)
     })
+
+    it('should return proxy from Set.add call', () => {
+      const set = reactive(new Set())
+      const result = set.add('a')
+      expect(result).toBe(set)
+    })
   })
 })
diff --git a/packages/reactivity/__tests__/collections/WeakMap.spec.ts b/packages/reactivity/__tests__/collections/WeakMap.spec.ts
index e3dde72e618..31613a06572 100644
--- a/packages/reactivity/__tests__/collections/WeakMap.spec.ts
+++ b/packages/reactivity/__tests__/collections/WeakMap.spec.ts
@@ -133,5 +133,10 @@ describe('reactivity/collections', () => {
       map.set(key, NaN)
       expect(mapSpy).toHaveBeenCalledTimes(1)
     })
+    it('should return proxy from WeakMap.set call', () => {
+      const map = reactive(new WeakMap())
+      const result = map.set({}, 'a')
+      expect(result).toBe(map)
+    })
   })
 })
diff --git a/packages/reactivity/__tests__/collections/WeakSet.spec.ts b/packages/reactivity/__tests__/collections/WeakSet.spec.ts
index 182298cf4b0..01512b0eeef 100644
--- a/packages/reactivity/__tests__/collections/WeakSet.spec.ts
+++ b/packages/reactivity/__tests__/collections/WeakSet.spec.ts
@@ -99,5 +99,11 @@ describe('reactivity/collections', () => {
       expect(observed.has(value)).toBe(true)
       expect(set.has(value)).toBe(false)
     })
+
+    it('should return proxy from WeakSet.add call', () => {
+      const set = reactive(new WeakSet())
+      const result = set.add({})
+      expect(result).toBe(set)
+    })
   })
 })
diff --git a/packages/reactivity/src/collectionHandlers.ts b/packages/reactivity/src/collectionHandlers.ts
index 913cd71f82b..e43d3ac59fd 100644
--- a/packages/reactivity/src/collectionHandlers.ts
+++ b/packages/reactivity/src/collectionHandlers.ts
@@ -76,11 +76,11 @@ function add(this: SetTypes, value: unknown) {
   const target = toRaw(this)
   const proto = getProto(target)
   const hadKey = proto.has.call(target, value)
-  const result = target.add(value)
+  target.add(value)
   if (!hadKey) {
     trigger(target, TriggerOpTypes.ADD, value, value)
   }
-  return result
+  return this
 }
 
 function set(this: MapTypes, key: unknown, value: unknown) {
@@ -97,13 +97,13 @@ function set(this: MapTypes, key: unknown, value: unknown) {
   }
 
   const oldValue = get.call(target, key)
-  const result = target.set(key, value)
+  target.set(key, value)
   if (!hadKey) {
     trigger(target, TriggerOpTypes.ADD, key, value)
   } else if (hasChanged(value, oldValue)) {
     trigger(target, TriggerOpTypes.SET, key, value, oldValue)
   }
-  return result
+  return this
 }
 
 function deleteEntry(this: CollectionTypes, key: unknown) {