fix(set-map): support custom Set and Map (#71)

src/current.ts

@@ -4,6 +4,8 @@ import {
   get,
   getProxyDraft,
   getType,
+  isBaseMapInstance,
+  isBaseSetInstance,
   isDraft,
   isDraftable,
   isEqual,
@@ -68,7 +70,9 @@ function getCurrent(target: any) {
   function ensureShallowCopy() {
     currentValue =
       type === DraftType.Map
-        ? new Map(target)
+        ? !isBaseMapInstance(target)
+          ? new (Object.getPrototypeOf(target).constructor)(target)
+          : new Map(target)
         : type === DraftType.Set
           ? Array.from(proxyDraft!.setMap!.values()!)
           : shallowCopy(target, proxyDraft?.options);
@@ -96,7 +100,13 @@ function getCurrent(target: any) {
       set(currentValue, key, newValue);
     }
   });
-  return type === DraftType.Set ? new Set(currentValue) : currentValue;
+  if (type === DraftType.Set) {
+    const value = proxyDraft?.original ?? currentValue;
+    return !isBaseSetInstance(value)
+      ? new (Object.getPrototypeOf(value).constructor)(currentValue)
+      : new Set(currentValue);
+  }
+  return currentValue;
 }
 
 /**

src/set.ts

@@ -154,6 +154,8 @@ export const setHandler = {
 if (Set.prototype.difference) {
   // for compatibility with new Set methods
   // https://github.com/tc39/proposal-set-methods
+  // And `https://github.com/tc39/proposal-set-methods/blob/main/details.md#symbolspecies` has some details about the `@@species` symbol.
+  // So we can't use SubSet instance constructor to get the constructor of the SubSet instance.
   Object.assign(setHandler, {
     intersection(this: Set<any>, other: ReadonlySetLike<any>): Set<any> {
       return Set.prototype.intersection.call(new Set(this.values()), other);

src/utils/copy.ts

@@ -1,6 +1,7 @@
 import type { Options, ProxyDraft } from '../interface';
 import { dataTypes } from '../constant';
 import { getValue, isDraft, isDraftable } from './draft';
+import { isBaseMapInstance, isBaseSetInstance } from './proto';
 
 function strictCopy(target: any) {
   const copy = Object.create(Object.getPrototypeOf(target));
@@ -34,10 +35,18 @@ export function shallowCopy(original: any, options?: Options<any, any>) {
   if (Array.isArray(original)) {
     return Array.prototype.concat.call(original);
   } else if (original instanceof Set) {
+    if (!isBaseSetInstance(original)) {
+      const SubClass = Object.getPrototypeOf(original).constructor;
+      return new SubClass(original.values());
+    }
     return Set.prototype.difference
       ? Set.prototype.difference.call(original, new Set())
       : new Set(original.values());
   } else if (original instanceof Map) {
+    if (!isBaseMapInstance(original)) {
+      const SubClass = Object.getPrototypeOf(original).constructor;
+      return new SubClass(original);
+    }
     return new Map(original);
   } else if (
     options?.mark &&
@@ -88,11 +97,25 @@ function deepClone<T>(target: T): T;
 function deepClone(target: any) {
   if (!isDraftable(target)) return getValue(target);
   if (Array.isArray(target)) return target.map(deepClone);
-  if (target instanceof Map)
-    return new Map(
-      Array.from(target.entries()).map(([k, v]) => [k, deepClone(v)])
-    );
-  if (target instanceof Set) return new Set(Array.from(target).map(deepClone));
+  if (target instanceof Map) {
+    const iterable = Array.from(target.entries()).map(([k, v]) => [
+      k,
+      deepClone(v),
+    ]) as Iterable<readonly [any, any]>;
+    if (!isBaseMapInstance(target)) {
+      const SubClass = Object.getPrototypeOf(target).constructor;
+      return new SubClass(iterable);
+    }
+    return new Map(iterable);
+  }
+  if (target instanceof Set) {
+    const iterable = Array.from(target).map(deepClone);
+    if (!isBaseSetInstance(target)) {
+      const SubClass = Object.getPrototypeOf(target).constructor;
+      return new SubClass(iterable);
+    }
+    return new Set(iterable);
+  }
   const copy = Object.create(Object.getPrototypeOf(target));
   for (const key in target) copy[key] = deepClone(target[key]);
   return copy;

src/utils/proto.ts

@@ -15,3 +15,11 @@ export function getDescriptor(target: object, key: PropertyKey) {
   }
   return;
 }
+
+export function isBaseSetInstance(obj: any) {
+  return Object.getPrototypeOf(obj) === Set.prototype;
+}
+
+export function isBaseMapInstance(obj: any) {
+  return Object.getPrototypeOf(obj) === Map.prototype;
+}

test/__snapshots__/current.test.ts.snap

@@ -0,0 +1,60 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`current() for Custom Set/Map draft 1`] = `
+{
+  "a": Set {
+    {
+      "id": 42,
+    },
+    {
+      "id": 43,
+    },
+  },
+  "b": Map {
+    1 => {
+      "id": 42,
+    },
+    2 => {
+      "id": 43,
+    },
+  },
+  "x": {
+    "y": {
+      "z": {
+        "k": 43,
+      },
+    },
+  },
+}
+`;
+
+exports[`current() for Custom Set/Map draft 2`] = `
+{
+  "a": Set {
+    {
+      "id": 42,
+    },
+    {
+      "id": 43,
+    },
+    Set {
+      {},
+    },
+  },
+  "b": Map {
+    1 => {
+      "id": 42,
+    },
+    2 => {
+      "id": 43,
+    },
+  },
+  "x": {
+    "y": {
+      "z": {
+        "k": 43,
+      },
+    },
+  },
+}
+`;

test/apply.test.ts

@@ -9,7 +9,7 @@
 /* eslint-disable no-param-reassign */
 /* eslint-disable @typescript-eslint/ban-ts-comment */
 import { create, apply, Patches, original } from '../src';
-import { deepClone } from '../src/utils';
+import { deepClone, set } from '../src/utils';
 
 test('classic case', () => {
   const data = {
@@ -1411,3 +1411,30 @@ test('modify deep object', () => {
   expect(base.map.get('set2')).toBe(set2);
   expect(first(state.map.get('set1'))).toEqual({ a: 2 });
 });
+
+test('#70 - deep copy patches with Custom Set/Map', () => {
+  class CustomSet<T> extends Set<T> {}
+  class CustomMap<K, V> extends Map<K, V> {}
+  const baseState = {
+    map: new CustomMap<any, any>(),
+    set: new CustomSet<any>(),
+  };
+  const [state, patches, inversePatches] = create(
+    baseState,
+    (draft) => {
+      draft.map = new CustomMap<any, any>([[1, 1]]);
+      draft.set = new CustomSet<any>([1]);
+    },
+    {
+      enablePatches: true,
+    }
+  );
+  const nextState = apply(baseState, patches);
+  expect(patches[0].value).toBeInstanceOf(CustomMap);
+  expect(patches[1].value).toBeInstanceOf(CustomSet);
+  expect(nextState).toEqual(state);
+  const prevState = apply(state, inversePatches);
+  expect(inversePatches[0].value).toBeInstanceOf(CustomMap);
+  expect(inversePatches[1].value).toBeInstanceOf(CustomSet);
+  expect(prevState).toEqual(baseState);
+});

test/current.test.ts

@@ -307,7 +307,6 @@ test('nested create() - Avoid deep copies', () => {
   });
 });
 
-
 test('#61 - type issue: current of Draft<T> type should return T type', () => {
   function test<T extends { x: { y: ReadonlySet<string> } }>(base: T): T {
     const [draft] = create(base);
@@ -316,3 +315,31 @@ test('#61 - type issue: current of Draft<T> type should return T type', () => {
     return currentValue0;
   }
 });
+
+test('current() for Custom Set/Map draft', () => {
+  class CustomSet<T> extends Set<T> {}
+  class CustomMap<T, P> extends Map<T, P> {}
+  const obj = { k: 42 };
+  const base = {
+    x: { y: { z: obj } },
+    a: new CustomSet([{ id: 42 }]),
+    b: new CustomMap([[1, { id: 42 }]]),
+  };
+  create(base, (draft) => {
+    const obj1 = draft.x.y.z;
+    const d = { id: 43 };
+    draft.a.add(d);
+    draft.b.set(2, { id: 43 });
+    draft.x.y.z = { k: 43 };
+    const c = current(draft);
+    expect(c.a.has(d)).toBeTruthy();
+    expect(c.b.get(2)).toEqual({ id: 43 });
+    expect(c).toMatchSnapshot();
+    // @ts-ignore
+    draft.a.add(new CustomSet([{}]));
+    // @ts-ignore
+    Array.from(draft.a)[2].value = obj1;
+    const f = current(draft);
+    expect(f).toMatchSnapshot();
+  });
+});

test/immer-non-support.test.ts

@@ -18,6 +18,7 @@ import {
   current as immerCurrent,
   createDraft,
   finishDraft,
+  immerable,
 } from 'immer';
 import { create, apply, current } from '../src';
 
@@ -642,3 +643,93 @@ test('set - new Set API', () => {
     });
   }
 });
+
+test('CustomSet', () => {
+  {
+    enableMapSet();
+    class CustomSet extends Set {
+      [immerable] = true;
+
+      getIdentity() {
+        return 'CustomSet';
+      }
+    }
+
+    const s = new CustomSet();
+    const newS = produce(s, (draft) => {
+      draft.add(1);
+      // @ts-ignore
+      expect(typeof draft.getIdentity === 'function').toBeFalsy(); // it should be `true`
+    });
+    // @ts-ignore
+    expect(typeof newS.getIdentity === 'function').toBeFalsy(); // it should be `true`
+  }
+  {
+    class CustomSet extends Set {
+      getIdentity() {
+        return 'CustomSet';
+      }
+    }
+
+    const s = new CustomSet();
+    const newS = create(
+      s,
+      (draft) => {
+        draft.add(1);
+        // @ts-ignore
+        expect(draft.getIdentity()).toBe('CustomSet');
+      },
+      {
+        mark: () => 'immutable',
+      }
+    );
+    expect(newS instanceof CustomSet).toBeTruthy();
+    // @ts-ignore
+    expect(newS.getIdentity()).toBe('CustomSet');
+  }
+});
+
+test('CustomMap', () => {
+  {
+    enableMapSet();
+    class CustomMap extends Map {
+      [immerable] = true;
+
+      getIdentity() {
+        return 'CustomMap';
+      }
+    }
+
+    const state = new CustomMap();
+    const newState = produce(state, (draft) => {
+      draft.set(1, 1);
+      // @ts-ignore
+      expect(typeof draft.getIdentity === 'function').toBeFalsy(); // it should be `true`
+    });
+    // @ts-ignore
+    expect(typeof newState.getIdentity === 'function').toBeFalsy(); // it should be `true`
+  }
+  {
+    class CustomMap extends Map {
+      getIdentity() {
+        return 'CustomMap';
+      }
+    }
+
+    const state = new CustomMap();
+    const newState = create(
+      state,
+      (draft) => {
+        draft.set(1, 1);
+        // @ts-ignore
+        expect(draft.getIdentity()).toBe('CustomMap');
+      },
+      {
+        mark: () => 'immutable',
+      }
+    );
+    expect(newState instanceof CustomMap).toBeTruthy();
+    // @ts-ignore
+    expect(newState.getIdentity()).toBe('CustomMap');
+  }
+});