Fix: Flickering page on route navigation

packages/router/src/page-loader.tsx

@@ -1,8 +1,14 @@
 import React, { useContext } from 'react'
 
 import isEqual from 'lodash.isequal'
+import { unstable_batchedUpdates } from 'react-dom'
 
 import { Spec } from './router'
+import {
+  RouterState,
+  useRouterState,
+  useRouterStateSetter,
+} from './router-context'
 import {
   createNamedContext,
   getAnnouncement,
@@ -31,23 +37,26 @@ export const usePageLoadingContext = () => {
 
 type synchonousLoaderSpec = () => { default: React.ComponentType }
 interface State {
-  Page?: React.ComponentType
-  pageName?: string
   slowModuleImport: boolean
-  params?: Record<string, string>
 }
 
-interface Props {
+interface PageLoaderProps {
   spec: Spec
   delay?: number
   params?: Record<string, string>
   whileLoadingPage?: () => React.ReactElement | null
 }
+interface Props extends PageLoaderProps {
+  currentRoute?: {
+    pageName: string
+    Page?: React.ComponentType | null
+    params?: Record<string, string>
+  }
+  setRouterState: React.Dispatch<Partial<RouterState>>
+}
 
-export class PageLoader extends React.Component<Props> {
+class PageLoaderWithRouterContext extends React.Component<Props> {
   state: State = {
-    Page: undefined,
-    pageName: undefined,
     slowModuleImport: false,
   }
 
@@ -60,11 +69,15 @@ export class PageLoader extends React.Component<Props> {
     return !isEqual(p1.params, p2.params)
   }
 
-  stateChanged = (s1: State, s2: State) => {
-    if (s1.pageName !== s2.pageName) {
+  stateChanged = (s1: State, s2: State) =>
+    s1.slowModuleImport !== s2.slowModuleImport
+
+  currentRouteChanged = (p1: Props, p2: Props) => {
+    if (p1.currentRoute?.pageName !== p2.currentRoute?.pageName) {
       return true
     }
-    return !isEqual(s1.params, s2.params)
+
+    return !isEqual(p1.currentRoute?.params, p2.currentRoute?.params)
   }
 
   shouldComponentUpdate(nextProps: Props, nextState: State) {
@@ -74,7 +87,10 @@ export class PageLoader extends React.Component<Props> {
       return false
     }
 
-    if (this.stateChanged(this.state, nextState)) {
+    if (
+      this.currentRouteChanged(this.props, nextProps) ||
+      this.stateChanged(this.state, nextState)
+    ) {
       return true
     }
 
@@ -94,7 +110,10 @@ export class PageLoader extends React.Component<Props> {
       this.startPageLoadTransition(this.props)
     }
 
-    if (this.stateChanged(prevState, this.state)) {
+    if (
+      this.currentRouteChanged(prevProps, this.props) ||
+      this.stateChanged(prevState, this.state)
+    ) {
       global?.scrollTo(0, 0)
       if (this.announcementRef.current) {
         this.announcementRef.current.innerText = getAnnouncement()
@@ -137,61 +156,91 @@ export class PageLoader extends React.Component<Props> {
     // Remove the timeout because the page has loaded.
     this.clearLoadingTimeout()
 
-    this.setState({
-      pageName: name,
-      Page: module.default,
-      slowModuleImport: false,
-      params: props.params,
+    // Update downloaded page in the Router Context to avoid
+    // blank page (page: undefined) when Route unmounts (and therefore PageLoader)
+
+    // Batched update is to avoid unnecessary re-rendering
+    unstable_batchedUpdates(() => {
+      this.props.setRouterState({
+        activeRoute: {
+          pageName: name,
+          Page: module.default,
+          params: props.params,
+        },
+      })
+
+      this.setState({
+        slowModuleImport: false,
+      })
     })
   }
 
   render() {
-    const { Page } = this.state
-
     if (global.__REDWOOD__PRERENDERING) {
       // babel autoloader plugin uses withStaticImport in prerender mode
       // override the types for this condition
+      const { params: newParams } = this.props
       const syncPageLoader = this.props.spec
         .loader as unknown as synchonousLoaderSpec
       const PageFromLoader = syncPageLoader().default
 
       return (
         <PageLoadingContext.Provider value={{ loading: false }}>
-          <PageFromLoader {...this.state.params} />
+          <PageFromLoader {...newParams} />
         </PageLoadingContext.Provider>
       )
     }
 
-    if (Page) {
-      return (
-        <PageLoadingContext.Provider
-          value={{ loading: this.state.slowModuleImport }}
-        >
-          <Page {...this.state.params} />
-          <div
-            id="redwood-announcer"
-            style={{
-              position: 'absolute',
-              top: 0,
-              width: 1,
-              height: 1,
-              padding: 0,
-              overflow: 'hidden',
-              clip: 'rect(0, 0, 0, 0)',
-              whiteSpace: 'nowrap',
-              border: 0,
-            }}
-            role="alert"
-            aria-live="assertive"
-            aria-atomic="true"
-            ref={this.announcementRef}
-          ></div>
-        </PageLoadingContext.Provider>
-      )
-    } else {
-      return this.state.slowModuleImport
-        ? this.props.whileLoadingPage?.() || null
-        : null
+    // Page will always be there, either old one or the latest (except first load / hard refresh)
+    // So we have to check if we want to show loading state before rendering the Page
+    const { slowModuleImport } = this.state
+    const { whileLoadingPage } = this.props
+    if (slowModuleImport && whileLoadingPage) {
+      return this.props.whileLoadingPage?.() || null
+    }
+
+    // currentRoute holds the last page until new page chunk is loaded
+    const { Page, params } = this.props.currentRoute || {}
+    if (!Page) {
+      return null
     }
+
+    return (
+      <PageLoadingContext.Provider
+        value={{ loading: this.state.slowModuleImport }}
+      >
+        <Page {...params} />
+        <div
+          id="redwood-announcer"
+          style={{
+            position: 'absolute',
+            top: 0,
+            width: 1,
+            height: 1,
+            padding: 0,
+            overflow: 'hidden',
+            clip: 'rect(0, 0, 0, 0)',
+            whiteSpace: 'nowrap',
+            border: 0,
+          }}
+          role="alert"
+          aria-live="assertive"
+          aria-atomic="true"
+          ref={this.announcementRef}
+        ></div>
+      </PageLoadingContext.Provider>
+    )
   }
 }
+
+export const PageLoader = (props: PageLoaderProps) => {
+  const { activeRoute } = useRouterState()
+  const setRouterState = useRouterStateSetter()
+  return (
+    <PageLoaderWithRouterContext
+      currentRoute={activeRoute}
+      setRouterState={setRouterState}
+      {...props}
+    />
+  )
+}

packages/router/src/router-context.tsx

@@ -10,6 +10,11 @@ export interface RouterState {
   paramTypes?: Record<string, ParamType>
   pageLoadingDelay?: number
   useAuth: typeof useAuth
+  activeRoute?: {
+    pageName: string
+    Page?: React.ComponentType
+    params?: Record<string, string>
+  }
 }
 
 const RouterStateContext = createContext<RouterState | undefined>(undefined)