perf(array): improve array method perf

package.json

@@ -31,6 +31,7 @@
     "benchmark": "yarn build && yarn benchmark:base && yarn benchmark:object && yarn benchmark:array && yarn benchmark:class",
     "all-benchmark": "yarn build && NODE_ENV='production' ts-node test/benchmark/index.ts",
     "benchmark:reducer": "NODE_ENV='production' ts-node test/performance/benchmark-reducer.ts",
+    "benchmark:reducer1": "NODE_ENV='production' node test/performance/benchmark-reducer1.mjs",
     "benchmark:base": "NODE_ENV='production' ts-node test/performance/benchmark.ts",
     "benchmark:object": "NODE_ENV='production' ts-node test/performance/benchmark-object.ts",
     "benchmark:array": "NODE_ENV='production' ts-node test/performance/benchmark-array.ts",

src/current.ts

@@ -1,4 +1,4 @@
-import { type Draft, DraftType, type ProxyDraft } from './interface';
+import { DraftType, type Draft, type ProxyDraft } from './interface';
 import {
   forEach,
   get,

src/draft.ts

@@ -3,8 +3,8 @@ import {
   Finalities,
   Patches,
   ProxyDraft,
-  Options,
   Operation,
+  DraftOptions,
 } from './interface';
 import { dataTypes, PROXY_DRAFT } from './constant';
 import { mapHandler, mapHandlerKeys } from './map';
@@ -29,17 +29,22 @@ import {
   finalizeSetValue,
   markFinalization,
   finalizePatches,
+  isDraft,
 } from './utils';
 import { checkReadable } from './unsafe';
 import { generatePatches } from './patch';
 
-const draftsCache = new WeakSet<object>();
+// The array methods that need to be handled by the draft.
+// `sort` is not included, because array items may be modified by mutations in the sort function, it has to be drafted.
+// `copyWithin` is not included, it would require implementing a complete check of array copy reference relationships,
+// which might result in limited performance gains and increased maintenance complexity.
+const proxyArrayMethods = ['splice', 'shift', 'unshift', 'reverse'];
 
 const proxyHandler: ProxyHandler<ProxyDraft> = {
   get(target: ProxyDraft, key: string | number | symbol, receiver: any) {
     const copy = target.copy?.[key];
     // Improve draft reading performance by caching the draft copy.
-    if (copy && draftsCache.has(copy)) {
+    if (copy && target.finalities.draftsCache.has(copy)) {
       return copy;
     }
     if (key === PROXY_DRAFT) return target;
@@ -61,6 +66,15 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
       }
     }
     const source = latest(target);
+    const skipFinalization = target.options.skipFinalization;
+
+    if (
+      skipFinalization &&
+      source[key] &&
+      target.finalities.draftsCache.has(source[key])
+    ) {
+      return source[key];
+    }
 
     if (source instanceof Map && mapHandlerKeys.includes(key as any)) {
       if (key === 'size') {
@@ -84,9 +98,35 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
 
     if (!has(source, key)) {
       const desc = getDescriptor(source, key);
+      const value = desc?.value;
+      if (
+        target.type === DraftType.Array &&
+        skipFinalization &&
+        proxyArrayMethods.includes(key as string)
+      ) {
+        return function (this: any, ...args: any[]) {
+          let returnValue: any;
+          target.finalities.arrayHandling = true;
+          try {
+            returnValue = value.apply(this, args);
+            if (isDraftable(returnValue, target.options)) {
+              returnValue = createDraft({
+                original: returnValue,
+                parentDraft: undefined,
+                key: undefined,
+                finalities: target.finalities,
+                options: target.options,
+              });
+            }
+            return returnValue;
+          } finally {
+            target.finalities.arrayHandling = false;
+          }
+        };
+      }
       return desc
         ? `value` in desc
-          ? desc.value
+          ? value
           : // !case: support for getter
             desc.get?.call(target.proxy)
         : undefined;
@@ -99,10 +139,17 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
       return value;
     }
     // Ensure that the assigned values are not drafted
-    if (value === peek(target.original, key)) {
+    if (
+      !target.finalities.arrayHandling &&
+      (value === peek(target.original, key) || skipFinalization?.has(value))
+    ) {
+      const shouldSkip = skipFinalization?.has(value);
+      if (shouldSkip) {
+        skipFinalization!.delete(value);
+      }
       ensureShallowCopy(target);
       target.copy![key] = createDraft({
-        original: target.original[key],
+        original: shouldSkip ? target.copy![key] : target.original[key],
         parentDraft: target,
         key: target.type === DraftType.Array ? Number(key) : key,
         finalities: target.finalities,
@@ -118,6 +165,15 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
       }
       return target.copy![key];
     }
+    if (
+      target.finalities.arrayHandling &&
+      skipFinalization &&
+      !isDraft(value) &&
+      isDraftable(value)
+    ) {
+      // !case: handle the case of assigning the original array item via array methods(`splice`, `shift``, `unshift`, `reverse`)
+      skipFinalization.add(value);
+    }
     return value;
   },
   set(target: ProxyDraft, key: string | number | symbol, value: any) {
@@ -223,7 +279,7 @@ export function createDraft<T extends object>(createDraftOptions: {
   parentDraft?: ProxyDraft | null;
   key?: string | number | symbol;
   finalities: Finalities;
-  options: Options<any, any>;
+  options: DraftOptions;
 }): T {
   const { original, parentDraft, key, finalities, options } =
     createDraftOptions;
@@ -252,7 +308,7 @@ export function createDraft<T extends object>(createDraftOptions: {
     proxyHandler
   );
   finalities.revoke.push(revoke);
-  draftsCache.add(proxy);
+  finalities.draftsCache.add(proxy);
   proxyDraft.proxy = proxy;
   if (parentDraft) {
     const target = parentDraft;
@@ -270,7 +326,11 @@ export function createDraft<T extends object>(createDraftOptions: {
         }
         finalizeSetValue(proxyDraft);
         finalizePatches(proxyDraft, generatePatches, patches, inversePatches);
-        if (__DEV__ && target.options.enableAutoFreeze) {
+        if (
+          __DEV__ &&
+          target.options.enableAutoFreeze &&
+          typeof updatedValue === 'object'
+        ) {
           target.options.updatedValues =
             target.options.updatedValues ?? new WeakMap();
           target.options.updatedValues.set(updatedValue, proxyDraft.original);
@@ -315,10 +375,10 @@ export function finalizeDraft<T>(
   const state = hasReturnedValue
     ? returnedValue[0]
     : proxyDraft
-    ? proxyDraft.operated
-      ? proxyDraft.copy
-      : proxyDraft.original
-    : result;
+      ? proxyDraft.operated
+        ? proxyDraft.copy
+        : proxyDraft.original
+      : result;
   if (proxyDraft) revokeProxy(proxyDraft);
   if (enableAutoFreeze) {
     deepFreeze(state, state, proxyDraft?.options.updatedValues);

src/draftify.ts

@@ -1,6 +1,6 @@
-import {
+import type {
+  DraftOptions,
   Finalities,
-  Options,
   Patches,
   PatchesOptions,
   Result,
@@ -12,15 +12,17 @@ import { dataTypes } from './constant';
 export function draftify<
   T extends object,
   O extends PatchesOptions = false,
-  F extends boolean = false
+  F extends boolean = false,
 >(
   baseState: T,
-  options: Options<O, F>
+  options: DraftOptions
 ): [T, (returnedValue: [T] | []) => Result<T, O, F>] {
   const finalities: Finalities = {
     draft: [],
     revoke: [],
     handledSet: new WeakSet<any>(),
+    draftsCache: new WeakSet<object>(),
+    arrayHandling: false,
   };
   let patches: Patches | undefined;
   let inversePatches: Patches | undefined;

src/interface.ts

@@ -32,6 +32,8 @@ export interface Finalities {
   draft: ((patches?: Patches, inversePatches?: Patches) => void)[];
   revoke: (() => void)[];
   handledSet: WeakSet<any>;
+  draftsCache: WeakSet<object>;
+  arrayHandling: boolean;
 }
 
 export interface ProxyDraft<T = any> {
@@ -42,7 +44,7 @@ export interface ProxyDraft<T = any> {
   copy: T | null;
   proxy: T | null;
   finalities: Finalities;
-  options: Options<any, any> & { updatedValues?: WeakMap<any, any> };
+  options: DraftOptions;
   parent?: ProxyDraft | null;
   key?: string | number | symbol;
   setMap?: Map<any, ProxyDraft>;
@@ -62,30 +64,30 @@ export type Patch<P extends PatchesOptions = any> = P extends {
       path: string;
     }
   : P extends true | object
-  ? IPatch & {
-      path: (string | number)[];
-    }
-  : IPatch & {
-      path: string | (string | number)[];
-    };
+    ? IPatch & {
+        path: (string | number)[];
+      }
+    : IPatch & {
+        path: string | (string | number)[];
+      };
 
 export type Patches<P extends PatchesOptions = any> = Patch<P>[];
 
 export type Result<
   T extends any,
   O extends PatchesOptions,
-  F extends boolean
+  F extends boolean,
 > = O extends true | object
   ? [F extends true ? Immutable<T> : T, Patches<O>, Patches<O>]
   : F extends true
-  ? Immutable<T>
-  : T;
+    ? Immutable<T>
+    : T;
 
 export type CreateResult<
   T extends any,
   O extends PatchesOptions,
   F extends boolean,
-  R extends void | Promise<void> | T | Promise<T>
+  R extends void | Promise<void> | T | Promise<T>,
 > = R extends Promise<void> | Promise<T>
   ? Promise<Result<T, O, F>>
   : Result<T, O, F>;
@@ -99,8 +101,8 @@ export type Mark<O extends PatchesOptions, F extends boolean> = (
 ) => O extends true | object
   ? BaseMark
   : F extends true
-  ? BaseMark
-  : MarkWithCopy;
+    ? BaseMark
+    : MarkWithCopy;
 
 export interface ApplyMutableOptions {
   /**
@@ -129,6 +131,17 @@ export interface Options<O extends PatchesOptions, F extends boolean> {
   mark?: Mark<O, F>;
 }
 
+export type DraftOptions = Options<any, any> & {
+  /**
+   * a collection for circular reference check
+   */
+  updatedValues?: WeakMap<any, any>;
+  /**
+   * a collection for array item skip deep check
+   */
+  skipFinalization?: WeakSet<any>;
+};
+
 export interface ExternalOptions<O extends PatchesOptions, F extends boolean> {
   /**
    * In strict mode, Forbid accessing non-draftable values and forbid returning a non-draft value.
@@ -147,6 +160,10 @@ export interface ExternalOptions<O extends PatchesOptions, F extends boolean> {
    * And it can also return a shallow copy function(AutoFreeze and Patches should both be disabled).
    */
   mark?: Mark<O, F>[] | Mark<O, F>;
+  /**
+   * Enable optimized array for improving performance.
+   */
+  enableOptimizedArray?: boolean;
 }
 
 // Exclude `symbol`
@@ -161,8 +178,8 @@ export type IfAvailable<T, Fallback = void> = true | false extends (
 )
   ? Fallback
   : keyof T extends never
-  ? Fallback
-  : T;
+    ? Fallback
+    : T;
 type WeakReferences =
   | IfAvailable<WeakMap<any, any>>
   | IfAvailable<WeakSet<any>>;
@@ -171,14 +188,14 @@ type AtomicObject = Function | Promise<any> | Date | RegExp;
 export type Immutable<T> = T extends Primitive | AtomicObject
   ? T
   : T extends IfAvailable<ReadonlyMap<infer K, infer V>>
-  ? ImmutableMap<K, V>
-  : T extends IfAvailable<ReadonlySet<infer V>>
-  ? ImmutableSet<V>
-  : T extends WeakReferences
-  ? T
-  : T extends object
-  ? ImmutableObject<T>
-  : T;
+    ? ImmutableMap<K, V>
+    : T extends IfAvailable<ReadonlySet<infer V>>
+      ? ImmutableSet<V>
+      : T extends WeakReferences
+        ? T
+        : T extends object
+          ? ImmutableObject<T>
+          : T;
 
 type DraftedMap<K, V> = Map<K, Draft<V>>;
 type DraftedSet<T> = Set<Draft<T>>;