Nitpicks for code splitting (#4346)

beta/src/components/MDX/MDXComponents.tsx

@@ -19,7 +19,6 @@ import Intro from './Intro';
 import Link from './Link';
 import {PackageImport} from './PackageImport';
 import Recap from './Recap';
-import dynamic from 'next/dynamic';
 import Sandpack from './Sandpack';
 import SimpleCallout from './SimpleCallout';
 import TerminalBlock from './TerminalBlock';
@@ -363,7 +362,6 @@ export const MDXComponents = {
   code: CodeBlock,
   // The code block renders <pre> so we just want a div here.
   pre: (p: JSX.IntrinsicElements['div']) => <div {...p} />,
-  // Scary: dynamic(() => import('./Scary')),
   APIAnatomy,
   AnatomyStep,
   CodeDiagram,

beta/src/components/MDX/Sandpack/SandpackWrapper.tsx → beta/src/components/MDX/Sandpack/SandpackRoot.tsx

@@ -2,14 +2,12 @@
  * Copyright (c) Facebook, Inc. and its affiliates.
  */
 
-import React from 'react';
+import * as React from 'react';
-import {
+import {SandpackProvider} from '@codesandbox/sandpack-react';
-  SandpackProvider,
-  SandpackSetup,
-  SandpackFile,
-} from '@codesandbox/sandpack-react';
-
 import {CustomPreset} from './CustomPreset';
+import {createFileMap} from './utils';
+
+import type {SandpackSetup} from '@codesandbox/sandpack-react';
 
 type SandpackProps = {
   children: React.ReactChildren;
@@ -17,7 +15,6 @@ type SandpackProps = {
   setup?: SandpackSetup;
   showDevTools?: boolean;
 };
-import {reducedCodeSnippet} from './utils';
 
 const sandboxStyle = `
 * {
@@ -65,13 +62,13 @@ ul {
 }
 `.trim();
 
-function SandpackWrapper(props: SandpackProps) {
+function SandpackRoot(props: SandpackProps) {
   let {children, setup, autorun = true, showDevTools = false} = props;
   const [devToolsLoaded, setDevToolsLoaded] = React.useState(false);
   let codeSnippets = React.Children.toArray(children) as React.ReactElement[];
   let isSingleFile = true;
 
-  const files = reducedCodeSnippet(codeSnippets);
+  const files = createFileMap(codeSnippets);
 
   files['/styles.css'] = {
     code: [sandboxStyle, files['/styles.css']?.code ?? ''].join('\n\n'),
@@ -97,6 +94,6 @@ function SandpackWrapper(props: SandpackProps) {
   );
 }
 
-SandpackWrapper.displayName = 'Sandpack';
+SandpackRoot.displayName = 'Sandpack';
 
-export default SandpackWrapper;
+export default SandpackRoot;

beta/src/components/MDX/Sandpack/index.tsx

@@ -1,8 +1,14 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ */
+
 import * as React from 'react';
-import {reducedCodeSnippet} from './utils';
+import dynamic from 'next/dynamic';
-const Sandpack = React.lazy(() => import('./SandpackWrapper'));
+import {createFileMap} from './utils';
+
+const SandpackRoot = dynamic(() => import('./SandpackRoot'), {suspense: true});
 
-const SandpackFallBack = ({code}: {code: string}) => (
+const SandpackGlimmer = ({code}: {code: string}) => (
   <div className="sandpack-container my-8">
     <div className="shadow-lg dark:shadow-lg-dark rounded-lg">
       <div className="bg-wash h-10 dark:bg-card-dark flex justify-between items-center relative z-10 border-b border-border dark:border-border-dark rounded-t-lg rounded-b-none">
@@ -41,11 +47,10 @@ const SandpackFallBack = ({code}: {code: string}) => (
 );
 
 export default React.memo(function SandpackWrapper(props: any): any {
-  const codeSnippet = reducedCodeSnippet(
+  const codeSnippet = createFileMap(React.Children.toArray(props.children));
-    React.Children.toArray(props.children)
-  );
 
-  // To set the active file in the fallback we have to find the active file first. If there are no active files we fallback to App.js as default
+  // To set the active file in the fallback we have to find the active file first.
+  // If there are no active files we fallback to App.js as default.
   let activeCodeSnippet = Object.keys(codeSnippet).filter(
     (fileName) =>
       codeSnippet[fileName]?.active === true &&
@@ -59,10 +64,8 @@ export default React.memo(function SandpackWrapper(props: any): any {
   }
 
   return (
-    <>
+    <React.Suspense fallback={<SandpackGlimmer code={activeCode} />}>
-      <React.Suspense fallback={<SandpackFallBack code={activeCode} />}>
+      <SandpackRoot {...props} />
-        <Sandpack {...props} />
+    </React.Suspense>
-      </React.Suspense>
-    </>
   );
 });

beta/src/components/MDX/Sandpack/utils.ts

@@ -48,8 +48,7 @@ export const computeViewportSize = (
   return viewport;
 };
 
-//TODO: revisit to reduce this (moved this to utils from sandpackWrapper since it is being used for finding active file for fallBack too)
+export const createFileMap = (codeSnippets: any) => {
-export const reducedCodeSnippet = (codeSnippets: any) => {
   let isSingleFile = true;
 
   return codeSnippets.reduce(