From 9cc760ba1b38aa9768b1d937e3e61fd154d35b42 Mon Sep 17 00:00:00 2001
From: Michael Lin <unadlib@gmail.com>
Date: Sat, 27 Jul 2024 04:51:59 +0800
Subject: [PATCH] revert(current): revert current() handling (#52)

---
 src/current.ts     | 32 +++++++++++++++++++++---------
 src/utils/copy.ts  |  4 ++--
 test/index.test.ts | 49 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 11 deletions(-)

diff --git a/src/current.ts b/src/current.ts
index f6d464e6..83610376 100644
--- a/src/current.ts
+++ b/src/current.ts
@@ -64,23 +64,37 @@ function getCurrent(target: any) {
   if (!isDraftable(target, proxyDraft?.options)) return target;
   const type = getType(target);
   if (proxyDraft && !proxyDraft.operated) return proxyDraft.original;
-  if (proxyDraft) proxyDraft.finalized = true;
-  let currentValue = target;
+  let currentValue: any;
+  function ensureShallowCopy() {
+    currentValue =
+      type === DraftType.Map
+        ? new Map(target)
+        : type === DraftType.Set
+        ? Array.from(proxyDraft!.setMap!.values()!)
+        : shallowCopy(target, proxyDraft?.options);
+  }
+
   if (proxyDraft) {
+    // It's a proxy draft, let's create a shallow copy eagerly
+    proxyDraft.finalized = true;
     try {
-      currentValue =
-        type === DraftType.Map
-          ? new Map(target)
-          : type === DraftType.Set
-          ? Array.from(proxyDraft.setMap!.values()!)
-          : shallowCopy(target, proxyDraft.options);
+      ensureShallowCopy();
     } finally {
       proxyDraft.finalized = false;
     }
+  } else {
+    // It's not a proxy draft, let's use the target directly and let's see
+    // lazily if we need to create a shallow copy
+    currentValue = target;
   }
+
   forEach(currentValue, (key, value) => {
     if (proxyDraft && isEqual(get(proxyDraft.original, key), value)) return;
-    set(currentValue, key, getCurrent(value));
+    const newValue = getCurrent(value);
+    if (newValue !== value) {
+      if (currentValue === target) ensureShallowCopy();
+      set(currentValue, key, newValue);
+    }
   });
   return type === DraftType.Set ? new Set(currentValue) : currentValue;
 }
diff --git a/src/utils/copy.ts b/src/utils/copy.ts
index 1af6d89f..f53aa257 100644
--- a/src/utils/copy.ts
+++ b/src/utils/copy.ts
@@ -29,7 +29,7 @@ function strictCopy(target: any) {
 
 const propIsEnum = Object.prototype.propertyIsEnumerable;
 
-export function shallowCopy(original: any, options: Options<any, any>) {
+export function shallowCopy(original: any, options?: Options<any, any>) {
   let markResult: any;
   if (Array.isArray(original)) {
     return Array.prototype.concat.call(original);
@@ -38,7 +38,7 @@ export function shallowCopy(original: any, options: Options<any, any>) {
   } else if (original instanceof Map) {
     return new Map(original);
   } else if (
-    options.mark &&
+    options?.mark &&
     ((markResult = options.mark(original, dataTypes)),
     markResult !== undefined) &&
     markResult !== dataTypes.mutable
diff --git a/test/index.test.ts b/test/index.test.ts
index f90c4124..3b457829 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -3928,3 +3928,52 @@ test('#37 - Proxy revoked error with `mark` option when performing chai deep equ
     })
   );
 });
+
+it("#51 doesn't assign values to frozen object", () => {
+  const frozen = Object.freeze({ test: 42 });
+  const base = { k: null };
+  create(base, (draft) => {
+    // @ts-ignore
+    draft.k = frozen;
+    const c = current(draft); // Boom! tries to set a value in the frozen object
+    expect(c.k).toBe(frozen);
+  });
+});
+
+it('#51-1 nested drafts work after current', () => {
+  const base = { k1: {}, k2: {} };
+  const result = create(base, (draft) => {
+    const obj = { x: draft.k2 };
+    draft.k1 = obj;
+    current(draft); // try to comment me
+    // @ts-ignore
+    obj.x.abc = 100;
+    // @ts-ignore
+    draft.k2.def = 200;
+  });
+  // @ts-ignore
+  expect(result.k1.x).toBe(result.k2);
+  expect(result).toEqual({
+    k1: { x: { abc: 100, def: 200 } },
+    k2: { abc: 100, def: 200 },
+  });
+});
+
+it('#51-2 nested drafts work after current', () => {
+  const base = { k1: {}, k2: {} };
+  const result = create(base, (draft) => {
+    const obj = { x: draft.k2 };
+    draft.k1 = obj;
+    // current(draft); // try to comment me
+    // @ts-ignore
+    obj.x.abc = 100;
+    // @ts-ignore
+    draft.k2.def = 200;
+  });
+  // @ts-ignore
+  expect(result.k1.x).toBe(result.k2);
+  expect(result).toEqual({
+    k1: { x: { abc: 100, def: 200 } },
+    k2: { abc: 100, def: 200 },
+  });
+});