From b1f194eedf191f4c8e8bce4474d2a0117b55471d Mon Sep 17 00:00:00 2001
From: "Mohammad H. Sattarian" <mh_sattarian@outlook.com>
Date: Sat, 4 May 2024 00:25:55 +0330
Subject: [PATCH 1/4] feat: allow providing a subscribe method to
 createJSONStorage util

---
 src/vanilla/utils/atomWithStorage.ts          | 81 +++++++++++----
 .../vanilla-utils/atomWithStorage.test.tsx    | 99 +++++++++++++++++++
 2 files changed, 159 insertions(+), 21 deletions(-)

diff --git a/src/vanilla/utils/atomWithStorage.ts b/src/vanilla/utils/atomWithStorage.ts
index 038c76feaf..bc939e08ce 100644
--- a/src/vanilla/utils/atomWithStorage.ts
+++ b/src/vanilla/utils/atomWithStorage.ts
@@ -5,8 +5,21 @@ import { RESET } from './constants.ts'
 const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
   typeof (x as any)?.then === 'function'
 
+type Subscribe<Value> = (
+  key: string,
+  callback: (value: Value) => void,
+  initialValue: Value,
+) => Unsubscribe
+
 type Unsubscribe = () => void
 
+type SubscribeHandler<Value> = (
+  subscribe: Subscribe<Value>,
+  key: string,
+  callback: (value: Value) => void,
+  initialValue: Value,
+) => Unsubscribe
+
 type SetStateActionWithReset<Value> =
   | Value
   | typeof RESET
@@ -38,12 +51,14 @@ export interface AsyncStringStorage {
   getItem: (key: string) => PromiseLike<string | null>
   setItem: (key: string, newValue: string) => PromiseLike<void>
   removeItem: (key: string) => PromiseLike<void>
+  subscribe?: Subscribe<string | null>
 }
 
 export interface SyncStringStorage {
   getItem: (key: string) => string | null
   setItem: (key: string, newValue: string) => void
   removeItem: (key: string) => void
+  subscribe?: Subscribe<string | null>
 }
 
 export function withStorageValidator<Value>(
@@ -114,6 +129,42 @@ export function createJSONStorage<Value>(
 ): AsyncStorage<Value> | SyncStorage<Value> {
   let lastStr: string | undefined
   let lastValue: Value
+
+  const webStorageSubscribe: Subscribe<Value> = (key, callback) => {
+    if (!(getStringStorage() instanceof window.Storage)) {
+      return () => {}
+    }
+    const storageEventCallback = (e: StorageEvent) => {
+      if (e.storageArea === getStringStorage() && e.key === key) {
+        callback((e.newValue || '') as Value)
+      }
+    }
+    window.addEventListener('storage', storageEventCallback)
+    return () => {
+      window.removeEventListener('storage', storageEventCallback)
+    }
+  }
+
+  const handleSubscribe: SubscribeHandler<Value> = (
+    subscriber,
+    key,
+    callback,
+    initialValue,
+  ) => {
+    function callbackWithParser(v: Value) {
+      let newValue: Value
+      try {
+        newValue = JSON.parse((v as string) || '')
+      } catch {
+        newValue = initialValue
+      }
+
+      callback(newValue as Value)
+    }
+
+    return subscriber(key, callbackWithParser, initialValue)
+  }
+
   const storage: AsyncStorage<Value> | SyncStorage<Value> = {
     getItem: (key, initialValue) => {
       const parse = (str: string | null) => {
@@ -141,30 +192,18 @@ export function createJSONStorage<Value>(
       ),
     removeItem: (key) => getStringStorage()?.removeItem(key),
   }
+
   if (
     typeof window !== 'undefined' &&
-    typeof window.addEventListener === 'function' &&
-    window.Storage
+    typeof window.addEventListener === 'function'
   ) {
-    storage.subscribe = (key, callback, initialValue) => {
-      if (!(getStringStorage() instanceof window.Storage)) {
-        return () => {}
-      }
-      const storageEventCallback = (e: StorageEvent) => {
-        if (e.storageArea === getStringStorage() && e.key === key) {
-          let newValue: Value
-          try {
-            newValue = JSON.parse(e.newValue || '')
-          } catch {
-            newValue = initialValue
-          }
-          callback(newValue)
-        }
-      }
-      window.addEventListener('storage', storageEventCallback)
-      return () => {
-        window.removeEventListener('storage', storageEventCallback)
-      }
+    if (getStringStorage()?.subscribe) {
+      storage.subscribe = handleSubscribe.bind(
+        null,
+        getStringStorage()!.subscribe as unknown as Subscribe<Value>,
+      )
+    } else if (window.Storage) {
+      storage.subscribe = handleSubscribe.bind(null, webStorageSubscribe)
     }
   }
   return storage
diff --git a/tests/react/vanilla-utils/atomWithStorage.test.tsx b/tests/react/vanilla-utils/atomWithStorage.test.tsx
index 1056f7ba29..4eb2194dc2 100644
--- a/tests/react/vanilla-utils/atomWithStorage.test.tsx
+++ b/tests/react/vanilla-utils/atomWithStorage.test.tsx
@@ -9,6 +9,7 @@ import {
   createJSONStorage,
   unstable_withStorageValidator as withStorageValidator,
 } from 'jotai/vanilla/utils'
+import type { SyncStringStorage } from 'jotai/vanilla/utils/atomWithStorage'
 
 const resolve: (() => void)[] = []
 
@@ -601,3 +602,101 @@ describe('withStorageValidator', () => {
     atomWithStorage('my-number', 0, withStorageValidator(isNumber)(storage))
   })
 })
+
+describe('with subscribe method in string storage', () => {
+  it('createJSONStorage subscriber is called correctly', async () => {
+    const store = createStore()
+
+    const subscribe = vi.fn()
+    const stringStorage = {
+      getItem: () => {
+        return null
+      },
+      setItem: () => {},
+      removeItem: () => {},
+      subscribe,
+    }
+
+    const dummyStorage = createJSONStorage<number>(() => stringStorage)
+
+    const dummyAtom = atomWithStorage<number>('dummy', 1, dummyStorage)
+
+    const DummyComponent = () => {
+      const [value] = useAtom(dummyAtom, { store })
+      return (
+        <>
+          <div>{value}</div>
+        </>
+      )
+    }
+
+    render(
+      <StrictMode>
+        <DummyComponent />
+      </StrictMode>,
+    )
+
+    expect(subscribe).toHaveBeenCalledWith('dummy', expect.any(Function), 1)
+  })
+
+  it('createJSONStorage subscriber responds to events correctly', async () => {
+    const storageData: Record<string, string> = {
+      count: '10',
+    }
+
+    const stringStorage = {
+      getItem: (key: string) => {
+        return storageData[key] || null
+      },
+      setItem: (key: string, newValue: string) => {
+        storageData[key] = newValue
+      },
+      removeItem: (key: string) => {
+        delete storageData[key]
+      },
+      subscribe(key, callback) {
+        function handler(event: CustomEvent<string>) {
+          callback(event.detail)
+        }
+
+        window.addEventListener('dummystoragechange', handler as EventListener)
+        return () =>
+          window.removeEventListener(
+            'dummystoragechange',
+            handler as EventListener,
+          )
+      },
+    } as SyncStringStorage
+
+    const dummyStorage = createJSONStorage<number>(() => stringStorage)
+
+    const countAtom = atomWithStorage('count', 1, dummyStorage)
+
+    const Counter = () => {
+      const [count] = useAtom(countAtom)
+      return (
+        <>
+          <div>count: {count}</div>
+        </>
+      )
+    }
+
+    const { findByText } = render(
+      <StrictMode>
+        <Counter />
+      </StrictMode>,
+    )
+
+    await findByText('count: 10')
+
+    storageData.count = '12'
+    fireEvent(
+      window,
+      new CustomEvent('dummystoragechange', {
+        detail: '12',
+      }),
+    )
+    await findByText('count: 12')
+    // expect(storageData.count).toBe('11')
+  })
+})

From 82af4a61573a9a709a4aeb2e015ba06cada01af3 Mon Sep 17 00:00:00 2001
From: "Mohammad H. Sattarian" <mh_sattarian@outlook.com>
Date: Thu, 23 May 2024 00:34:12 +0330
Subject: [PATCH 2/4] refactor: convert handleSubscribe to a HOF

---
 src/vanilla/utils/atomWithStorage.ts | 44 +++++++++++-----------------
 1 file changed, 17 insertions(+), 27 deletions(-)

diff --git a/src/vanilla/utils/atomWithStorage.ts b/src/vanilla/utils/atomWithStorage.ts
index bc939e08ce..7656c9b0ce 100644
--- a/src/vanilla/utils/atomWithStorage.ts
+++ b/src/vanilla/utils/atomWithStorage.ts
@@ -1,5 +1,5 @@
-import { atom } from '../../vanilla.ts'
 import type { WritableAtom } from '../../vanilla.ts'
+import { atom } from '../../vanilla.ts'
 import { RESET } from './constants.ts'
 
 const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
@@ -13,13 +13,6 @@ type Subscribe<Value> = (
 
 type Unsubscribe = () => void
 
-type SubscribeHandler<Value> = (
-  subscribe: Subscribe<Value>,
-  key: string,
-  callback: (value: Value) => void,
-  initialValue: Value,
-) => Unsubscribe
-
 type SetStateActionWithReset<Value> =
   | Value
   | typeof RESET
@@ -145,26 +138,24 @@ export function createJSONStorage<Value>(
     }
   }
 
-  const handleSubscribe: SubscribeHandler<Value> = (
-    subscriber,
-    key,
-    callback,
-    initialValue,
-  ) => {
-    function callbackWithParser(v: Value) {
-      let newValue: Value
-      try {
-        newValue = JSON.parse((v as string) || '')
-      } catch {
-        newValue = initialValue
+  const createHandleSubscribe =
+    (subscriber: Subscribe<Value>) =>
+    (...params: Parameters<Subscribe<Value>>) => {
+      const [key, callback, initialValue] = params
+      function callbackWithParser(v: Value) {
+        let newValue: Value
+        try {
+          newValue = JSON.parse((v as string) || '')
+        } catch {
+          newValue = initialValue
+        }
+
+        callback(newValue as Value)
       }
 
-      callback(newValue as Value)
+      return subscriber(key, callbackWithParser, initialValue)
     }
 
-    return subscriber(key, callbackWithParser, initialValue)
-  }
-
   const storage: AsyncStorage<Value> | SyncStorage<Value> = {
     getItem: (key, initialValue) => {
       const parse = (str: string | null) => {
@@ -198,12 +189,11 @@ export function createJSONStorage<Value>(
     typeof window.addEventListener === 'function'
   ) {
     if (getStringStorage()?.subscribe) {
-      storage.subscribe = handleSubscribe.bind(
-        null,
+      storage.subscribe = createHandleSubscribe(
         getStringStorage()!.subscribe as unknown as Subscribe<Value>,
       )
     } else if (window.Storage) {
-      storage.subscribe = handleSubscribe.bind(null, webStorageSubscribe)
+      storage.subscribe = createHandleSubscribe(webStorageSubscribe)
     }
   }
   return storage

From 1044dfa2f0431ac3226d9b6297646576a58e7ed7 Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Fri, 24 May 2024 13:53:26 +0900
Subject: [PATCH 3/4] refactor

---
 src/vanilla/utils/atomWithStorage.ts          | 104 ++++++++----------
 .../vanilla-utils/atomWithStorage.test.tsx    |   2 +-
 2 files changed, 47 insertions(+), 59 deletions(-)

diff --git a/src/vanilla/utils/atomWithStorage.ts b/src/vanilla/utils/atomWithStorage.ts
index 7656c9b0ce..db0c6c99a4 100644
--- a/src/vanilla/utils/atomWithStorage.ts
+++ b/src/vanilla/utils/atomWithStorage.ts
@@ -1,17 +1,22 @@
-import type { WritableAtom } from '../../vanilla.ts'
 import { atom } from '../../vanilla.ts'
+import type { WritableAtom } from '../../vanilla.ts'
 import { RESET } from './constants.ts'
 
 const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
   typeof (x as any)?.then === 'function'
 
+type Unsubscribe = () => void
+
 type Subscribe<Value> = (
   key: string,
   callback: (value: Value) => void,
   initialValue: Value,
 ) => Unsubscribe
 
-type Unsubscribe = () => void
+type StringSubscribe = (
+  key: string,
+  callback: (value: string | null) => void,
+) => Unsubscribe
 
 type SetStateActionWithReset<Value> =
   | Value
@@ -22,36 +27,28 @@ export interface AsyncStorage<Value> {
   getItem: (key: string, initialValue: Value) => PromiseLike<Value>
   setItem: (key: string, newValue: Value) => PromiseLike<void>
   removeItem: (key: string) => PromiseLike<void>
-  subscribe?: (
-    key: string,
-    callback: (value: Value) => void,
-    initialValue: Value,
-  ) => Unsubscribe
+  subscribe?: Subscribe<Value>
 }
 
 export interface SyncStorage<Value> {
   getItem: (key: string, initialValue: Value) => Value
   setItem: (key: string, newValue: Value) => void
   removeItem: (key: string) => void
-  subscribe?: (
-    key: string,
-    callback: (value: Value) => void,
-    initialValue: Value,
-  ) => Unsubscribe
+  subscribe?: Subscribe<Value>
 }
 
 export interface AsyncStringStorage {
   getItem: (key: string) => PromiseLike<string | null>
   setItem: (key: string, newValue: string) => PromiseLike<void>
   removeItem: (key: string) => PromiseLike<void>
-  subscribe?: Subscribe<string | null>
+  subscribe?: StringSubscribe
 }
 
 export interface SyncStringStorage {
   getItem: (key: string) => string | null
   setItem: (key: string, newValue: string) => void
   removeItem: (key: string) => void
-  subscribe?: Subscribe<string | null>
+  subscribe?: StringSubscribe
 }
 
 export function withStorageValidator<Value>(
@@ -123,39 +120,6 @@ export function createJSONStorage<Value>(
   let lastStr: string | undefined
   let lastValue: Value
 
-  const webStorageSubscribe: Subscribe<Value> = (key, callback) => {
-    if (!(getStringStorage() instanceof window.Storage)) {
-      return () => {}
-    }
-    const storageEventCallback = (e: StorageEvent) => {
-      if (e.storageArea === getStringStorage() && e.key === key) {
-        callback((e.newValue || '') as Value)
-      }
-    }
-    window.addEventListener('storage', storageEventCallback)
-    return () => {
-      window.removeEventListener('storage', storageEventCallback)
-    }
-  }
-
-  const createHandleSubscribe =
-    (subscriber: Subscribe<Value>) =>
-    (...params: Parameters<Subscribe<Value>>) => {
-      const [key, callback, initialValue] = params
-      function callbackWithParser(v: Value) {
-        let newValue: Value
-        try {
-          newValue = JSON.parse((v as string) || '')
-        } catch {
-          newValue = initialValue
-        }
-
-        callback(newValue as Value)
-      }
-
-      return subscriber(key, callbackWithParser, initialValue)
-    }
-
   const storage: AsyncStorage<Value> | SyncStorage<Value> = {
     getItem: (key, initialValue) => {
       const parse = (str: string | null) => {
@@ -184,17 +148,41 @@ export function createJSONStorage<Value>(
     removeItem: (key) => getStringStorage()?.removeItem(key),
   }
 
-  if (
-    typeof window !== 'undefined' &&
-    typeof window.addEventListener === 'function'
-  ) {
-    if (getStringStorage()?.subscribe) {
-      storage.subscribe = createHandleSubscribe(
-        getStringStorage()!.subscribe as unknown as Subscribe<Value>,
-      )
-    } else if (window.Storage) {
-      storage.subscribe = createHandleSubscribe(webStorageSubscribe)
-    }
+  const createHandleSubscribe =
+    (subscriber: StringSubscribe): Subscribe<Value> =>
+    (key, callback, initialValue) =>
+      subscriber(key, (v) => {
+        let newValue: Value
+        try {
+          newValue = JSON.parse(v || '')
+        } catch {
+          newValue = initialValue
+        }
+        callback(newValue)
+      })
+
+  const subscriber: StringSubscribe | undefined =
+    getStringStorage()?.subscribe ||
+    (typeof window !== 'undefined' &&
+      typeof window.addEventListener === 'function' &&
+      ((key, callback) => {
+        if (!(getStringStorage() instanceof window.Storage)) {
+          return () => {}
+        }
+        const storageEventCallback = (e: StorageEvent) => {
+          if (e.storageArea === getStringStorage() && e.key === key) {
+            callback(e.newValue)
+          }
+        }
+        window.addEventListener('storage', storageEventCallback)
+        return () => {
+          window.removeEventListener('storage', storageEventCallback)
+        }
+      })) ||
+    undefined
+
+  if (subscriber) {
+    storage.subscribe = createHandleSubscribe(subscriber)
   }
   return storage
 }
diff --git a/tests/react/vanilla-utils/atomWithStorage.test.tsx b/tests/react/vanilla-utils/atomWithStorage.test.tsx
index 4eb2194dc2..2325e1e250 100644
--- a/tests/react/vanilla-utils/atomWithStorage.test.tsx
+++ b/tests/react/vanilla-utils/atomWithStorage.test.tsx
@@ -636,7 +636,7 @@ describe('with subscribe method in string storage', () => {
       </StrictMode>,
     )
 
-    expect(subscribe).toHaveBeenCalledWith('dummy', expect.any(Function), 1)
+    expect(subscribe).toHaveBeenCalledWith('dummy', expect.any(Function))
   })
 
   it('createJSONStorage subscriber responds to events correctly', async () => {

From 38bb47ea2241aa0bc3b8eddd80afd8a6d9d7e512 Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Fri, 24 May 2024 14:02:33 +0900
Subject: [PATCH 4/4] refactor 2

---
 src/vanilla/utils/atomWithStorage.ts | 37 ++++++++++++++--------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/src/vanilla/utils/atomWithStorage.ts b/src/vanilla/utils/atomWithStorage.ts
index db0c6c99a4..b508854f37 100644
--- a/src/vanilla/utils/atomWithStorage.ts
+++ b/src/vanilla/utils/atomWithStorage.ts
@@ -161,25 +161,26 @@ export function createJSONStorage<Value>(
         callback(newValue)
       })
 
-  const subscriber: StringSubscribe | undefined =
-    getStringStorage()?.subscribe ||
-    (typeof window !== 'undefined' &&
-      typeof window.addEventListener === 'function' &&
-      ((key, callback) => {
-        if (!(getStringStorage() instanceof window.Storage)) {
-          return () => {}
+  let subscriber = getStringStorage()?.subscribe
+  if (
+    !subscriber &&
+    typeof window !== 'undefined' &&
+    typeof window.addEventListener === 'function' &&
+    window.Storage &&
+    getStringStorage() instanceof window.Storage
+  ) {
+    subscriber = (key, callback) => {
+      const storageEventCallback = (e: StorageEvent) => {
+        if (e.storageArea === getStringStorage() && e.key === key) {
+          callback(e.newValue)
         }
-        const storageEventCallback = (e: StorageEvent) => {
-          if (e.storageArea === getStringStorage() && e.key === key) {
-            callback(e.newValue)
-          }
-        }
-        window.addEventListener('storage', storageEventCallback)
-        return () => {
-          window.removeEventListener('storage', storageEventCallback)
-        }
-      })) ||
-    undefined
+      }
+      window.addEventListener('storage', storageEventCallback)
+      return () => {
+        window.removeEventListener('storage', storageEventCallback)
+      }
+    }
+  }
 
   if (subscriber) {
     storage.subscribe = createHandleSubscribe(subscriber)