From f9843eb589ea9752f9021f9ebcfc49f6659350d2 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Fri, 2 Sep 2022 14:40:28 +0200
Subject: [PATCH] feat(ssr): handle Maps and Sets

Close #1608
---
 packages/pinia/__tests__/state.spec.ts | 91 ++++++++++++++++++++++++++
 packages/pinia/src/store.ts            | 12 ++++
 2 files changed, 103 insertions(+)

diff --git a/packages/pinia/__tests__/state.spec.ts b/packages/pinia/__tests__/state.spec.ts
index d10ca572..1c171b73 100644
--- a/packages/pinia/__tests__/state.spec.ts
+++ b/packages/pinia/__tests__/state.spec.ts
@@ -279,6 +279,97 @@ describe('State', () => {
     expect(store.double).toBe(30)
   })
 
+  it('hydrates Set in option stores', async () => {
+    const useStore = defineStore('main', {
+      state: () => ({ set: new Set() }),
+    })
+
+    const pinia = createPinia()
+    pinia.state.value.main = {
+      set: new Set([1, 2]),
+    }
+    setActivePinia(pinia)
+
+    const store = useStore()
+    expect(store.set).toBeInstanceOf(Set)
+    expect([...store.set.values()]).toEqual([1, 2])
+  })
+
+  it('hydrates Set in setup stores', async () => {
+    const useStore = defineStore('main', () => {
+      const setRef = ref(new Set())
+      const setReactive = reactive(new Set())
+      return { setRef, setReactive }
+    })
+
+    const pinia = createPinia()
+    pinia.state.value.main = {
+      setRef: new Set([1, 2]),
+      setReactive: new Set([3, 4]),
+    }
+    setActivePinia(pinia)
+
+    const store = useStore()
+    expect(store.setRef).toBeInstanceOf(Set)
+    expect(store.setReactive).toBeInstanceOf(Set)
+    expect([...store.setRef.values()]).toEqual([1, 2])
+    expect([...store.setReactive.values()]).toEqual([3, 4])
+  })
+
+  it('hydrates Map in option stores', async () => {
+    const useStore = defineStore('main', {
+      state: () => ({ map: new Map() }),
+    })
+
+    const map = new Map()
+
+    map.set('ssr', 'test')
+    map.set('other', 'test2')
+
+    const pinia = createPinia()
+    pinia.state.value.main = { map }
+    setActivePinia(pinia)
+
+    const store = useStore()
+    expect(store.map).toBeInstanceOf(Map)
+    expect([...store.map.entries()]).toEqual([
+      ['ssr', 'test'],
+      ['other', 'test2'],
+    ])
+  })
+
+  it('hydrates Map in setup stores', async () => {
+    const useStore = defineStore('main', () => {
+      const mapRef = ref(new Map())
+      const mapReactive = reactive(new Map())
+      return { mapRef, mapReactive }
+    })
+
+    const mapRef = new Map()
+    const mapReactive = new Map()
+
+    mapRef.set('ref:ssr', 'test')
+    mapRef.set('ref:other', 'test2')
+    mapReactive.set('reactive:ssr', 'test')
+    mapReactive.set('reactive:other', 'test2')
+
+    const pinia = createPinia()
+    pinia.state.value.main = { mapRef, mapReactive }
+    setActivePinia(pinia)
+
+    const store = useStore()
+    expect(store.mapRef).toBeInstanceOf(Map)
+    expect(store.mapReactive).toBeInstanceOf(Map)
+    expect([...store.mapRef.entries()]).toEqual([
+      ['ref:ssr', 'test'],
+      ['ref:other', 'test2'],
+    ])
+    expect([...store.mapReactive.entries()]).toEqual([
+      ['reactive:ssr', 'test'],
+      ['reactive:other', 'test2'],
+    ])
+  })
+
   describe('custom refs', () => {
     let spy!: Mock
     function useCustomRef() {
diff --git a/packages/pinia/src/store.ts b/packages/pinia/src/store.ts
index d7156a0e..59683448 100644
--- a/packages/pinia/src/store.ts
+++ b/packages/pinia/src/store.ts
@@ -57,6 +57,15 @@ function mergeReactiveObjects<T extends StateTree>(
   target: T,
   patchToApply: _DeepPartial<T>
 ): T {
+  // Handle Map instances
+  if (target instanceof Map && patchToApply instanceof Map) {
+    patchToApply.forEach((value, key) => target.set(key, value))
+  }
+  // Handle Set instances
+  if (target instanceof Set && patchToApply instanceof Set) {
+    patchToApply.forEach(target.add, target)
+  }
+
   // no need to go through symbols because they cannot be serialized anyway
   for (const key in patchToApply) {
     if (!patchToApply.hasOwnProperty(key)) continue
@@ -69,6 +78,9 @@ function mergeReactiveObjects<T extends StateTree>(
       !isRef(subPatch) &&
       !isReactive(subPatch)
     ) {
+      // NOTE: here I wanted to warn about inconsistent types but it's not possible because in setup stores one might
+      // start the value of a property as a certain type e.g. a Map, and then for some reason, during SSR, change that
+      // to `undefined`. When trying to hydrate, we want to override the Map with `undefined`.
       target[key] = mergeReactiveObjects(targetValue, subPatch)
     } else {
       // @ts-expect-error: subPatch is a valid value
-- 
2.47.2