Repository: immerjs/immer
Pull request #370 (2019-05-19 20:15:57)

feat: non-primitive keys for Maps, hooks for Maps and Sets


Patch:
diff --git a/readme.md b/readme.md
index ba4fcd6..993c520 100644
--- a/readme.md
+++ b/readme.md
@@ -715,6 +715,10 @@ Make sure you don't return any promises as state, because `produce` will actuall
 
 When using TypeScript, for curried reducers that are typed in the form `produce<Type>((arg) => { })`, rewrite this to `produce((arg: Type) => { })` or `produce((arg: Draft<Type>) => { })` for correct inference.
 
+**Immer 3.\* -> 4.0**
+
+`Maps` and `Sets` are now supported by `immer` and, therefore, immutable. Your current code, where you kept immutability of `Maps` and `Sets` on your own, should keep working as before in simple cases. However, consider covering it with a sufficient number of unit tests to avoid any unexpected outcome.
+
 ## FAQ
 
 _(for those who skimmed the above instead of actually reading)_
diff --git a/src/common.js b/src/common.js
index 7f70d6b..4919b46 100644
--- a/src/common.js
+++ b/src/common.js
@@ -33,42 +33,31 @@ export function original(value) {
 
 // We use Maps as `drafts` for Sets, not Objects
 // See proxy.js
-export function assignSet(target, ...mapOverrides) {
-    mapOverrides.forEach(override => {
-        for (const value of override.values()) {
-            // When we add new drafts we have to remove their originals if present
-            const originalValue = original(value)
-            if (originalValue) {
-                target.delete(originalValue)
-            }
-            target.add(value)
-        }
+export function assignSet(target, override) {
+    override.forEach(value => {
+        // When we add new drafts we have to remove their originals if present
+        const prev = original(value)
+        if (prev) target.delete(prev)
+        target.add(value)
     })
     return target
 }
 
-export function assignMap(target, ...objOverrides) {
-    objOverrides.forEach(override => {
-        for (let key in override) {
-            if (has(override, key)) {
-                target.set(key, override[key])
-            }
-        }
-    })
+// We use Maps as `drafts` for Maps, not Objects
+// See proxy.js
+export function assignMap(target, override) {
+    override.forEach((value, key) => target.set(key, value))
     return target
 }
 
-function assignObjectLegacy(target, ...objOverrides) {
-    objOverrides.forEach(function(override) {
-        for (let key in override) {
-            if (has(override, key)) {
-                target[key] = override[key]
-            }
-        }
+export const assign =
+    Object.assign ||
+    ((target, ...overrides) => {
+        overrides.forEach(override =>
+            Object.keys(override).forEach(key => (target[key] = override[key]))
+        )
+        return target
     })
-    return target
-}
-export const assign = Object.assign || assignObjectLegacy
 
 export const ownKeys =
     typeof Reflect !== "undefined" && Reflect.ownKeys
@@ -110,12 +99,12 @@ export function shallowCopy(base, invokeGetters = false) {
     return clone
 }
 
-export function each(value, cb) {
-    if (Array.isArray(value) || isMap(value) || isSet(value)) {
-        value.forEach((entry, index) => cb(index, entry, value))
-        return
+export function each(obj, iter) {
+    if (Array.isArray(obj) || isMap(obj) || isSet(obj)) {
+        obj.forEach((entry, index) => iter(index, entry, obj))
+    } else {
+        ownKeys(obj).forEach(key => iter(key, obj[key], obj))
     }
-    ownKeys(value).forEach(key => cb(key, value[key], value))
 }
 
 export function isEnumerable(base, prop) {
@@ -124,7 +113,13 @@ export function isEnumerable(base, prop) {
 }
 
 export function has(thing, prop) {
-    return Object.prototype.hasOwnProperty.call(thing, prop)
+    return isMap(thing)
+        ? thing.has(prop)
+        : Object.prototype.hasOwnProperty.call(thing, prop)
+}
+
+export function get(thing, prop) {
+    return isMap(thing) ? thing.get(prop) : thing[prop]
 }
 
 export function is(x, y) {
diff --git a/src/es5.js b/src/es5.js
index c6a6f5d..fab25f2 100644
--- a/src/es5.js
+++ b/src/es5.js
@@ -64,7 +64,7 @@ function revoke() {
     this.revoked = true
 }
 
-function source(state) {
+function latest(state) {
     return state.copy || state.base
 }
 
@@ -82,7 +82,7 @@ function peek(draft, prop) {
 
 function get(state, prop) {
     assertUnrevoked(state)
-    const value = peek(source(state), prop)
+    const value = peek(latest(state), prop)
     if (state.finalizing) return value
     // Create a draft if the value is unmodified.
     if (value === peek(state.base, prop) && isDraftable(value)) {
@@ -96,7 +96,7 @@ function set(state, prop, value) {
     assertUnrevoked(state)
     state.assigned[prop] = true
     if (!state.modified) {
-        if (is(value, peek(source(state), prop))) return
+        if (is(value, peek(latest(state), prop))) return
         markChanged(state)
         prepareCopy(state)
     }
@@ -148,7 +148,7 @@ function assertUnrevoked(state) {
     if (state.revoked === true)
         throw new Error(
             "Cannot use a proxy that has been revoked. Did you pass an object from inside an immer function to an async process? " +
-                JSON.stringify(source(state))
+                JSON.stringify(latest(state))
         )
 }
 
diff --git a/src/immer.js b/src/immer.js
index 7cfece7..c2437f2 100644
--- a/src/immer.js
+++ b/src/immer.js
@@ -4,6 +4,7 @@ import {applyPatches, generatePatches} from "./patches"
 import {
     assign,
     each,
+    get,
     has,
     is,
     isDraft,
@@ -13,7 +14,8 @@ import {
     shallowCopy,
     DRAFT_STATE,
     NOTHING,
-    isSet
+    isSet,
+    original
 } from "./common"
 import {ImmerScope} from "./scope"
 
@@ -188,14 +190,14 @@ export class Immer {
             state.finalized = true
             this.finalizeTree(state.draft, path, scope)
 
-            // TODO: It won't fire for Sets because they don't use `assigned`. Is it an issue?
-            if (this.onDelete) {
+            // We cannot really delete anything inside of a Set. We can only replace the whole Set.
+            if (this.onDelete && !isSet(state.base)) {
                 // The `assigned` object is unreliable with ES5 drafts.
                 if (this.useProxies) {
                     const {assigned} = state
-                    for (const prop in assigned) {
-                        if (!assigned[prop]) this.onDelete(state, prop)
-                    }
+                    each(assigned, (prop, exists) => {
+                        if (!exists) this.onDelete(state, prop)
+                    })
                 } else {
                     // TODO: Figure it out for Maps and Sets if we need to support ES5
                     const {base, copy} = state
@@ -247,32 +249,31 @@ export class Immer {
 
             // In the `finalizeTree` method, only the `root` object may be a draft.
             const isDraftProp = !!state && parent === root
+            const isSetMember = isSet(parent)
 
             if (isDraft(value)) {
                 const path =
                     isDraftProp &&
                     needPatches &&
-                    !isSet(parent) &&
-                    !state.assigned[prop]
+                    !isSetMember && // Set objects are atomic since they have no keys.
+                    !has(state.assigned, prop) // Skip deep patches for assigned keys.
                         ? rootPath.concat(prop)
                         : null
 
                 // Drafts owned by `scope` are finalized here.
                 value = this.finalize(value, path, scope)
+                replace(parent, prop, value)
 
                 // Drafts from another scope must prevent auto-freezing.
                 if (isDraft(value)) {
                     scope.canAutoFreeze = false
                 }
 
-                setProperty(parent, prop, value)
-
                 // Unchanged drafts are never passed to the `onAssign` hook.
-                // TODO: Add tests and support for Maps and Sets
-                if (isDraftProp && value === state.base[prop]) return
+                if (isDraftProp && value === get(state.base, prop)) return
             }
             // Unchanged draft properties are ignored.
-            else if (isDraftProp && is(value, state.base[prop])) {
+            else if (isDraftProp && is(value, get(state.base, prop))) {
                 return
             }
             // Search new objects for unfinalized drafts. Frozen objects should never contain drafts.
@@ -280,7 +281,8 @@ export class Immer {
                 each(value, finalizeProperty)
             }
 
-            if (isDraftProp && this.onAssign) {
+            // We cannot really assign anything inside of a Set. We can only replace the whole Set.
+            if (isDraftProp && this.onAssign && !isSetMember) {
                 this.onAssign(state, prop, value)
             }
         }
@@ -290,21 +292,21 @@ export class Immer {
     }
 }
 
-function setProperty(parent, prop, value) {
-    // Preserve non-enumerable properties.
-    if (Array.isArray(parent) || isEnumerable(parent, prop)) {
-        parent[prop] = value
-        return
-    }
+function replace(parent, prop, value) {
     if (isMap(parent)) {
         parent.set(prop, value)
-        return
-    }
-    if (isSet(parent)) {
-        // Here prop is a proxied value
+    } else if (isSet(parent)) {
+        // In this case, the `prop` is actually a draft.
         parent.delete(prop)
         parent.add(value)
-        return
+    } else if (Array.isArray(parent) || isEnumerable(parent, prop)) {
+        // Preserve non-enumerable properties.
+        parent[prop] = value
+    } else {
+        Object.defineProperty(parent, prop, {
+            value,
+            writable: true,
+            configurable: true
+        })
     }
-    Object.defineProperty(parent, prop, {value})
 }
diff --git a/src/patches.js b/src/patches.js
index 91b336b..dd89767 100644
--- a/src/patches.js
+++ b/src/patches.js
@@ -1,19 +1,11 @@
-import {each, isMap, isSet} from "./common"
+import {get, each, isMap, isSet, has} from "./common"
 
 export function generatePatches(state, basePath, patches, inversePatches) {
     const generatePatchesFn = Array.isArray(state.base)
         ? generateArrayPatches
-        : isMap(state.base)
-        ? generatePatchesFromAssigned(
-              (map, key) => map.get(key),
-              (map, key) => map.has(key)
-          )
         : isSet(state.base)
         ? generateSetPatches
-        : generatePatchesFromAssigned(
-              (obj, key) => obj[key],
-              (obj, key) => key in obj
-          )
+        : generatePatchesFromAssigned
 
     generatePatchesFn(state, basePath, patches, inversePatches)
 }
@@ -87,29 +79,28 @@ function generateArrayPatches(state, basePath, patches, inversePatches) {
     }
 }
 
-function generatePatchesFromAssigned(getValueByKey, hasKey) {
-    return function(state, basePath, patches, inversePatches) {
-        const {base, copy} = state
-        each(state.assigned, (key, assignedValue) => {
-            const origValue = getValueByKey(base, key)
-            const value = getValueByKey(copy, key)
-            const op = !assignedValue
-                ? "remove"
-                : hasKey(base, key)
-                ? "replace"
-                : "add"
-            if (origValue === value && op === "replace") return
-            const path = basePath.concat(key)
-            patches.push(op === "remove" ? {op, path} : {op, path, value})
-            inversePatches.push(
-                op === "add"
-                    ? {op: "remove", path}
-                    : op === "remove"
-                    ? {op: "add", path, value: origValue}
-                    : {op: "replace", path, value: origValue}
-            )
-        })
-    }
+// This is used for both Map objects and normal objects.
+function generatePatchesFromAssigned(state, basePath, patches, inversePatches) {
+    const {base, copy} = state
+    each(state.assigned, (key, assignedValue) => {
+        const origValue = get(base, key)
+        const value = get(copy, key)
+        const op = !assignedValue
+            ? "remove"
+            : has(base, key)
+            ? "replace"
+            : "add"
+        if (origValue === value && op === "replace") return
+        const path = basePath.concat(key)
+        patches.push(op === "remove" ? {op, path} : {op, path, value})
+        inversePatches.push(
+            op === "add"
+                ? {op: "remove", path}
+                : op === "remove"
+                ? {op: "add", path, value: origValue}
+                : {op: "replace", path, value: origValue}
+        )
+    })
 }
 
 function generateSetPatches(state, basePath, patches, inversePatches) {
@@ -160,11 +151,7 @@ export function applyPatches(draft, patches) {
         } else {
             let base = draft
             for (let i = 0; i < path.length - 1; i++) {
-                if (isMap(base)) {
-                    base = base.get(path[i])
-                } else {
-                    base = base[path[i]]
-                }
+                base = get(base, path[i])
                 if (!base || typeof base !== "object")
                     throw new Error("Cannot apply patch, path doesn't resolve: " + path.join("/")) // prettier-ignore
             }
diff --git a/src/proxy.js b/src/proxy.js
index faec0a4..e1e312b 100644
--- a/src/proxy.js
+++ b/src/proxy.js
@@ -49,20 +49,26 @@ export function createProxy(base, parent) {
         revoke: null
     }
 
-    let proxyTarget = state
+    let target = state
     let traps = objectTraps
     if (Array.isArray(base)) {
-        proxyTarget = [state]
+        target = [state]
         traps = arrayTraps
-    } else if (isMap(base)) {
+    }
+    // Map drafts must support object keys, so we use Map objects to track changes.
+    else if (isMap(base)) {
         traps = mapTraps
-    } else if (isSet(base)) {
+        state.drafts = new Map()
+        state.assigned = new Map()
+    }
+    // Set drafts use a Map object to track which of its values are drafted.
+    // And we don't need the "assigned" property, because Set objects have no keys.
+    else if (isSet(base)) {
         traps = setTraps
-        // We use values of a Set as keys and objects do not support other objects as keys
         state.drafts = new Map()
     }
 
-    const {revoke, proxy} = Proxy.revocable(proxyTarget, traps)
+    const {revoke, proxy} = Proxy.revocable(target, traps)
 
     state.draft = proxy
     state.revoke = revoke
@@ -85,7 +91,7 @@ const objectTraps = {
             return drafts[prop]
         }
 
-        const value = source(state)[prop]
+        const value = latest(state)[prop]
         if (state.finalized || !isDraftable(value)) {
             return value
         }
@@ -101,10 +107,10 @@ const objectTraps = {
         return (drafts[prop] = createProxy(value, state))
     },
     has(state, prop) {
-        return prop in source(state)
+        return prop in latest(state)
     },
     ownKeys(state) {
-        return Reflect.ownKeys(source(state))
+        return Reflect.ownKeys(latest(state))
     },
     set(state, prop, value) {
         if (!state.modified) {
@@ -134,7 +140,7 @@ const objectTraps = {
     // Note: We never coerce `desc.value` into an Immer draft, because we can't make
     // the same guarantee in ES5 mode.
     getOwnPropertyDescriptor(state, prop) {
-        const owner = source(state)
+        const owner = latest(state)
         const desc = Reflect.getOwnPropertyDescriptor(owner, prop)
         if (desc) {
             desc.writable = true
@@ -196,64 +202,54 @@ const reflectTraps = makeReflectTraps([
 
 const mapTraps = makeTrapsForGetters({
     [DRAFT_STATE]: state => state,
-    size: state => source(state).size,
-    has(state, prop) {
-        state = source(state)
-        return state.has.bind(state)
-    },
+    size: state => latest(state).size,
+    has: state => key => latest(state).has(key),
     set: state => (key, value) => {
-        const stateSource = source(state)
-        if (!stateSource.has(key) || stateSource.get(key) !== value) {
+        const values = latest(state)
+        if (!values.has(key) || values.get(key) !== value) {
             markChanged(state)
-            state.assigned[key] = true
+            state.assigned.set(key, true)
             state.copy.set(key, value)
         }
         return state.draft
     },
     delete: state => key => {
-        markChanged(state)
-        state.assigned[key] = false
-        return state.copy.delete(key)
+        if (latest(state).has(key)) {
+            markChanged(state)
+            state.assigned.set(key, false)
+            return state.copy.delete(key)
+        }
+        return false
     },
     clear: state => () => {
         markChanged(state)
-        state.assigned = {}
-        for (const key of source(state).keys()) {
-            state.assigned[key] = false
+        state.assigned = new Map()
+        for (const key of latest(state).keys()) {
+            state.assigned.set(key, false)
         }
         return state.copy.clear()
     },
     forEach: (state, _, receiver) => (cb, thisArg) =>
-        source(state).forEach((_, key, map) => {
+        latest(state).forEach((_, key, map) => {
             const value = receiver.get(key)
             cb.call(thisArg, value, key, map)
         }),
     get: state => key => {
-        if (!state.modified && has(state.drafts, key)) {
-            return state.drafts[key]
-        }
-
-        if (state.modified && state.copy.has(key)) {
-            return state.copy.get(key)
+        const drafts = state[state.modified ? "copy" : "drafts"]
+        if (drafts.has(key)) {
+            return drafts.get(key)
         }
 
-        const value = source(state).get(key)
-
+        const value = latest(state).get(key)
         if (state.finalized || !isDraftable(value)) {
             return value
         }
 
         const draft = createProxy(value, state)
-
-        if (!state.modified) {
-            state.drafts[key] = draft
-        } else {
-            state.copy.set(key, draft)
-        }
-
+        drafts.set(key, draft)
         return draft
     },
-    keys: state => () => source(state).keys(),
+    keys: state => () => latest(state).keys(),
     values: iterateMapValues,
     entries: iterateMapValues,
     [Symbol.iterator]: iterateMapValues
@@ -261,17 +257,15 @@ const mapTraps = makeTrapsForGetters({
 
 /** Map.prototype.values _-or-_ Map.prototype.entries */
 function iterateMapValues(state, prop, receiver) {
-    const getYieldable =
-        prop === "values" ? (key, value) => value : (key, value) => [key, value]
-
+    const isEntries = prop !== "values"
     return () => {
-        const iterator = source(state)[Symbol.iterator]()
+        const iterator = latest(state)[Symbol.iterator]()
         return makeIterable(() => {
             const result = iterator.next()
             if (!result.done) {
                 const [key] = result.value
                 const value = receiver.get(key)
-                result.value = getYieldable(key, value)
+                result.value = isEntries ? [key, value] : value
             }
             return result
         })
@@ -284,13 +278,10 @@ function iterateMapValues(state, prop, receiver) {
 
 const setTraps = makeTrapsForGetters({
     [DRAFT_STATE]: state => state,
-    size: state => source(state).size,
-    has(state, prop) {
-        state = source(state)
-        return state.has.bind(state)
-    },
+    size: state => latest(state).size,
+    has: state => key => latest(state).has(key),
     add: state => value => {
-        if (!source(state).has(value)) {
+        if (!latest(state).has(value)) {
             markChanged(state)
             state.copy.add(value)
         }
@@ -319,17 +310,14 @@ const setTraps = makeTrapsForGetters({
 })
 
 function iterateSetValues(state, prop) {
+    const isEntries = prop === "entries"
     return () => {
-        const iterator = source(state)[Symbol.iterator]()
-
+        const iterator = latest(state)[Symbol.iterator]()
         return makeIterable(() => {
             const result = iterator.next()
             if (!result.done) {
-                const valueWrapped = wrapSetValue(state, result.value)
-                result.value = valueWrapped
-                if (prop === "entries") {
-                    result.value = [valueWrapped, valueWrapped]
-                }
+                const value = wrapSetValue(state, result.value)
+                result.value = isEntries ? [value, value] : value
             }
             return result
         })
@@ -356,8 +344,8 @@ function wrapSetValue(state, value) {
  * Helpers
  */
 
-// returns the object we should be reading the current value from, which is base, until some change has been made
-function source(state) {
+// Retrieve the latest values of the draft.
+function latest(state) {
     return state.copy || state.base
 }
 
@@ -365,7 +353,7 @@ function source(state) {
 function peek(draft, prop) {
     const state = draft[DRAFT_STATE]
     const desc = Reflect.getOwnPropertyDescriptor(
-        state ? source(state) : draft,
+        state ? latest(state) : draft,
         prop
     )
     return desc && desc.value
@@ -404,10 +392,10 @@ function makeIterable(next) {
     })
 }
 
-/** Create traps that all use the `Reflect` API on the `source(state)` */
+/** Create traps that all use the `Reflect` API on the `latest(state)` */
 function makeReflectTraps(names) {
     return names.reduce((traps, name) => {
-        traps[name] = (state, ...args) => Reflect[name](source(state), ...args)
+        traps[name] = (state, ...args) => Reflect[name](latest(state), ...args)
         return traps
     }, {})
 }


Test patch:
diff --git a/__tests__/__snapshots__/hooks.js.snap b/__tests__/__snapshots__/hooks.js.snap
index f703ed4..822fcb4 100644
--- a/__tests__/__snapshots__/hooks.js.snap
+++ b/__tests__/__snapshots__/hooks.js.snap
@@ -151,6 +151,91 @@ Array [
 ]
 `;
 
+exports[`hooks (proxy) - onAssign() when draft is a Map assign 1`] = `
+Array [
+  Array [
+    "a",
+    10,
+  ],
+  Array [
+    Object {
+      "prop": "val1",
+    },
+    11,
+  ],
+]
+`;
+
+exports[`hooks (proxy) - onAssign() when draft is a Map nested assignments 1`] = `
+Array [
+  Array [
+    "b",
+    2,
+  ],
+  Array [
+    Object {
+      "prop": "val1",
+    },
+    Map {
+      "b" => 2,
+      "d" => 1,
+    },
+  ],
+  Array [
+    "a",
+    Map {
+      Object {
+        "prop": "val1",
+      } => Map {
+        "b" => 2,
+        "d" => 1,
+      },
+    },
+  ],
+]
+`;
+
+exports[`hooks (proxy) - onAssign() when draft is a Set assign 1`] = `
+Array [
+  Array [
+    "a",
+    Set {
+      1,
+      2,
+      3,
+      4,
+    },
+  ],
+]
+`;
+
+exports[`hooks (proxy) - onAssign() when draft is a Set delete 1`] = `
+Array [
+  Array [
+    "a",
+    Set {
+      2,
+      3,
+    },
+  ],
+]
+`;
+
+exports[`hooks (proxy) - onAssign() when draft is a Set nested assignments 1`] = `
+Array [
+  Array [
+    "a",
+    Set {
+      "a",
+      Set {
+        1,
+        2,
+      },
+    },
+  ],
+]
+`;
+
 exports[`hooks (proxy) - onAssign() when draft is an array assign 1`] = `
 Array [
   Array [
@@ -259,6 +344,27 @@ Array [
 ]
 `;
 
+exports[`hooks (proxy) - onDelete() when draft is a Map - delete 1`] = `
+Array [
+  Array [
+    "a",
+  ],
+  Array [
+    Object {
+      "prop": "val1",
+    },
+  ],
+]
+`;
+
+exports[`hooks (proxy) - onDelete() when draft is a Map - nested deletions 1`] = `
+Array [
+  Array [
+    "b",
+  ],
+]
+`;
+
 exports[`hooks (proxy) - onDelete() when draft is an array - length = 0 1`] = `Array []`;
 
 exports[`hooks (proxy) - onDelete() when draft is an array - pop 1`] = `
diff --git a/__tests__/base.js b/__tests__/base.js
index 36694ab..025cb24 100644
--- a/__tests__/base.js
+++ b/__tests__/base.js
@@ -309,6 +309,16 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                     expect(nextState).toBe(baseState)
                 })
 
+                it("supports key access for non-primitive keys", () => {
+                    const key = {prop: "val"}
+                    const base = new Map([[key, {id: 1, a: 1}]])
+                    const value = base.get(key)
+                    const nextState = produce(base, s => {
+                        expect(s.get(key)).toEqual(value)
+                    })
+                    expect(nextState).toBe(base)
+                })
+
                 it("supports iteration", () => {
                     const base = new Map([
                         ["first", {id: 1, a: 1}],
@@ -392,24 +402,39 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                 })
 
                 it("supports forEach", () => {
+                    const key1 = {prop: "val1"}
+                    const key2 = {prop: "val2"}
                     const base = new Map([
                         ["first", {id: 1, a: 1}],
-                        ["second", {id: 2, a: 1}]
+                        ["second", {id: 2, a: 1}],
+                        [key1, {id: 3, a: 1}],
+                        [key2, {id: 4, a: 1}]
                     ])
                     const result = produce(base, draft => {
                         let sum1 = 0
                         draft.forEach(({a}) => {
                             sum1 += a
                         })
-                        expect(sum1).toBe(2)
+                        expect(sum1).toBe(4)
                         let sum2 = 0
                         draft.get("first").a = 10
                         draft.get("second").a = 20
+                        draft.get(key1).a = 30
+                        draft.get(key2).a = 40
                         draft.forEach(({a}) => {
                             sum2 += a
                         })
-                        expect(sum2).toBe(30)
+                        expect(sum2).toBe(100)
                     })
+                    expect(result).not.toBe(base)
+                    expect(base.get("first").a).toEqual(1)
+                    expect(base.get("second").a).toEqual(1)
+                    expect(base.get(key1).a).toEqual(1)
+                    expect(base.get(key2).a).toEqual(1)
+                    expect(result.get("first").a).toEqual(10)
+                    expect(result.get("second").a).toEqual(20)
+                    expect(result.get(key1).a).toEqual(30)
+                    expect(result.get(key2).a).toEqual(40)
                 })
 
                 it("supports forEach mutation", () => {
@@ -438,6 +463,18 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                     expect(nextState.aMap.get("force")).toEqual(true)
                 })
 
+                it("can assign by a non-primitive key", () => {
+                    const key = {prop: "val"}
+                    const value = {id: 1, a: 1}
+                    const base = new Map([[key, value]])
+                    const nextState = produce(base, s => {
+                        s.set(key, true)
+                    })
+                    expect(nextState).not.toBe(base)
+                    expect(base.get(key)).toEqual(value)
+                    expect(nextState.get(key)).toEqual(true)
+                })
+
                 it("state stays the same if the the same item is assigned by key", () => {
                     const nextState = produce(baseState, s => {
                         s.aMap.set("jediTotal", 42)
diff --git a/__tests__/hooks.js b/__tests__/hooks.js
index 759b41e..454630c 100644
--- a/__tests__/hooks.js
+++ b/__tests__/hooks.js
@@ -1,6 +1,7 @@
 "use strict"
 import {Immer, setUseProxies} from "../src/index"
 import matchers from "expect/build/matchers"
+import {isSet} from "../src/common"
 
 describe("hooks (proxy) -", () => createHookTests(true))
 describe("hooks (es5) -", () => createHookTests(false))
@@ -12,9 +13,9 @@ function createHookTests(useProxies) {
         ;({produce, onAssign, onDelete, onCopy} = new Immer({
             autoFreeze: true,
             useProxies,
-            onAssign: defuseProxies(jest.fn().mockName("onAssign")),
-            onDelete: defuseProxies(jest.fn().mockName("onDelete")),
-            onCopy: defuseProxies(jest.fn().mockName("onCopy"))
+            onAssign: jest.fn().mockName("onAssign"),
+            onDelete: jest.fn().mockName("onDelete"),
+            onCopy: jest.fn().mockName("onCopy")
         }))
     })
 
@@ -118,6 +119,92 @@ function createHookTests(useProxies) {
                 expect(calls).toEqual(["b", "d", "c", "x"])
             })
         })
+
+        if (useProxies) {
+            describe("when draft is a Map", () => {
+                test("assign", () => {
+                    const key1 = {prop: "val1"}
+                    const key2 = {prop: "val2"}
+                    produce(new Map([["a", 0], [key1, 1], [key2, 2]]), s => {
+                        s.set("a", 10)
+                        s.set(key1, 11)
+                    })
+                    expectCalls(onAssign)
+                })
+                test("assign (no change)", () => {
+                    produce(new Map([["a", 0]]), s => {
+                        s.set("a", 0)
+                    })
+                    expect(onAssign).not.toBeCalled()
+                })
+                test("delete", () => {
+                    produce(new Map([["a", 0]]), s => {
+                        s.delete("a")
+                    })
+                    expect(onAssign).not.toBeCalled()
+                })
+                test("nested assignments", () => {
+                    const key1 = {prop: "val1"}
+                    produce(
+                        new Map([
+                            [
+                                "a",
+                                new Map([
+                                    [
+                                        key1,
+                                        new Map([["b", 1], ["c", 1], ["d", 1]])
+                                    ]
+                                ])
+                            ]
+                        ]),
+                        s => {
+                            const nested = s.get("a").get(key1)
+                            nested.set("b", 2)
+                            nested.delete("c")
+                            nested.set("d", 1) // no-op
+                        }
+                    )
+                    expectCalls(onAssign)
+                })
+            })
+
+            describe("when draft is a Set", () => {
+                test("assign", () => {
+                    produce({a: new Set([1, 2, 3])}, s => {
+                        s.a.add(4)
+                    })
+                    expectCalls(onAssign)
+                })
+                test("assign (no change)", () => {
+                    produce({a: new Set([1, 2, 3])}, s => {
+                        s.a.add(3)
+                    })
+                    expect(onAssign).not.toBeCalled()
+                })
+                // Any mutation of a Set results in a new assignment. Including deletes.
+                test("delete", () => {
+                    produce({a: new Set([1, 2, 3])}, s => {
+                        s.a.delete(1)
+                    })
+                    expectCalls(onAssign)
+                })
+                test("nested assignments", () => {
+                    const val1 = {prop: "val1"}
+                    produce({a: new Set(["a", new Set([val1, 1])])}, s => {
+                        let nested
+                        s.a.forEach(value => {
+                            if (isSet(value)) {
+                                nested = value
+                            }
+                        })
+                        nested.add(2)
+                        nested.delete(val1)
+                        nested.add(1) // no-op
+                    })
+                    expectCalls(onAssign)
+                })
+            })
+        }
     })
 
     describe("onDelete()", () => {
@@ -163,21 +250,87 @@ function createHookTests(useProxies) {
                 expectCalls(onDelete)
             })
         })
+
+        if (useProxies) {
+            describe("when draft is a Map -", () => {
+                test("delete", () => {
+                    const key1 = {prop: "val1"}
+                    const key2 = {prop: "val2"}
+                    produce(new Map([["a", 0], [key1, 1], [key2, 2]]), s => {
+                        s.delete("a")
+                        s.delete(key1)
+                    })
+                    expectCalls(onDelete)
+                })
+                test("delete (no change)", () => {
+                    produce(new Map(), s => {
+                        s.delete("a")
+                    })
+                    expect(onDelete).not.toBeCalled()
+                })
+                test("nested deletions", () => {
+                    const key1 = {prop: "val1"}
+                    produce(
+                        new Map([
+                            ["a", new Map([[key1, new Map([["b", 1]])]])]
+                        ]),
+                        s => {
+                            s.get("a")
+                                .get(key1)
+                                .delete("b")
+                        }
+                    )
+                    expectCalls(onDelete)
+                })
+            })
+
+            describe("when draft is a Set -", () => {
+                test("delete", () => {
+                    produce({a: new Set([1])}, s => {
+                        s.a.delete(1)
+                    })
+                    expect(onDelete).not.toBeCalled()
+                })
+            })
+        }
     })
 
     describe("onCopy()", () => {
-        useSharedTests(() => onCopy)
-        it("is called in the right order", () => {
-            const calls = []
+        let calls
+        beforeEach(() => {
+            calls = []
             onCopy.mockImplementation(s => {
                 calls.push(s.base)
             })
+        })
+
+        useSharedTests(() => onCopy)
+        it("is called in the right order for objects", () => {
             const base = {a: {b: {c: 1}}}
             produce(base, s => {
                 delete s.a.b.c
             })
             expect(calls).toShallowEqual([base.a.b, base.a, base])
         })
+
+        if (useProxies) {
+            it("is called in the right order for Maps", () => {
+                const base = new Map([["a", new Map([["b", 0]])]])
+                produce(base, s => {
+                    s.get("a").delete("b")
+                })
+                expect(calls).toShallowEqual([base.get("a"), base])
+            })
+
+            it("is called in the right order for Sets", () => {
+                const item1 = {a: 0}
+                const base = new Set([item1])
+                produce(base, s => {
+                    s.forEach(item => item.a++)
+                })
+                expect(calls).toShallowEqual([item1, base])
+            })
+        }
     })
 
     function useSharedTests(getHook) {
@@ -205,15 +358,6 @@ function expectCalls(hook) {
     ).toMatchSnapshot()
 }
 
-// For defusing draft proxies.
-function defuseProxies(fn) {
-    return Object.assign((...args) => {
-        expect(args[0].finalized).toBeTruthy()
-        args[0].draft = args[0].drafts = null
-        fn(...args)
-    }, fn)
-}
-
 expect.extend({
     toShallowEqual(received, expected) {
         const match = matchers.toBe(received, expected)
diff --git a/__tests__/patch.js b/__tests__/patch.js
index 1b8f8b1..587b9ec 100644
--- a/__tests__/patch.js
+++ b/__tests__/patch.js
@@ -158,6 +158,27 @@ describe("simple assignment - 6", () => {
     )
 })
 
+describe("simple assignment - 7", () => {
+    const key1 = {prop: "val1"}
+    const key2 = {prop: "val2"}
+    runPatchTest(
+        {x: new Map([[key1, 4]])},
+        d => {
+            d.x.set(key1, 5)
+            d.x.set(key2, 6)
+        },
+        [
+            {op: "replace", path: ["x", key1], value: 5},
+            {op: "add", path: ["x", key2], value: 6}
+        ],
+        [
+            {op: "replace", path: ["x", key1], value: 4},
+            {op: "remove", path: ["x", key2]}
+        ],
+        true
+    )
+})
+
 describe("delete 1", () => {
     runPatchTest(
         {x: {y: 4}},
@@ -192,7 +213,25 @@ describe("delete 3", () => {
     )
 })
 
-describe("delete 4", () => {
+describe("delete 5", () => {
+    const key1 = {prop: "val1"}
+    const key2 = {prop: "val2"}
+    runPatchTest(
+        {x: new Map([[key1, 1], [key2, 2]])},
+        d => {
+            d.x.delete(key1)
+            d.x.delete(key2)
+        },
+        [{op: "remove", path: ["x", key1]}, {op: "remove", path: ["x", key2]}],
+        [
+            {op: "add", path: ["x", key1], value: 1},
+            {op: "add", path: ["x", key2], value: 2}
+        ],
+        true
+    )
+})
+
+describe("delete 6", () => {
     runPatchTest(
         new Set(["x", 1]),
         d => {
@@ -204,7 +243,7 @@ describe("delete 4", () => {
     )
 })
 
-describe("delete 5", () => {
+describe("delete 7", () => {
     runPatchTest(
         {x: new Set(["y", 1])},
         d => {