unadlib · unadlib · Jan 25, 2023 · Jan 23, 2023 · Jan 24, 2023 · Jan 24, 2023
diff --git a/package.json b/package.json
@@ -20,7 +20,7 @@
     "clean": "rimraf dist",
     "test:coverage": "jest --coverage && coveralls < coverage/lcov.info",
     "perf": "cd test/performance && ts-node add-data.ts && ts-node todo.ts && ts-node incremental.ts",
-    "benchmark": "cd test/performance && ts-node benchmark.ts",
+    "benchmark": "yarn build && ts-node test/performance/benchmark.ts",
     "performance:basic": "cd test/performance && ts-node index.ts",
     "performance:set-map": "cd test/performance && ts-node set-map.ts",
     "performance:big-object": "cd test/performance && ts-node big-object.ts",

diff --git a/src/apply.ts b/src/apply.ts
@@ -31,6 +31,17 @@ export function apply<T extends object, F extends boolean = false>(
     Exclude<keyof Options<boolean, F>, 'enablePatches'>
   >
 ) {
+  let i: number;
+  for (i = patches.length - 1; i >= 0; i -= 1) {
+    const { value, op, path } = patches[i];
+    if (!path.length && op === Operation.Replace) {
+      state = value;
+      break;
+    }
+  }
+  if (i > -1) {
+    patches = patches.slice(i + 1);
+  }
   const mutate = (draft: Draft<T>) => {
     patches.forEach((patch) => {
       const { path, op } = patch;
@@ -39,8 +50,10 @@ export function apply<T extends object, F extends boolean = false>(
         const parentType = getType(base);
         const key = String(path[index]);
         if (
-          (parentType === DraftType.Object || parentType === DraftType.Array) &&
-          (key === '__proto__' || key === 'constructor')
+          ((parentType === DraftType.Object ||
+            parentType === DraftType.Array) &&
+            (key === '__proto__' || key === 'constructor')) ||
+          (typeof base === 'function' && key === 'prototype')
         ) {
           throw new Error(
             `Patching reserved attributes like __proto__ and constructor is not allowed.`
@@ -103,7 +116,7 @@ export function apply<T extends object, F extends boolean = false>(
     });
   };
   if (isDraft(state)) {
-    if (typeof applyOptions !== 'undefined') {
+    if (applyOptions !== undefined) {
       throw new Error(`Cannot apply patches with options to a draft.`);
     }
     mutate(state as Draft<T>);

diff --git a/src/create.ts b/src/create.ts
@@ -1,8 +1,14 @@
 import { CreateResult, Draft, Options, Result } from './interface';
 import { draftify } from './draftify';
-import { dataTypes } from './constant';
-import { getProxyDraft, isDraft, revokeProxy } from './utils';
-import { current } from './current';
+import {
+  getProxyDraft,
+  isDraft,
+  isDraftable,
+  isEqual,
+  revokeProxy,
+} from './utils';
+import { current, handleReturnValue } from './current';
+import { safeReturnValue } from './safeReturn';
 
 /**
  * `create(baseState, callback, options)` to create the next state
@@ -27,17 +33,17 @@ import { current } from './current';
  * ```
  */
 function create<
-  T extends object,
+  T extends any,
   F extends boolean = false,
   O extends boolean = false,
-  R extends void | Promise<void> = void
+  R extends void | Promise<void> | T | Promise<T> = void
 >(
   base: T,
   mutate: (draft: Draft<T>) => R,
   options?: Options<O, F>
 ): CreateResult<T, O, F, R>;
 function create<
-  T extends object,
+  T extends any,
   F extends boolean = false,
   O extends boolean = false,
   R extends void | Promise<void> = void
@@ -47,7 +53,7 @@ function create<
   options?: Options<O, F>
 ): CreateResult<T, O, F, R>;
 function create<
-  T extends object,
+  T extends any,
   P extends any[] = [],
   F extends boolean = false,
   O extends boolean = false,
@@ -57,7 +63,7 @@ function create<
   options?: Options<O, F>
 ): (base: T, ...args: P) => CreateResult<T, O, F, R>;
 function create<
-  T extends object,
+  T extends any,
   O extends boolean = false,
   F extends boolean = false
 >(base: T, options?: Options<O, F>): [T, () => Result<T, O, F>];
@@ -86,40 +92,87 @@ function create(arg0: any, arg1: any, arg2?: any): any {
     );
   }
   const state = isDraft(base) ? current(base) : base;
-  if (options?.mark?.(state, dataTypes) === dataTypes.mutable) {
-    const finalization = options.enablePatches ? [state, [], []] : state;
-    if (typeof arg1 !== 'function') {
-      return [state, () => finalization];
-    }
-    const result = mutate(state);
-    if (result instanceof Promise) {
-      return result.then(() => finalization);
-    }
-    return finalization;
+  const mark = options?.mark;
+  const enablePatches = options?.enablePatches ?? false;
+  const strict = options?.strict ?? false;
+  const enableAutoFreeze = options?.enableAutoFreeze ?? false;
+  const _options = {
+    enableAutoFreeze,
+    mark,
+    strict,
+    enablePatches,
+  };
+  safeReturnValue.length = 0;
+  if (
+    !isDraftable(state, _options) &&
+    typeof state === 'object' &&
+    state !== null
+  ) {
+    throw new Error(
+      `Invalid base state: create() only supports plain objects, arrays, Set, Map or using mark() to mark the state as immutable.`
+    );
   }
-  const [draft, finalize] = draftify(state, options);
+  const [draft, finalize] = draftify(state, _options);
   if (typeof arg1 !== 'function') {
+    if (!isDraftable(state, _options)) {
+      throw new Error(
+        `Invalid base state: create() only supports plain objects, arrays, Set, Map or using mark() to mark the state as immutable.`
+      );
+    }
     return [draft, finalize];
   }
   let result;
   try {
     result = mutate(draft);
   } catch (error) {
-    revokeProxy(getProxyDraft(draft)!);
+    revokeProxy(getProxyDraft(draft));
     throw error;
   }
+  const returnValue = (value: any) => {
+    const proxyDraft = getProxyDraft(draft)!;
+    if (!isDraft(value)) {
+      if (
+        value !== undefined &&
+        !isEqual(value, draft) &&
+        proxyDraft?.operated
+      ) {
+        throw new Error(
+          `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);
+        }
+        return finalize([_value]);
+      }
+      if (value !== undefined) {
+        if (_options.strict && typeof value === 'object' && value !== null) {
+          handleReturnValue(value, true);
+        }
+        return finalize([value]);
+      }
+    }
+    if (value === draft || value === undefined) {
+      return finalize([]);
+    }
+    const returnedProxyDraft = getProxyDraft(value)!;
+    if (_options === returnedProxyDraft.options) {
+      if (returnedProxyDraft.operated) {
+        throw new Error(`Cannot return a modified child draft.`);
+      }
+      return finalize([current(value)]);
+    }
+    return finalize([value]);
+  };
   if (result instanceof Promise) {
-    return result.then(finalize, (error) => {
+    return result.then(returnValue, (error) => {
       revokeProxy(getProxyDraft(draft)!);
       throw error;
     });
   }
-  if (result !== undefined) {
-    throw new Error(
-      `The create() callback must return 'void' or 'Promise<void>'.`
-    );
-  }
-  return finalize();
+  return returnValue(result);
 }
 
 export { create };
diff --git a/src/current.ts b/src/current.ts
@@ -11,6 +11,29 @@ import {
   shallowCopy,
 } from './utils';
 
+export function handleReturnValue<T extends object>(value: T, warning = false) {
+  forEach(value, (key, item, source) => {
+    const proxyDraft = getProxyDraft(item);
+    if (proxyDraft) {
+      if (warning) {
+        console.warn(
+          `The return value contains drafts, please use safeReturn() to wrap the return value.`
+        );
+      }
+      const currentValue = proxyDraft.original;
+      if (source instanceof Set) {
+        const arr = Array.from(source);
+        source.clear();
+        arr.forEach((item) => source.add(key === item ? currentValue : item));
+      } else {
+        set(source, key, currentValue);
+      }
+    } else if (typeof item === 'object' && item !== null) {
+      handleReturnValue(item, warning);
+    }
+  });
+}
+
 function getCurrent(target: any) {
   const proxyDraft = getProxyDraft(target);
   if (!isDraftable(target, proxyDraft?.options)) return target;

diff --git a/src/draft.ts b/src/draft.ts
@@ -4,6 +4,7 @@ import {
   Patches,
   ProxyDraft,
   Options,
+  Operation,
 } from './interface';
 import { dataTypes, PROXY_DRAFT } from './constant';
 import { mapHandler, mapHandlerKeys } from './map';
@@ -131,6 +132,7 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
     if (currentProxyDraft && isEqual(currentProxyDraft.original, value)) {
       // !case: ignore the case of assigning the original draftable value to a draft
       target.copy![key] = value;
+      target.assignedMap = target.assignedMap ?? new Map();
       target.assignedMap.set(key, false);
       return true;
     }
@@ -144,9 +146,9 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
     markChanged(target);
     if (has(target.original, key) && isEqual(value, target.original[key])) {
       // !case: handle the case of assigning the original non-draftable value to a draft
-      target.assignedMap.delete(key);
+      target.assignedMap!.delete(key);
     } else {
-      target.assignedMap.set(key, true);
+      target.assignedMap!.set(key, true);
     }
     target.copy![key] = value;
     markFinalization(target, key, value);
@@ -186,8 +188,9 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
       // !case: delete an existing key
       ensureShallowCopy(target);
       markChanged(target);
-      target.assignedMap.set(key, false);
+      target.assignedMap!.set(key, false);
     } else {
+      target.assignedMap = target.assignedMap ?? new Map();
       // The original non-existent key has been deleted
       target.assignedMap.delete(key);
     }
@@ -215,27 +218,26 @@ export function createDraft<T extends object>(createDraftOptions: {
     proxy: null,
     finalities,
     options,
-    assignedMap: new Map(),
     // Mapping of draft Set items to their corresponding draft values.
     setMap:
       type === DraftType.Set
         ? new Map((original as Set<any>).entries())
         : undefined,
   };
   // !case: undefined as a draft map key
-  if (Object.hasOwnProperty.call(createDraftOptions, 'key')) {
+  if (key || 'key' in createDraftOptions) {
     proxyDraft.key = key;
   }
   const { proxy, revoke } = Proxy.revocable<any>(
-    Array.isArray(original) ? Object.assign([], proxyDraft) : proxyDraft,
+    type === DraftType.Array ? Object.assign([], proxyDraft) : proxyDraft,
     proxyHandler
   );
-  finalities.revoke.unshift(revoke);
+  finalities.revoke.push(revoke);
   proxyDraft.proxy = proxy;
   if (parentDraft) {
     const target = parentDraft;
-    const oldProxyDraft = getProxyDraft(proxy)!;
-    target.finalities.draft.unshift((patches, inversePatches) => {
+    target.finalities.draft.push((patches, inversePatches) => {
+      const oldProxyDraft = getProxyDraft(proxy)!;
       // if target is a Set draft, `setMap` is the real Set copies proxy mapping.
       const proxyDraft = getProxyDraft(
         get(target.type === DraftType.Set ? target.setMap : target.copy, key!)
@@ -257,7 +259,7 @@ export function createDraft<T extends object>(createDraftOptions: {
   } else {
     // !case: handle the root draft
     const target = getProxyDraft(proxy)!;
-    target.finalities.draft.unshift((patches, inversePatches) => {
+    target.finalities.draft.push((patches, inversePatches) => {
       finalizePatches(target, patches, inversePatches);
     });
   }
@@ -268,21 +270,38 @@ internal.createDraft = createDraft;
 
 export function finalizeDraft<T>(
   result: T,
+  returnedValue: [T] | [],
   patches?: Patches,
-  inversePatches?: Patches
+  inversePatches?: Patches,
+  enableAutoFreeze?: boolean
 ) {
-  const proxyDraft = getProxyDraft(result as any)!;
-  for (const finalize of proxyDraft.finalities.draft) {
-    finalize(patches, inversePatches);
+  const proxyDraft = getProxyDraft(result);
+  const original = proxyDraft?.original ?? result;
+  const hasReturnedValue = !!returnedValue.length;
+  if (proxyDraft?.operated) {
+    while (proxyDraft.finalities.draft.length > 0) {
+      const finalize = proxyDraft.finalities.draft.pop()!;
+      finalize(patches, inversePatches);
+    }
   }
-  const state = proxyDraft.operated ? proxyDraft.copy : proxyDraft.original;
-  revokeProxy(proxyDraft);
-  if (proxyDraft.options.enableAutoFreeze) {
+  const state = hasReturnedValue
+    ? returnedValue[0]
+    : proxyDraft
+    ? proxyDraft.operated
+      ? proxyDraft.copy
+      : proxyDraft.original
+    : result;
+  if (proxyDraft) revokeProxy(proxyDraft);
+  if (enableAutoFreeze) {
     deepFreeze(state);
   }
-  return [state, patches, inversePatches] as [
-    T,
-    Patches | undefined,
-    Patches | undefined
-  ];
+  return [
+    state,
+    patches && hasReturnedValue
+      ? [{ op: Operation.Replace, path: [], value: returnedValue[0] }]
+      : patches,
+    inversePatches && hasReturnedValue
+      ? [{ op: Operation.Replace, path: [], value: original }]
+      : inversePatches,
+  ] as [T, Patches | undefined, Patches | undefined];
 }