feat(jsx/dom): support getServerSnapshot in useSyncExternalStore (#2646)

* feat(jsx/dom): support getServerSnapshot in useSyncExternalStore

* chore: denoify

* test: add tests for hooks with `.toString()`

* test: fix typo.

* refactor: set the flag to false at the end of top level render

* chore: denoify

deno_dist/jsx/dom/render.ts

@@ -82,8 +82,10 @@ export type Context =
       PendingType, // PendingType
       boolean, // got an error
       UpdateHook, // update hook
-      boolean // is in view transition
+      boolean, // is in view transition
+      boolean // is in top level render
     ]
+  | [PendingType, boolean, UpdateHook, boolean]
   | [PendingType, boolean, UpdateHook]
   | [PendingType, boolean]
   | [PendingType]
@@ -572,7 +574,10 @@ export const update = async (
 export const render = (jsxNode: unknown, container: Container) => {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   const node = buildNode({ tag: '', props: { children: jsxNode } } as any) as NodeObject
-  build([], node, undefined)
+  const context: Context = []
+  ;(context as Context)[4] = true // start top level render
+  build(context, node, undefined)
+  ;(context as Context)[4] = false // finish top level render
 
   const fragment = document.createDocumentFragment()
   apply(node, fragment)

deno_dist/jsx/hooks/index.ts

@@ -388,25 +388,32 @@ export const useImperativeHandle = <T>(
   }, deps)
 }
 
-let useSyncExternalStoreGetServerSnapshotNotified = false
 export const useSyncExternalStore = <T>(
   subscribe: (callback: (value: T) => void) => () => void,
   getSnapshot: () => T,
   getServerSnapshot?: () => T
 ): T => {
-  const [state, setState] = useState(getSnapshot())
-  useEffect(
-    () =>
-      subscribe(() => {
-        setState(getSnapshot())
-      }),
-    []
-  )
-
-  if (getServerSnapshot && !useSyncExternalStoreGetServerSnapshotNotified) {
-    useSyncExternalStoreGetServerSnapshotNotified = true
-    console.info('`getServerSnapshot` is not supported yet.')
+  const buildData = buildDataStack.at(-1) as [Context, unknown]
+  if (!buildData) {
+    // now a stringify process, maybe in server side
+    if (!getServerSnapshot) {
+      throw new Error('getServerSnapshot is required for server side rendering')
+    }
+    return getServerSnapshot()
   }
 
+  const [serverSnapshotIsUsed] = useState<boolean>(!!(buildData[0][4] && getServerSnapshot))
+  const [state, setState] = useState(() =>
+    serverSnapshotIsUsed ? (getServerSnapshot as () => T)() : getSnapshot()
+  )
+  useEffect(() => {
+    if (serverSnapshotIsUsed) {
+      setState(getSnapshot())
+    }
+    return subscribe(() => {
+      setState(getSnapshot())
+    })
+  }, [])
+
   return state
 }

src/jsx/dom/render.ts

@@ -82,8 +82,10 @@ export type Context =
       PendingType, // PendingType
       boolean, // got an error
       UpdateHook, // update hook
-      boolean // is in view transition
+      boolean, // is in view transition
+      boolean // is in top level render
     ]
+  | [PendingType, boolean, UpdateHook, boolean]
   | [PendingType, boolean, UpdateHook]
   | [PendingType, boolean]
   | [PendingType]
@@ -572,7 +574,10 @@ export const update = async (
 export const render = (jsxNode: unknown, container: Container) => {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   const node = buildNode({ tag: '', props: { children: jsxNode } } as any) as NodeObject
-  build([], node, undefined)
+  const context: Context = []
+  ;(context as Context)[4] = true // start top level render
+  build(context, node, undefined)
+  ;(context as Context)[4] = false // finish top level render
 
   const fragment = document.createDocumentFragment()
   apply(node, fragment)

src/jsx/hooks/dom.test.tsx

@@ -563,5 +563,37 @@ describe('Hooks', () => {
       await new Promise((r) => setTimeout(r))
       expect(root.innerHTML).toBe('<div>1</div><button>toggle</button>')
     })
+
+    it('with getServerSnapshot', async () => {
+      let count = 0
+      const unsubscribe = vi.fn()
+      const subscribe = vi.fn(() => unsubscribe)
+      const getSnapshot = vi.fn(() => count++)
+      const getServerSnapshot = vi.fn(() => 100)
+      const SubApp = () => {
+        const count = useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot)
+        return <div>{count}</div>
+      }
+      const App = () => {
+        const [show, setShow] = useState(true)
+        return (
+          <>
+            {show && <SubApp />}
+            <button onClick={() => setShow((s) => !s)}>toggle</button>
+          </>
+        )
+      }
+      render(<App />, root)
+      expect(root.innerHTML).toBe('<div>100</div><button>toggle</button>')
+      await new Promise((r) => setTimeout(r))
+      expect(root.innerHTML).toBe('<div>0</div><button>toggle</button>')
+      root.querySelector('button')?.click()
+      await new Promise((r) => setTimeout(r))
+      expect(root.innerHTML).toBe('<button>toggle</button>')
+      expect(unsubscribe).toBeCalled()
+      root.querySelector('button')?.click()
+      await new Promise((r) => setTimeout(r))
+      expect(root.innerHTML).toBe('<div>1</div><button>toggle</button>')
+    })
   })
 })

src/jsx/hooks/index.ts

@@ -388,25 +388,32 @@ export const useImperativeHandle = <T>(
   }, deps)
 }
 
-let useSyncExternalStoreGetServerSnapshotNotified = false
 export const useSyncExternalStore = <T>(
   subscribe: (callback: (value: T) => void) => () => void,
   getSnapshot: () => T,
   getServerSnapshot?: () => T
 ): T => {
-  const [state, setState] = useState(getSnapshot())
-  useEffect(
-    () =>
-      subscribe(() => {
-        setState(getSnapshot())
-      }),
-    []
-  )
-
-  if (getServerSnapshot && !useSyncExternalStoreGetServerSnapshotNotified) {
-    useSyncExternalStoreGetServerSnapshotNotified = true
-    console.info('`getServerSnapshot` is not supported yet.')
+  const buildData = buildDataStack.at(-1) as [Context, unknown]
+  if (!buildData) {
+    // now a stringify process, maybe in server side
+    if (!getServerSnapshot) {
+      throw new Error('getServerSnapshot is required for server side rendering')
+    }
+    return getServerSnapshot()
   }
 
+  const [serverSnapshotIsUsed] = useState<boolean>(!!(buildData[0][4] && getServerSnapshot))
+  const [state, setState] = useState(() =>
+    serverSnapshotIsUsed ? (getServerSnapshot as () => T)() : getSnapshot()
+  )
+  useEffect(() => {
+    if (serverSnapshotIsUsed) {
+      setState(getSnapshot())
+    }
+    return subscribe(() => {
+      setState(getSnapshot())
+    })
+  }, [])
+
   return state
 }

src/jsx/hooks/string.test.tsx

@@ -0,0 +1,40 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { jsx, useState, useSyncExternalStore } from '..'
+
+describe('useState', () => {
+  it('should be rendered with initial state', () => {
+    const Component = () => {
+      const [state] = useState('hello')
+      return <span>{state}</span>
+    }
+    const template = <Component />
+    expect(template.toString()).toBe('<span>hello</span>')
+  })
+})
+
+describe('useSyncExternalStore', () => {
+  it('should be rendered with result of getServerSnapshot()', () => {
+    const unsubscribe = vi.fn()
+    const subscribe = vi.fn(() => unsubscribe)
+    const getSnapshot = vi.fn()
+    const getServerSnapshot = vi.fn(() => 100)
+    const App = () => {
+      const count = useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot)
+      return <div>{count}</div>
+    }
+    const template = <App />
+    expect(template.toString()).toBe('<div>100</div>')
+    expect(unsubscribe).not.toBeCalled()
+    expect(subscribe).not.toBeCalled()
+    expect(getSnapshot).not.toBeCalled()
+  })
+
+  it('should raise an error if getServerShot() is not provided', () => {
+    const App = () => {
+      const count = useSyncExternalStore(vi.fn(), vi.fn())
+      return <div>{count}</div>
+    }
+    const template = <App />
+    expect(() => template.toString()).toThrowError()
+  })
+})