From c62d7846977f19f50e696ea7281ae9490c939cb8 Mon Sep 17 00:00:00 2001
From: Olivier Tassinari <olivier.tassinari@gmail.com>
Date: Wed, 9 Nov 2022 00:11:30 +0100
Subject: [PATCH] [docs] Fix bugs with live edit demos

---
 docs/src/modules/components/Demo.js | 32 +++++++++++++++--------------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/docs/src/modules/components/Demo.js b/docs/src/modules/components/Demo.js
index 2be1e42b13ad85..8aca997c4935ab 100644
--- a/docs/src/modules/components/Demo.js
+++ b/docs/src/modules/components/Demo.js
@@ -98,7 +98,6 @@ function useDemoElement({
   demoOptions,
   demoData,
   editorCode,
-  initialEditorCode,
   setDebouncedError,
 }) {
   const debouncedSetError = React.useMemo(
@@ -114,17 +113,9 @@ function useDemoElement({
 
   // Memoize to avoid rendering the demo more than it needs to be.
   // For example, avoid a render when the demo is hovered.
-  return React.useMemo(() => {
-    if (
-      // No need for a live environment if the code matches with the component rendered server-side.
-      editorCode.value === initialEditorCode ||
-      // A limitation from https://github.com/nihgwu/react-runner, we can inject the `window` of the iframe
-      demoOptions.disableLiveEdit
-    ) {
-      return <demoData.Component />;
-    }
-
-    return (
+  const BundledComponent = React.useMemo(() => <demoData.Component />, [demoData]);
+  const LiveComponent = React.useMemo(
+    () => (
       <ReactRunner
         scope={demoData.scope}
         onError={debouncedSetError}
@@ -137,8 +128,17 @@ function useDemoElement({
             : editorCode.value
         }
       />
-    );
-  }, [demoData, demoOptions.disableLiveEdit, editorCode, initialEditorCode, debouncedSetError]);
+    ),
+    [demoData, debouncedSetError, editorCode.isPreview, editorCode.value],
+  );
+
+  const renderBundledComponent =
+    // No need for a live environment if the code matches with the component rendered server-side.
+    editorCode.value === editorCode.initialEditorCode ||
+    // A limitation from https://github.com/nihgwu/react-runner, we can inject the `window` of the iframe
+    demoOptions.disableLiveEdit;
+
+  return renderBundledComponent ? BundledComponent : LiveComponent;
 }
 
 const Root = styled('div')(({ theme }) => ({
@@ -374,12 +374,14 @@ export default function Demo(props) {
   const [editorCode, setEditorCode] = React.useState({
     value: initialEditorCode,
     isPreview,
+    initialEditorCode,
   });
 
   const resetDemo = () => {
     setEditorCode({
       value: initialEditorCode,
       isPreview,
+      initialEditorCode,
     });
     setDemoKey();
   };
@@ -388,6 +390,7 @@ export default function Demo(props) {
     setEditorCode({
       value: initialEditorCode,
       isPreview,
+      initialEditorCode,
     });
   }, [initialEditorCode, isPreview]);
 
@@ -397,7 +400,6 @@ export default function Demo(props) {
     demoOptions,
     demoData,
     editorCode,
-    initialEditorCode,
     setDebouncedError,
   });