Refactor to useSyncExternalStore

src/index.ts

@@ -1,4 +1,4 @@
-import { useEffect, useLayoutEffect, useRef, useState } from "react"
+import { useCallback, useRef, useSyncExternalStore } from "react"
 
 /*
 
@@ -155,61 +155,36 @@ export const makeStore = <T extends IState>(initialState: T): Store<T> => {
                                        ▀
 */
 
-/**
- * If a component is loaded in a SSR context and imports the useStore hook,
- * React will trigger a warning that says: "useLayoutEffect does nothing on
- * the server". To surpress this warning, we need to check if window is
- * defined.
- */
-const useIsomorphicLayoutEffect = typeof window !== "undefined" ? useLayoutEffect : useEffect
-
 /**
  * Provides reactive read access to a Statery store. Returns a proxy object that
  * provides direct access to the store's state and makes sure that the React component
- * it was invoked from automaticaly re-renders when any of the data it uses is updated.
+ * it was invoked from automatically re-renders when any of the data it uses is updated.
  *
  * @param store The Statery store to access.
  */
 export const useStore = <T extends IState>(store: Store<T>): T => {
-  /* A cheap version state that we will bump in order to re-render the component. */
-  const [v, setVersion] = useState(0)
-
   /* A set containing all props that we're interested in. */
   const subscribedProps = useConst(() => new Set<keyof T>())
+  const prevSnapshot = useRef(store.state)
 
-  /* Grab a copy of the state at the time the component is rendering; then, in an effect,
-  check if there have already been any updates. This can happen because something that
-  was rendered alongside this component wrote into the store immediately, possibly
-  through a function ref. If we detect a change related to the props we're interested in,
-  force the component to reload. */
-  const initialState = useConst(() => store.state)
+  const subscribe = useCallback((listener: () => void) => {
+    store.subscribe(listener)
+    return () => store.unsubscribe(listener)
+  }, [])
 
-  useIsomorphicLayoutEffect(() => {
-    if (store.state === initialState) return
+  const getSnapshot = useCallback(() => {
+    const props = Array.from(subscribedProps)
+    const hasChanged = props.some((prop) => store.state[prop] !== prevSnapshot.current[prop])
 
-    subscribedProps.forEach((prop) => {
-      if (initialState[prop] !== store.state[prop]) {
-        setVersion((v) => v + 1)
-        return
-      }
-    })
-  }, [store])
-
-  /* Subscribe to changes in the store. */
-  useIsomorphicLayoutEffect(() => {
-    const listener: Listener<T> = (updates: Partial<T>) => {
-      /* If there is at least one prop being updated that we're interested in,
-         bump our local version. */
-      if (Object.keys(updates).find((prop) => subscribedProps.has(prop))) {
-        setVersion((v) => v + 1)
-      }
+    if (hasChanged) {
+      prevSnapshot.current = store.state
     }
 
-    /* Mount & unmount the listener */
-    store.subscribe(listener)
-    return () => void store.unsubscribe(listener)
+    return prevSnapshot.current
   }, [store])
 
+  const snapshot = useSyncExternalStore(subscribe, getSnapshot)
+
   return new Proxy<Record<any, any>>(
     {},
     {
@@ -218,7 +193,7 @@ export const useStore = <T extends IState>(store: Store<T>): T => {
         subscribedProps.add(prop)
 
         /* Return the current value of the property. */
-        return store.state[prop]
+        return snapshot[prop]
       }
     }
   )

test/hooks.test.tsx

@@ -204,4 +204,28 @@ describe("useStore", () => {
     fireEvent.click(page.getByText("Toggle"))
     await page.findByText("Active: No")
   })
+
+  it("re-renders if a property is changed during the render phase", async () => {
+    let changedDuringRender = false
+    let renders = 0
+
+    const store = makeStore({ lightning: "Slow" })
+
+    const Lightning = () => {
+      renders++
+      const { lightning } = useStore(store)
+
+      if (!changedDuringRender) {
+        store.set({ lightning: "Fast" })
+        changedDuringRender = true
+      }
+
+      return <p>Lightning: {lightning}</p>
+    }
+
+    const page = render(<Lightning />)
+
+    await page.findByText("Lightning: Fast")
+    expect(renders).toBe(2)
+  })
 })