From 56e04f869bd2b6a5b15780353a6d2f36f0f30d18 Mon Sep 17 00:00:00 2001
From: Andrey Polischuk <me@andrepolischuk.com>
Date: Thu, 14 Dec 2023 09:58:05 +0300
Subject: [PATCH] fix: don't use react lazy after preloading the component

---
 .size-limit.json        |  4 ++--
 src/common/types.ts     |  4 ++--
 src/components/lazy.tsx | 39 +++++++++++++++++++++++++++------------
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/.size-limit.json b/.size-limit.json
index 4fffbe6..edc78e4 100644
--- a/.size-limit.json
+++ b/.size-limit.json
@@ -1,10 +1,10 @@
 [
   {
     "path": "dist/client/index.js",
-    "limit": "7.1 KB"
+    "limit": "7.15 KB"
   },
   {
     "path": "dist/server/index.js",
-    "limit": "7.1 KB"
+    "limit": "7.05 KB"
   }
 ]
diff --git a/src/common/types.ts b/src/common/types.ts
index b51b967..d1e079a 100644
--- a/src/common/types.ts
+++ b/src/common/types.ts
@@ -78,6 +78,7 @@ export type GetMetaData = Loader<MetaData, {data: InitialData}>
  * ```
  */
 
+/** Page component */
 export type PageComponent<P = any> = React.ComponentType<
   P & {isLoading: boolean}
 > & {
@@ -86,8 +87,7 @@ export type PageComponent<P = any> = React.ComponentType<
 }
 
 /** Lazy page component */
-export interface LazyPageComponent<P = any>
-  extends React.LazyExoticComponent<PageComponent<P>> {
+export type LazyPageComponent<P = any> = PageComponent<P> & {
   preload: Loader<void>
   getMetaData: GetMetaData
   getInitialData: GetInitialData
diff --git a/src/components/lazy.tsx b/src/components/lazy.tsx
index 3357b03..36ce9a8 100644
--- a/src/components/lazy.tsx
+++ b/src/components/lazy.tsx
@@ -1,4 +1,4 @@
-import {lazy as reactLazy} from 'react'
+import React, {lazy as reactLazy} from 'react'
 import type {
   Context,
   PageComponent,
@@ -8,9 +8,14 @@ import type {
   InitialData
 } from '../common/types'
 
-/** Component factory */
+/** Page component factory */
+export interface ComponentModule {
+  default: PageComponent
+}
+
+/** Page component factory */
 export interface ComponentFactory {
-  (): Promise<{default: PageComponent}>
+  (): Promise<ComponentModule>
 }
 
 /** Data factory */
@@ -32,12 +37,16 @@ export interface DataFactory<T, C = any> {
  * ```
  */
 export const lazy = (componentFactory: ComponentFactory): LazyPageComponent => {
-  let promise: ReturnType<ComponentFactory>
+  let componentModule: ComponentModule
+  let componentPromise: ReturnType<ComponentFactory>
 
-  const onceFactory: ComponentFactory = () => {
-    promise ??= componentFactory()
+  const onceFactory: ComponentFactory = async () => {
+    if (!componentModule) {
+      componentPromise ??= componentFactory()
+      componentModule = await componentPromise
+    }
 
-    return promise
+    return componentModule
   }
 
   function loaderFactory<T, C = any>(
@@ -50,17 +59,23 @@ export const lazy = (componentFactory: ComponentFactory): LazyPageComponent => {
     }
   }
 
-  const Component = reactLazy(onceFactory) as LazyPageComponent
+  const ReactLazy = reactLazy(onceFactory)
+
+  const Lazy: LazyPageComponent = (props) => {
+    const Component = componentModule?.default ?? ReactLazy
+
+    return <Component {...props} />
+  }
 
-  Component.getMetaData = loaderFactory<MetaData, {data: InitialData}>(
+  Lazy.getMetaData = loaderFactory<MetaData, {data: InitialData}>(
     ({getMetaData}, context) => getMetaData?.(context)
   )
 
-  Component.getInitialData = loaderFactory<InitialData>(
+  Lazy.getInitialData = loaderFactory<InitialData>(
     ({getInitialData}, context) => getInitialData?.(context)
   )
 
-  Component.preload = loaderFactory<void>()
+  Lazy.preload = loaderFactory<void>()
 
-  return Component
+  return Lazy
 }