diff --git a/src/draft.ts b/src/draft.ts index 2c99c6b2..8fdeb933 100644 --- a/src/draft.ts +++ b/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(); - const proxyHandler: ProxyHandler = { 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 = { } 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(createDraftOptions: { proxyHandler ); finalities.revoke.push(revoke); - draftsCache.add(proxy); proxyDraft.proxy = proxy; if (parentDraft) { const target = parentDraft; diff --git a/src/draftify.ts b/src/draftify.ts index c0c8502f..d4dedfd2 100644 --- a/src/draftify.ts +++ b/src/draftify.ts @@ -21,6 +21,7 @@ export function draftify< draft: [], revoke: [], handledSet: new WeakSet(), + draftsCache: new WeakSet(), }; let patches: Patches | undefined; let inversePatches: Patches | undefined; diff --git a/src/interface.ts b/src/interface.ts index d38d969a..cccef186 100644 --- a/src/interface.ts +++ b/src/interface.ts @@ -32,6 +32,7 @@ export interface Finalities { draft: ((patches?: Patches, inversePatches?: Patches) => void)[]; revoke: (() => void)[]; handledSet: WeakSet; + draftsCache: WeakSet; } export interface ProxyDraft {