fix(vanilla): setters do not freeze new values (#2476)

Co-authored-by: Daishi Kato <[email protected]>

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))
 }

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')
 })