Track lazy loaded components

packages/signia-react/src/track.test.tsx

@@ -1,4 +1,4 @@
-import { createRef, forwardRef, memo, useEffect, useImperativeHandle } from 'react'
+import { createRef, forwardRef, lazy, memo, Suspense, useEffect, useImperativeHandle } from 'react'
 import { act, create, ReactTestRenderer } from 'react-test-renderer'
 import { atom } from 'signia'
 import { track } from './track.js'
@@ -131,17 +131,19 @@ test('tracked components can use refs', async () => {
 	expect(ref.current?.handle).toBe('world')
 })
 
-test('tracked components update when the state they refernce updates', async () => {
+test('tracked components update when the state they reference updates', async () => {
 	const a = atom('a', 1)
 
-	const C = track(function Component() {
+	const Component = function Component() {
 		return <>{a.value}</>
-	})
+	}
+
+	const Tracked = track(Component)
 
 	let view: ReactTestRenderer
 
 	await act(() => {
-		view = create(<C />)
+		view = create(<Tracked />)
 	})
 
 	expect(view!.toJSON()).toMatchInlineSnapshot(`"1"`)
@@ -225,3 +227,98 @@ test("tracked zombie-children don't throw", async () => {
 		]
 	`)
 })
+
+describe('lazy components', () => {
+	test("are memo'd when tracked", async () => {
+		let numRenders = 0
+		const Component = function Component({ a, b, c }: { a: string; b: string; c: string }) {
+			numRenders++
+			return (
+				<>
+					{a}
+					{b}
+					{c}
+				</>
+			)
+		}
+
+		const Lazy = lazy(() => Promise.resolve({ default: Component }))
+		const TrackedLazy = track(Lazy)
+
+		let view: ReactTestRenderer
+		await act(() => {
+			view = create(
+				<Suspense>
+					<TrackedLazy a="a" b="b" c="c" />
+				</Suspense>
+			)
+		})
+
+		expect(view!.toJSON()).toMatchInlineSnapshot(`
+			[
+			  "a",
+			  "b",
+			  "c",
+			]
+		`)
+
+		expect(numRenders).toBe(1)
+
+		await act(() => {
+			view!.update(
+				<Suspense>
+					<TrackedLazy a="a" b="b" c="c" />
+				</Suspense>
+			)
+		})
+
+		expect(numRenders).toBe(1)
+
+		await act(() => {
+			view!.update(
+				<Suspense>
+					<TrackedLazy a="a" b="b" c="d" />
+				</Suspense>
+			)
+		})
+
+		expect(numRenders).toBe(2)
+
+		expect(view!.toJSON()).toMatchInlineSnapshot(`
+			[
+			  "a",
+			  "b",
+			  "d",
+			]
+		`)
+	})
+
+	test('update when the state they reference updates', async () => {
+		const a = atom('a', 1)
+
+		const Component = function Component() {
+			return <>{a.value}</>
+		}
+
+		const Lazy = lazy(() => Promise.resolve({ default: Component }))
+		const TrackedLazy = track(Lazy)
+
+		let view: ReactTestRenderer
+
+		await act(() => {
+			view = create(
+				<Suspense>
+					<TrackedLazy />
+				</Suspense>
+			)
+		})
+
+		expect(view!.toJSON()).toMatchInlineSnapshot(`"1"`)
+
+		await act(() => {
+			a.set(2)
+		})
+
+		expect(view!.toJSON()).toMatchInlineSnapshot(`"2"`)
+	})
+})

packages/signia-react/src/track.ts

@@ -1,4 +1,4 @@
-import React, { forwardRef, FunctionComponent, memo } from 'react'
+import React, { forwardRef, FunctionComponent, lazy, LazyExoticComponent, memo } from 'react'
 import { useStateTracking } from './useStateTracking.js'
 
 export const ProxyHandlers = {
@@ -20,9 +20,15 @@ export const ProxyHandlers = {
 	},
 }
 
+export const ReactLazySymbol = Symbol.for('react.lazy')
 export const ReactMemoSymbol = Symbol.for('react.memo')
 export const ReactForwardRefSymbol = Symbol.for('react.forward_ref')
 
+interface LazyFunctionComponent<T extends FunctionComponent<any>> extends LazyExoticComponent<T> {
+	_init: (arg: unknown) => FunctionComponent
+	_payload: { status: number; _result: FunctionComponent }
+}
+
 /**
  * Returns a tracked version of the given component.
  * Any signals whose values are read while the component renders will be tracked.
@@ -54,6 +60,21 @@ export function track<T extends FunctionComponent<any>>(
 	if ($$typeof === ReactForwardRefSymbol) {
 		return memo(forwardRef(new Proxy((baseComponent as any).render, ProxyHandlers) as any)) as any
 	}
+	if ($$typeof === ReactLazySymbol) {
+		let result: undefined | FunctionComponent
+
+		return memo(
+			lazy(() => {
+				if (!result) {
+					const { _init: init, _payload: payload } =
+						baseComponent as unknown as LazyFunctionComponent<any>
+					const loaded = init(payload)
+					result = new Proxy(loaded, ProxyHandlers)
+				}
+				return Promise.resolve({ default: result })
+			})
+		) as any
+	}
 
 	return memo(new Proxy(baseComponent, ProxyHandlers) as any, compare) as any
 }

packages/signia-react/src/wrapJsx.ts

@@ -28,6 +28,7 @@ import { track } from './track.js'
 
 const ReactMemoType = Symbol.for('react.memo') // https://github.com/facebook/react/blob/346c7d4c43a0717302d446da9e7423a8e28d8996/packages/shared/ReactSymbols.js#L30
 const ReactForwardRefType = Symbol.for('react.forward_ref')
+const ReactLazyType = Symbol.for('react.lazy')
 const ProxyInstance = new WeakMap<FunctionComponent<any>, FunctionComponent<any>>()
 
 function proxyFunctionalComponent(Component: FunctionComponent<any>) {
@@ -50,6 +51,8 @@ export function wrapJsx<T>(jsx: T): T {
 				type = proxyFunctionalComponent(type.type)
 			} else if (type.$$typeof === ReactForwardRefType) {
 				type = proxyFunctionalComponent(type)
+			} else if (type.$$typeof === ReactLazyType) {
+				type = proxyFunctionalComponent(type)
 			}
 		}