Merge 08f795a into 62614f0

unadlib · Mar 26, 2023 · f2df3d1 · f2df3d1
2 parents 62614f0 + 08f795a
commit f2df3d1
Show file tree

Hide file tree

Showing 15 changed files with 296 additions and 204 deletions.
diff --git a/README.md b/README.md
@@ -146,7 +146,7 @@ In this basic example, the changes to the draft are 'mutative' within the draft
 
   > Forbid accessing non-draftable values in strict mode(unless using [unsafe()](#unsafe)).
 
-  > It is recommended to enable `strict` in development mode and disable `strict` in production mode. This will ensure safe returns and also keep good performance in the production build. If the return value is mixed drafts or `undefined`, then use [safeReturn()](#safereturn).
+  > It is recommended to enable `strict` in development mode and disable `strict` in production mode. This will ensure safe explicit returns and also keep good performance in the production build. If the value that does not mix any current draft or is `undefined` is returned, then use [rawReturn()](#rawreturn).
 
 - enablePatches - `boolean | { pathAsArray?: boolean; arrayLengthAssignment?: boolean; }`, the default is false.
 
@@ -302,19 +302,33 @@ expect(isDraftable(baseState.list)).toBeTruthy();
 
 > You can set a mark to determine if the value is draftable, and the mark function should be the same as passing in `create()` mark option.
 
-### `safeReturn()`
+### `rawReturn()`
 
-It is used as a safe return value to ensure that this value replaces the finalized draft value or use it to return `undefined` explicitly.
+It is used as a raw return value to ensure that this value replaces the finalized draft value or use it to return `undefined` explicitly.
 
 ```ts
 const baseState = { id: 'test' };
 const state = create(baseState as { id: string } | undefined, (draft) => {
-  return safeReturn(undefined);
+  return rawReturn(undefined);
 });
 expect(state).toBe(undefined);
 ```
 
-> You don't need to use `safeReturn()` when the return value doesn't have any draft.
+> You don't need to use `rawReturn()` when the return value have any draft.
+
+```ts
+const baseState = { a: 1, b: { c: 1 } };
+const state = create(baseState, (draft) => {
+  if (draft.b.c === 1) {
+    return {
+      ...draft,
+      a: 2,
+    };
+  }
+});
+expect(state).toEqual({ a: 2, b: { c: 1 } });
+expect(isDraft(state.b)).toBeFalsy();
+```
 
 [View more API docs](./docs/README.md).
 
@@ -440,10 +454,10 @@ const nextState = produce(baseState, (draft) => {
 Use Mutative
 
 ```ts
-import { create, safeReturn } from 'mutative';
+import { create, rawReturn } from 'mutative';
 
 const nextState = create(baseState, (draft) => {
-  return safeReturn(undefined);
+  return rawReturn(undefined);
 });
 ```
 

diff --git a/src/constant.ts b/src/constant.ts
@@ -1,5 +1,6 @@
 // Don't use `Symbol()` just for 3rd party access the draft
 export const PROXY_DRAFT = Symbol.for('__MUTATIVE_PROXY_DRAFT__');
+export const RAW_RETURN_SYMBOL = Symbol('__MUTATIVE_RAW_RETURN_SYMBOL__');
 
 export const iteratorSymbol: typeof Symbol.iterator = Symbol.iterator;
 

diff --git a/src/create.ts b/src/create.ts
@@ -14,7 +14,7 @@ import {
   revokeProxy,
 } from './utils';
 import { current, handleReturnValue } from './current';
-import { safeReturnValue } from './safeReturn';
+import { RAW_RETURN_SYMBOL } from './constant';
 
 /**
  * `create(baseState, callback, options)` to create the next state
@@ -108,7 +108,6 @@ function create(arg0: any, arg1: any, arg2?: any): any {
     strict,
     enablePatches,
   };
-  safeReturnValue.length = 0;
   if (
     !isDraftable(state, _options) &&
     typeof state === 'object' &&
@@ -146,16 +145,17 @@ function create(arg0: any, arg1: any, arg2?: any): any {
           `Either the value is returned as a new non-draft value, or only the draft is modified without returning any value.`
         );
       }
-      if (safeReturnValue.length) {
-        const _value = safeReturnValue.pop();
-        if (typeof value === 'object' && value !== null) {
-          handleReturnValue(value);
+      const rawReturnValue = value?.[RAW_RETURN_SYMBOL] as [any] | undefined;
+      if (rawReturnValue) {
+        const _value = rawReturnValue[0];
+        if (_options.strict && typeof value === 'object' && value !== null) {
+          handleReturnValue(proxyDraft, value, true);
         }
         return finalize([_value]);
       }
       if (value !== undefined) {
-        if (_options.strict && typeof value === 'object' && value !== null) {
-          handleReturnValue(value, true);
+        if (typeof value === 'object' && value !== null) {
+          handleReturnValue(proxyDraft, value);
         }
         return finalize([value]);
       }

diff --git a/src/current.ts b/src/current.ts
@@ -1,4 +1,4 @@
-import { DraftType } from './interface';
+import { DraftType, ProxyDraft } from './interface';
 import {
   forEach,
   get,
@@ -11,13 +11,20 @@ import {
   shallowCopy,
 } from './utils';
 
-export function handleReturnValue<T extends object>(value: T, warning = false) {
+export function handleReturnValue<T extends object>(
+  rootDraft: ProxyDraft<any> | undefined,
+  value: T,
+  warning = false
+) {
+  let isContainDraft = false;
   forEach(value, (key, item, source) => {
     const proxyDraft = getProxyDraft(item);
-    if (proxyDraft) {
+    // just handle the draft which is created by the same rootDraft
+    if (proxyDraft && proxyDraft?.finalities === rootDraft?.finalities) {
+      isContainDraft = true;
       if (__DEV__ && warning) {
         console.warn(
-          `The return value contains drafts, please use safeReturn() to wrap the return value.`
+          `The return value contains drafts, please don't use 'rawReturn()' to wrap the return value.`
         );
       }
       const currentValue = proxyDraft.original;
@@ -31,9 +38,14 @@ export function handleReturnValue<T extends object>(value: T, warning = false) {
         set(source, key, currentValue);
       }
     } else if (typeof item === 'object' && item !== null) {
-      handleReturnValue(item, warning);
+      handleReturnValue(rootDraft, item, warning);
     }
   });
+  if (__DEV__ && warning && !isContainDraft) {
+    console.warn(
+      `The return value does not contain any draft, please use 'rawReturn()' to wrap the return value to improve performance.`
+    );
+  }
 }
 
 function getCurrent(target: any) {

diff --git a/src/index.ts b/src/index.ts
@@ -3,7 +3,7 @@ export { apply } from './apply';
 export { original } from './original';
 export { current } from './current';
 export { unsafe } from './unsafe';
-export { safeReturn } from './safeReturn';
+export { rawReturn } from './rawReturn';
 export { isDraft } from './utils/draft';
 export { isDraftable } from './utils/draft';
 

diff --git a/src/rawReturn.ts b/src/rawReturn.ts
@@ -0,0 +1,40 @@
+import { RAW_RETURN_SYMBOL } from './constant';
+
+/**
+ * Use rawReturn() to wrap the return value to skip the draft check and thus improve performance.
+ *
+ * ## Example
+ *
+ * ```ts
+ * import { create, rawReturn } from '../index';
+ *
+ * const baseState = { foo: { bar: 'str' }, arr: [] };
+ * const state = create(
+ *   baseState,
+ *   (draft) => {
+ *     return rawReturn(baseState);
+ *   },
+ * );
+ * expect(state).toBe(baseState);
+ * ```
+ */
+export function rawReturn<T extends object | undefined>(value: T): T {
+  if (arguments.length === 0) {
+    throw new Error('rawReturn() must be called with a value.');
+  }
+  if (arguments.length > 1) {
+    throw new Error('rawReturn() must be called with one argument.');
+  }
+  if (
+    __DEV__ &&
+    value !== undefined &&
+    (typeof value !== 'object' || value === null)
+  ) {
+    console.warn(
+      'rawReturn() must be called with an object or `undefined`, other types do not need to be returned via rawReturn().'
+    );
+  }
+  return {
+    [RAW_RETURN_SYMBOL]: [value],
+  } as never;
+}
diff --git a/src/safeReturn.ts b/src/safeReturn.ts
diff --git a/test/immer/base.test.ts b/test/immer/base.test.ts
@@ -1,6 +1,6 @@
 // @ts-nocheck
 'use strict';
-import { create, original, isDraft, safeReturn } from '../../src';
+import { create, original, isDraft, rawReturn } from '../../src';
 import deepFreeze from 'deep-freeze';
 import * as lodash from 'lodash';
 import { getType } from '../../src/utils';
@@ -1531,28 +1531,26 @@ function runBaseTest(
         const base = {};
         const result = create(base, (s1) =>
           // ! it's different from mutative
-          safeReturn(
-            create(
-              { s1 },
-              (s2) => {
-                expect(original(s2.s1)).toBe(s1);
-                s2.n = 1;
-                s2.s1 = create(
-                  { s2 },
-                  (s3) => {
-                    expect(original(s3.s2)).toBe(s2);
-                    expect(original(s3.s2.s1)).toBe(s2.s1);
-                    return s3.s2.s1;
-                  },
-                  // ! it's different from mutative
-                  { enableAutoFreeze: false }
-                );
-              },
-              {
+          create(
+            { s1 },
+            (s2) => {
+              expect(original(s2.s1)).toBe(s1);
+              s2.n = 1;
+              s2.s1 = create(
+                { s2 },
+                (s3) => {
+                  expect(original(s3.s2)).toBe(s2);
+                  expect(original(s3.s2.s1)).toBe(s2.s1);
+                  return s3.s2.s1;
+                },
                 // ! it's different from mutative
-                enableAutoFreeze: false,
-              }
-            )
+                { enableAutoFreeze: false }
+              );
+            },
+            {
+              // ! it's different from mutative
+              enableAutoFreeze: false,
+            }
           )
         );
         expect(result.n).toBe(1);
@@ -1805,8 +1803,7 @@ function runBaseTest(
       it('can return an object with two references to an unmodified draft', () => {
         const base = { a: {} };
         const next = produce(base, (d) => {
-          // ! it's different from mutative
-          return safeReturn([d.a, d.a]);
+          return [d.a, d.a];
         });
         expect(next[0]).toBe(base.a);
         expect(next[0]).toBe(next[1]);
@@ -1928,7 +1925,7 @@ function runBaseTest(
         produce(state, (draft) => {
           switch (action.type) {
             case 'SET_STARTING_DOTS':
-              return safeReturn(draft.availableStartingDots.map((a) => a));
+              return draft.availableStartingDots.map((a) => a);
             default:
               break;
           }
@@ -1951,9 +1948,9 @@ function runBaseTest(
         produce(state, (draft) => {
           switch (action.type) {
             case 'SET_STARTING_DOTS':
-              return safeReturn({
+              return {
                 dots: draft.availableStartingDots.map((a) => a),
-              });
+              };
             default:
               break;
           }
@@ -2053,15 +2050,15 @@ function runBaseTest(
       expect(create(base, () => null)).toBe(null);
       expect(create(base, () => undefined)).toBe(3);
       expect(create(base, () => {})).toBe(3);
-      expect(create(base, () => safeReturn(undefined))).toBe(undefined);
+      expect(create(base, () => rawReturn(undefined))).toBe(undefined);
 
       expect(create({}, () => undefined)).toEqual({});
-      expect(create({}, () => safeReturn(undefined))).toBe(undefined);
-      expect(create(3, () => safeReturn(undefined))).toBe(undefined);
+      expect(create({}, () => rawReturn(undefined))).toBe(undefined);
+      expect(create(3, () => rawReturn(undefined))).toBe(undefined);
 
       expect(create(() => undefined)({})).toEqual({});
-      expect(create(() => safeReturn(undefined))({})).toBe(undefined);
-      expect(create(() => safeReturn(undefined))(3)).toBe(undefined);
+      expect(create(() => rawReturn(undefined))({})).toBe(undefined);
+      expect(create(() => rawReturn(undefined))(3)).toBe(undefined);
     });
 
     describe('base state type', () => {