feat: support async recipes

Drafts won't be revoked until the returned promise is fulfilled or rejected.

As suggested by @Yurickh in immerjs#302

__tests__/base.js

@@ -815,7 +815,7 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
                 expect(next.obj).toBe(next.arr[0])
             })
 
-            it("can return an object with two references to any pristine draft", () => {
+            it("can return an object with two references to an unmodified draft", () => {
                 const base = {a: {}}
                 const next = produce(base, d => {
                     return [d.a, d.a]
@@ -833,6 +833,36 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
             })
         })
 
+        describe("async recipe function", () => {
+            it("can modify the draft", async () => {
+                const base = {a: 0, b: 0}
+                const res = await produce(base, async d => {
+                    d.a = 1
+                    await Promise.resolve()
+                    d.b = 1
+                })
+                expect(res).not.toBe(base)
+                expect(res).toEqual({a: 1, b: 1})
+            })
+
+            it("works with rejected promises", async () => {
+                let draft
+                const base = {a: 0, b: 0}
+                const err = new Error("passed")
+                try {
+                    await produce(base, async d => {
+                        draft = d
+                        draft.b = 1
+                        await Promise.reject(err)
+                    })
+                    throw "failed"
+                } catch (e) {
+                    expect(e).toBe(err)
+                    expect(() => draft.a).toThrowError(/revoked/)
+                }
+            })
+        })
+
         it("throws when the draft is modified and another object is returned", () => {
             const base = {x: 3}
             expect(() => {

src/es5.js

@@ -1,6 +1,4 @@
 "use strict"
-// @ts-check
-
 import {
     each,
     has,
@@ -11,20 +9,20 @@ import {
     shallowCopy,
     DRAFT_STATE
 } from "./common"
+import {ImmerScope} from "./scope"
 
 const descriptors = {}
 
-// For nested produce calls:
-export const scopes = []
-export const currentScope = () => scopes[scopes.length - 1]
-
-export function willFinalize(result, baseDraft, needPatches) {
-    const scope = currentScope()
-    scope.forEach(state => (state.finalizing = true))
-    if (result === undefined || result === baseDraft) {
-        if (needPatches) markChangesRecursively(baseDraft)
+export function willFinalize(scope, result, isReplaced) {
+    scope.drafts.forEach(draft => {
+        draft[DRAFT_STATE].finalizing = true
+    })
+    if (!isReplaced) {
+        if (scope.patches) {
+            markChangesRecursively(scope.drafts[0])
+        }
         // This is faster when we don't care about which attributes changed.
-        markChangesSweep(scope)
+        markChangesSweep(scope.drafts)
     }
 }
 
@@ -36,8 +34,9 @@ export function createDraft(base, parent) {
     })
 
     // See "proxy.js" for property documentation.
+    const scope = parent ? parent.scope : ImmerScope.current
     const state = {
-        scope: parent ? parent.scope : currentScope(),
+        scope,
         modified: false,
         finalizing: false, // es5 only
         finalized: false,
@@ -51,7 +50,7 @@ export function createDraft(base, parent) {
     }
 
     createHiddenProperty(draft, DRAFT_STATE, state)
-    state.scope.push(state)
+    scope.drafts.push(draft)
     return draft
 }
 
@@ -135,14 +134,14 @@ function assertUnrevoked(state) {
 }
 
 // This looks expensive, but only proxies are visited, and only objects without known changes are scanned.
-function markChangesSweep(scope) {
+function markChangesSweep(drafts) {
     // The natural order of drafts in the `scope` array is based on when they
     // were accessed. By processing drafts in reverse natural order, we have a
     // better chance of processing leaf nodes first. When a leaf node is known to
     // have changed, we can avoid any traversal of its ancestor nodes.
-    for (let i = scope.length - 1; i >= 0; i--) {
-        const state = scope[i]
-        if (state.modified === false) {
+    for (let i = drafts.length - 1; i >= 0; i--) {
+        const state = drafts[i][DRAFT_STATE]
+        if (!state.modified) {
             if (Array.isArray(state.base)) {
                 if (hasArrayChanges(state)) markChanged(state)
             } else if (hasObjectChanges(state)) markChanged(state)

src/immer.js

@@ -13,6 +13,7 @@ import {
     DRAFT_STATE,
     NOTHING
 } from "./common"
+import {ImmerScope} from "./scope"
 
 function verifyMinified() {}
 
@@ -51,62 +52,37 @@ export class Immer {
         }
 
         let result
-        // Only create proxies for plain objects/arrays.
-        if (!isDraftable(base)) {
-            result = recipe(base)
-            if (result === undefined) return base
-        }
-        // The given value must be proxied.
-        else {
-            this.scopes.push([])
+
+        // Only plain objects, arrays, and "immerable classes" are drafted.
+        if (isDraftable(base)) {
+            const scope = ImmerScope.enter()
             const baseDraft = this.createDraft(base)
             try {
                 result = recipe.call(baseDraft, baseDraft)
-                this.willFinalize(result, baseDraft, !!patchListener)
-
-                // Never generate patches when no listener exists.
-                var patches = patchListener && [],
-                    inversePatches = patchListener && []
-
-                // Finalize the modified draft...
-                if (result === undefined || result === baseDraft) {
-                    result = this.finalize(
-                        baseDraft,
-                        [],
-                        patches,
-                        inversePatches
-                    )
-                }
-                // ...or use a replacement value.
-                else {
-                    // Users must never modify the draft _and_ return something else.
-                    if (baseDraft[DRAFT_STATE].modified)
-                        throw new Error("An immer producer returned a new value *and* modified its draft. Either return a new value *or* modify the draft.") // prettier-ignore
-
-                    // Finalize the replacement in case it contains (or is) a subset of the draft.
-                    if (isDraftable(result)) result = this.finalize(result)
-
-                    if (patchListener) {
-                        patches.push({
-                            op: "replace",
-                            path: [],
-                            value: result
-                        })
-                        inversePatches.push({
-                            op: "replace",
-                            path: [],
-                            value: base
-                        })
+            } catch (error) {
+                scope.revoke()
+                throw error
+            }
+            scope.leave()
+            if (result instanceof Promise) {
+                return result.then(
+                    result => {
+                        scope.usePatches(patchListener)
+                        return this.processResult(result, scope)
+                    },
+                    error => {
+                        scope.revoke()
+                        throw error
                     }
-                }
-            } finally {
-                this.currentScope().forEach(state => state.revoke())
-                this.scopes.pop()
+                )
             }
-            patchListener && patchListener(patches, inversePatches)
+            scope.usePatches(patchListener)
+            return this.processResult(result, scope)
+        } else {
+            result = recipe(base)
+            if (result === undefined) return base
+            return result !== NOTHING ? result : undefined
         }
-        // Normalize the result.
-        return result === NOTHING ? undefined : result
     }
     setAutoFreeze(value) {
         this.autoFreeze = value
@@ -123,25 +99,64 @@ export class Immer {
         // Otherwise, produce a copy of the base state.
         return this.produce(base, draft => applyPatches(draft, patches))
     }
+    /** @internal */
+    processResult(result, scope) {
+        const baseDraft = scope.drafts[0]
+        const isReplaced = result !== undefined && result !== baseDraft
+        this.willFinalize(scope, result, isReplaced)
+        if (isReplaced) {
+            if (baseDraft[DRAFT_STATE].modified) {
+                scope.revoke()
+                throw new Error("An immer producer returned a new value *and* modified its draft. Either return a new value *or* modify the draft.") // prettier-ignore
+            }
+            if (isDraftable(result)) {
+                // Finalize the result in case it contains (or is) a subset of the draft.
+                result = this.finalize(result, null, scope)
+            }
+            if (scope.patches) {
+                scope.patches.push({
+                    op: "replace",
+                    path: [],
+                    value: result
+                })
+                scope.inversePatches.push({
+                    op: "replace",
+                    path: [],
+                    value: baseDraft[DRAFT_STATE].base
+                })
+            }
+        } else {
+            // Finalize the base draft.
+            result = this.finalize(baseDraft, [], scope)
+        }
+        scope.revoke()
+        if (scope.patches) {
+            scope.patchListener(scope.patches, scope.inversePatches)
+        }
+        return result !== NOTHING ? result : undefined
+    }
     /**
      * @internal
      * Finalize a draft, returning either the unmodified base state or a modified
      * copy of the base state.
      */
-    finalize(draft, path, patches, inversePatches) {
+    finalize(draft, path, scope) {
         const state = draft[DRAFT_STATE]
         if (!state) {
             if (Object.isFrozen(draft)) return draft
-            return this.finalizeTree(draft)
+            return this.finalizeTree(draft, null, scope)
         }
-        // Never finalize drafts owned by an outer scope.
-        if (state.scope !== this.currentScope()) {
+        // Never finalize drafts owned by another scope.
+        if (state.scope !== scope) {
             return draft
         }
-        if (!state.modified) return state.base
+        if (!state.modified) {
+            return state.base
+        }
         if (!state.finalized) {
             state.finalized = true
-            this.finalizeTree(state.draft, path, patches, inversePatches)
+            this.finalizeTree(state.draft, path, scope)
+
             if (this.onDelete) {
                 // The `assigned` object is unreliable with ES5 drafts.
                 if (this.useProxies) {
@@ -156,23 +171,32 @@ export class Immer {
                     })
                 }
             }
-            if (this.onCopy) this.onCopy(state)
+            if (this.onCopy) {
+                this.onCopy(state)
+            }
 
             // Nested producers must never auto-freeze their result,
             // because it may contain drafts from parent producers.
-            if (this.autoFreeze && this.scopes.length === 1) {
+            if (this.autoFreeze && !scope.parent) {
                 Object.freeze(state.copy)
             }
 
-            if (patches) generatePatches(state, path, patches, inversePatches)
+            if (path && scope.patches) {
+                generatePatches(
+                    state,
+                    path,
+                    scope.patches,
+                    scope.inversePatches
+                )
+            }
         }
         return state.copy
     }
     /**
      * @internal
      * Finalize all drafts in the given state tree.
      */
-    finalizeTree(root, path, patches, inversePatches) {
+    finalizeTree(root, path, scope) {
         const state = root[DRAFT_STATE]
         if (state) {
             if (!this.useProxies) {
@@ -193,11 +217,14 @@ export class Immer {
             const inDraft = !!state && parent === root
 
             if (isDraft(value)) {
-                value =
-                    // Patches are never generated for assigned properties.
-                    patches && inDraft && !state.assigned[prop]
-                        ? this.finalize(value, path.concat(prop), patches, inversePatches) // prettier-ignore
-                        : this.finalize(value)
+                const needPatches =
+                    inDraft && path && scope.patches && !state.assigned[prop]
+
+                value = this.finalize(
+                    value,
+                    needPatches ? path.concat(prop) : null,
+                    scope
+                )
 
                 // Preserve non-enumerable properties.
                 if (Array.isArray(parent) || isEnumerable(parent, prop)) {

src/proxy.js

@@ -1,6 +1,4 @@
 "use strict"
-// @ts-check
-
 import {
     assign,
     each,
@@ -11,18 +9,16 @@ import {
     shallowCopy,
     DRAFT_STATE
 } from "./common"
-
-// For nested produce calls:
-export const scopes = []
-export const currentScope = () => scopes[scopes.length - 1]
+import {ImmerScope} from "./scope"
 
 // Do nothing before being finalized.
 export function willFinalize() {}
 
 export function createDraft(base, parent) {
+    const scope = parent ? parent.scope : ImmerScope.current
     const state = {
         // Track which produce call this is associated with.
-        scope: parent ? parent.scope : currentScope(),
+        scope,
         // True for both shallow and deep changes.
         modified: false,
         // Used during finalization.
@@ -50,7 +46,7 @@ export function createDraft(base, parent) {
     state.draft = proxy
     state.revoke = revoke
 
-    state.scope.push(state)
+    scope.drafts.push(proxy)
     return proxy
 }