Change useTheme logic

Author: forestream

packages/react/src/hooks/use-theme.ts

@@ -1,48 +1,13 @@
 'use client'
-import { useId, useState } from 'react'
 
-import type { DevupTheme } from '../types/theme'
-import { useSafeEffect } from './use-safe-effect'
+import { useSyncExternalStore } from 'react'
+import { themeStore } from 'src/stores/themeStore'
 
-let observer: null | MutationObserver = null
-const setThemeMap: Record<string, React.Dispatch<keyof DevupTheme>> = {}
-let globalTheme: keyof DevupTheme | null = null
-
-export function useTheme(): keyof DevupTheme | null {
-  const id = useId()
-  const [theme, setTheme] = useState<keyof DevupTheme | null>(globalTheme)
-  useSafeEffect(() => {
-    if (globalTheme !== null) return
-    const currentTheme = document.documentElement.getAttribute('data-theme')
-    if (currentTheme !== null && currentTheme !== theme)
-      setTheme(currentTheme as keyof DevupTheme)
-  }, [theme])
-  useSafeEffect(() => {
-    const targetNode = document.documentElement
-    setThemeMap[id] = setTheme
-    if (!observer) {
-      observer = new MutationObserver(() => {
-        const theme = document.documentElement.getAttribute('data-theme')
-        globalTheme = theme as keyof DevupTheme
-        for (const key in setThemeMap)
-          setThemeMap[key](theme as keyof DevupTheme)
-      })
-      observer.observe(targetNode, {
-        attributes: true,
-        attributeFilter: ['data-theme'],
-        childList: false,
-        subtree: false,
-        characterData: false,
-        attributeOldValue: false,
-        characterDataOldValue: false,
-      })
-    }
-
-    return () => {
-      delete setThemeMap[id]
-      if (observer && Object.keys(setThemeMap).length === 0)
-        observer.disconnect()
-    }
-  }, [id])
+export function useTheme() {
+  const theme = useSyncExternalStore(
+    themeStore.subscribe,
+    themeStore.get,
+    themeStore.get,
+  )
   return theme
 }

packages/react/src/stores/themeStore.ts

@@ -0,0 +1,57 @@
+'use client'
+
+import { DevupTheme } from '@devup-ui/react'
+
+type Theme = keyof DevupTheme | null
+type StoreChangeEvent = (newTheme: Theme) => void
+
+const initTheme = null
+
+export const themeStore = (() => {
+  if (typeof window === 'undefined')
+    return {
+      get: () => initTheme,
+      set: () => {},
+      subscribe: () => () => {},
+    }
+  const el = document.documentElement
+
+  const subscribers: Set<StoreChangeEvent> = new Set()
+  let theme: Theme = initTheme
+  const get = () => {
+    return theme
+  }
+  const set = (newTheme: Theme) => {
+    theme = newTheme
+    subscribers.forEach((subscriber) => subscriber(theme))
+  }
+
+  const subscribe = (onStoreChange: StoreChangeEvent) => {
+    subscribers.add(onStoreChange)
+    set(el.getAttribute('data-theme') as Theme)
+    return () => subscribers.delete(onStoreChange)
+  }
+
+  const mo = new MutationObserver((mutations) => {
+    mutations.forEach((mutation) => {
+      if (
+        mutation.type === 'attributes' &&
+        mutation.target instanceof HTMLElement
+      ) {
+        set(mutation.target.getAttribute('data-theme') as Theme)
+      }
+    })
+  })
+
+  mo.observe(el, {
+    attributes: true,
+    attributeFilter: ['data-theme'],
+    subtree: false,
+  })
+
+  return {
+    get,
+    set,
+    subscribe,
+  }
+})()