Fix syntax highlighter flicker by lazy loading theme and component (#1846)

Instead of just dynamically importing the theme and separately loading lazily the prism highlighter, we do it together and we avoid the flickering issue

Author: learningbizz

packages/web-ui/src/ds/atoms/CodeBlock/SyntaxHightlighter/LazyPrism/index.tsx

@@ -0,0 +1,44 @@
+'use client'
+import { Prism as PrismHighlighter } from 'react-syntax-highlighter'
+import { CurrentTheme } from '../../../../../constants'
+import { ComponentProps, use } from 'react'
+import { cn } from '../../../../../lib/utils'
+
+type Props = Omit<ComponentProps<typeof PrismHighlighter>, 'style'> & {
+  currentTheme?: string
+}
+
+const darkThemePromise = import(
+  'react-syntax-highlighter/dist/esm/styles/prism/one-dark'
+).then((module) => module.default)
+
+const lightThemePromise = import(
+  'react-syntax-highlighter/dist/esm/styles/prism/one-light'
+).then((module) => module.default)
+
+export default function LazyPrism({
+  children,
+  className,
+  language,
+  currentTheme,
+}: Props) {
+  const style = use(
+    currentTheme === CurrentTheme.Dark ? darkThemePromise : lightThemePromise,
+  )
+
+  return (
+    <PrismHighlighter
+      language={language}
+      style={style}
+      className={cn('text-xs', className)}
+      customStyle={{
+        borderRadius: '0.375rem',
+        padding: '1rem',
+        lineHeight: '1.25rem',
+        margin: '0',
+      }}
+    >
+      {children}
+    </PrismHighlighter>
+  )
+}

packages/web-ui/src/ds/atoms/CodeBlock/SyntaxHightlighter/index.tsx

@@ -1,73 +1,29 @@
 'use client'
 
-import React, {
-  lazy,
-  ComponentProps,
-  Suspense,
-  useEffect,
-  useState,
-  useMemo,
-} from 'react'
-import { cn } from '../../../../lib/utils'
-import { CurrentTheme } from '../../../../constants'
+import { ComponentProps, Suspense, lazy } from 'react'
+import { Prism as PrismHighlighter } from 'react-syntax-highlighter'
+import { Skeleton } from '../../Skeleton'
 
-const LazyPrism = lazy(() =>
-  import('react-syntax-highlighter').then((mod) => ({ default: mod.Prism })),
-)
-
-type PrismThemeStyle = { [key: string]: React.CSSProperties } | undefined
-
-type Props = Omit<ComponentProps<typeof LazyPrism>, 'style'> & {
-  currentTheme: string | undefined
+type Props = Omit<ComponentProps<typeof PrismHighlighter>, 'style'> & {
+  currentTheme?: string
 }
 
+const LazyPrism = lazy(() => import('./LazyPrism'))
+
 export function SyntaxHighlighter({
   children,
   className,
   language,
   currentTheme,
 }: Props) {
-  const [style, setStyle] = useState<PrismThemeStyle>(undefined)
-  useEffect(() => {
-    let active = true
-
-    const loadTheme = async () => {
-      let themeModule
-      if (currentTheme === CurrentTheme.Dark) {
-        themeModule = await import(
-          'react-syntax-highlighter/dist/esm/styles/prism/one-dark'
-        )
-      } else {
-        themeModule = await import(
-          'react-syntax-highlighter/dist/esm/styles/prism/one-light'
-        )
-      }
-
-      if (active) {
-        setStyle(themeModule.default)
-      }
-    }
-    loadTheme()
-
-    return () => {
-      active = false
-    }
-  }, [currentTheme])
-
   return (
-    <Suspense fallback={<div className='p-4 text-sm'>Loading code...</div>}>
+    <Suspense fallback={<Skeleton height='h5' />}>
       <LazyPrism
-        className={cn('text-xs', className)}
+        className={className}
         language={language}
-        style={style}
-        customStyle={{
-          borderRadius: '0.375rem',
-          padding: '1rem',
-          lineHeight: '1.25rem',
-          margin: '0',
-        }}
+        currentTheme={currentTheme}
       >
-        {useMemo(() => children, [children])}
+        {children}
       </LazyPrism>
     </Suspense>
   )

packages/web-ui/src/ds/atoms/CodeBlock/index.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import { useTheme } from 'next-themes'
-import { memo, ReactNode } from 'react'
+import { useMemo, memo, ReactNode } from 'react'
 
 import { CurrentTheme } from '../../../constants'
 import { cn } from '../../../lib/utils'
@@ -19,8 +19,10 @@ interface CodeBlockProps {
   bgColor?: string
 }
 
-export function useCodeBlockBackgroundColor(override?: string) {
-  const { resolvedTheme } = useTheme()
+function getCodeBlockBackgroundColor(
+  override?: string,
+  resolvedTheme?: string,
+) {
   if (override) return override
   if (resolvedTheme === CurrentTheme.Light) return 'bg-backgroundCode'
   return 'bg-[#282c34]'
@@ -37,7 +39,12 @@ const Content = memo(
     bgColor: overrideBgColor,
   }: CodeBlockProps) => {
     const { resolvedTheme } = useTheme()
-    const bgColor = useCodeBlockBackgroundColor(overrideBgColor)
+
+    const bgColor = useMemo(
+      () => getCodeBlockBackgroundColor(overrideBgColor, resolvedTheme),
+      [overrideBgColor, resolvedTheme],
+    )
+
     return (
       <div
         className={cn(