From e9546369ddd0e22607e1627d0dfaf4ca8fc204ba Mon Sep 17 00:00:00 2001
From: Oliver Hoff <oliver@hofff.com>
Date: Tue, 2 Apr 2024 07:19:53 +0200
Subject: [PATCH] fix(vanilla): setters do not freeze new values

---
 src/vanilla/utils/freezeAtom.ts               | 62 ++++++++++++-------
 tests/react/vanilla-utils/freezeAtom.test.tsx | 44 ++++++++++---
 2 files changed, 74 insertions(+), 32 deletions(-)

diff --git a/src/vanilla/utils/freezeAtom.ts b/src/vanilla/utils/freezeAtom.ts
index e779c2d4aa..585274bf17 100644
--- a/src/vanilla/utils/freezeAtom.ts
+++ b/src/vanilla/utils/freezeAtom.ts
@@ -1,9 +1,6 @@
-import { atom } from '../../vanilla.ts'
-import type { Atom } from '../../vanilla.ts'
+import type { Atom, WritableAtom } from '../../vanilla.ts'
 
-const cache1 = new WeakMap()
-const memo1 = <T>(create: () => T, dep1: object): T =>
-  (cache1.has(dep1) ? cache1 : cache1.set(dep1, create())).get(dep1)
+const frozenAtoms = new WeakSet<Atom<any>>()
 
 const deepFreeze = (obj: unknown) => {
   if (typeof obj !== 'object' || obj === null) return
@@ -18,25 +15,44 @@ const deepFreeze = (obj: unknown) => {
 
 export function freezeAtom<AtomType extends Atom<unknown>>(
   anAtom: AtomType,
-): AtomType {
-  return memo1(() => {
-    const frozenAtom = atom(
-      (get) => deepFreeze(get(anAtom)),
-      (_get, set, arg) => set(anAtom as never, arg),
-    )
-    return frozenAtom as never
-  }, anAtom)
+): AtomType
+export function freezeAtom(
+  anAtom: WritableAtom<unknown, unknown[], unknown>,
+): WritableAtom<unknown, unknown[], unknown> {
+  if (frozenAtoms.has(anAtom)) {
+    return anAtom
+  }
+  frozenAtoms.add(anAtom)
+
+  const origRead = anAtom.read
+  anAtom.read = function (get, options) {
+    return deepFreeze(origRead.call(this, get, options))
+  }
+  if ('write' in anAtom) {
+    const origWrite = anAtom.write
+    anAtom.write = function (get, set, ...args) {
+      return origWrite.call(
+        this,
+        get,
+        (...setArgs) => {
+          if (setArgs[0] === anAtom) {
+            setArgs[1] = deepFreeze(setArgs[1])
+          }
+
+          return set(...setArgs)
+        },
+        ...args,
+      )
+    }
+  }
+  return anAtom
 }
 
 export function freezeAtomCreator<
-  CreateAtom extends (...params: never[]) => Atom<unknown>,
->(createAtom: CreateAtom): CreateAtom {
-  return ((...params: never[]) => {
-    const anAtom = createAtom(...params)
-    const origRead = anAtom.read
-    anAtom.read = function (get, options) {
-      return deepFreeze(origRead.call(this, get, options))
-    }
-    return anAtom
-  }) as never
+  CreateAtom extends (...args: unknown[]) => Atom<unknown>,
+>(createAtom: CreateAtom): CreateAtom
+export function freezeAtomCreator(
+  createAtom: (...args: unknown[]) => WritableAtom<unknown, unknown[], unknown>,
+): (...args: unknown[]) => WritableAtom<unknown, unknown[], unknown> {
+  return (...args) => freezeAtom(createAtom(...args))
 }
diff --git a/tests/react/vanilla-utils/freezeAtom.test.tsx b/tests/react/vanilla-utils/freezeAtom.test.tsx
index dfe4e7e0cb..52541827c2 100644
--- a/tests/react/vanilla-utils/freezeAtom.test.tsx
+++ b/tests/react/vanilla-utils/freezeAtom.test.tsx
@@ -1,43 +1,69 @@
 import { StrictMode } from 'react'
-import { render } from '@testing-library/react'
+import { fireEvent, render } from '@testing-library/react'
 import { it } from 'vitest'
 import { useAtom } from 'jotai/react'
 import { atom } from 'jotai/vanilla'
 import { freezeAtom, freezeAtomCreator } from 'jotai/vanilla/utils'
 
 it('freezeAtom basic test', async () => {
-  const objAtom = atom({ count: 0 })
+  const objAtom = atom({ deep: {} }, (_get, set, _ignored?) => {
+    set(objAtom, { deep: {} })
+  })
 
   const Component = () => {
-    const [obj] = useAtom(freezeAtom(objAtom))
+    const [obj, setObj] = useAtom(freezeAtom(objAtom))
 
-    return <div>isFrozen: {`${Object.isFrozen(obj)}`}</div>
+    return (
+      <>
+        <button onClick={setObj}>change</button>
+        <div>
+          isFrozen: {`${Object.isFrozen(obj) && Object.isFrozen(obj.deep)}`}
+        </div>
+      </>
+    )
   }
 
-  const { findByText } = render(
+  const { getByText, findByText } = render(
     <StrictMode>
       <Component />
     </StrictMode>,
   )
 
   await findByText('isFrozen: true')
+
+  fireEvent.click(getByText('change'))
+
+  await findByText('isFrozen: true')
 })
 
 it('freezeAtomCreator basic test', async () => {
   const createFrozenAtom = freezeAtomCreator(atom)
-  const objAtom = createFrozenAtom({ count: 0 })
+  const objAtom = createFrozenAtom({ deep: {} }, (_get, set, _ignored?) => {
+    set(objAtom, { deep: {} })
+  })
 
   const Component = () => {
-    const [obj] = useAtom(objAtom)
+    const [obj, setObj] = useAtom(objAtom)
 
-    return <div>isFrozen: {`${Object.isFrozen(obj)}`}</div>
+    return (
+      <>
+        <button onClick={setObj}>change</button>
+        <div>
+          isFrozen: {`${Object.isFrozen(obj) && Object.isFrozen(obj.deep)}`}
+        </div>
+      </>
+    )
   }
 
-  const { findByText } = render(
+  const { getByText, findByText } = render(
     <StrictMode>
       <Component />
     </StrictMode>,
   )
 
   await findByText('isFrozen: true')
+
+  fireEvent.click(getByText('change'))
+
+  await findByText('isFrozen: true')
 })