perf(cache): improve drafts cache for perf

src/draft.ts

@@ -29,17 +29,16 @@ import {
   finalizeSetValue,
   markFinalization,
   finalizePatches,
+  isDraft,
 } from './utils';
 import { checkReadable } from './unsafe';
 import { generatePatches } from './patch';
 
-const draftsCache = new WeakSet<object>();
-
 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;
@@ -118,6 +117,9 @@ const proxyHandler: ProxyHandler<ProxyDraft> = {
       }
       return target.copy![key];
     }
+    if (isDraft(value)) {
+      target.finalities.draftsCache.add(value);
+    }
     return value;
   },
   set(target: ProxyDraft, key: string | number | symbol, value: any) {
@@ -252,7 +254,6 @@ export function createDraft<T extends object>(createDraftOptions: {
     proxyHandler
   );
   finalities.revoke.push(revoke);
-  draftsCache.add(proxy);
   proxyDraft.proxy = proxy;
   if (parentDraft) {
     const target = parentDraft;

src/draftify.ts

@@ -21,6 +21,7 @@ export function draftify<
     draft: [],
     revoke: [],
     handledSet: new WeakSet<any>(),
+    draftsCache: new WeakSet<object>(),
   };
   let patches: Patches | undefined;
   let inversePatches: Patches | undefined;

src/interface.ts

@@ -32,6 +32,7 @@ export interface Finalities {
   draft: ((patches?: Patches, inversePatches?: Patches) => void)[];
   revoke: (() => void)[];
   handledSet: WeakSet<any>;
+  draftsCache: WeakSet<object>;
 }
 
 export interface ProxyDraft<T = any> {