[ThemeProvider] Optimize theme changes when enabling CSS theme vari…

…ables (mui#44588)
sai6855 · Nov 29, 2024 · d0a5989 · d0a5989
1 parent 7c0b7e8
commit d0a5989
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 43 deletions.
diff --git a/packages/mui-material/src/styles/ThemeProviderWithVars.test.js b/packages/mui-material/src/styles/ThemeProviderWithVars.test.js
@@ -408,4 +408,35 @@ describe('[Material UI] ThemeProviderWithVars', () => {
       fireEvent.click(screen.getByText('Dark'));
     }).not.toErrorDev();
   });
+
+  it('theme should remain the same when ThemeProvider rerenders', () => {
+    const theme = createTheme({ cssVariables: true });
+
+    function Inner() {
+      const upperTheme = useTheme();
+      const themeRef = React.useRef(upperTheme);
+      const [changed, setChanged] = React.useState(false);
+      React.useEffect(() => {
+        if (themeRef.current !== upperTheme) {
+          setChanged(true);
+        }
+      }, [upperTheme]);
+      return changed ? <div data-testid="theme-changed" /> : null;
+    }
+    function App() {
+      const [, setState] = React.useState({});
+      const rerender = () => setState({});
+      return (
+        <ThemeProvider theme={theme}>
+          <button onClick={() => rerender()}>rerender</button>
+          <Inner />
+        </ThemeProvider>
+      );
+    }
+    render(<App />);
+
+    fireEvent.click(screen.getByRole('button'));
+
+    expect(screen.queryByTestId('theme-changed')).to.equal(null);
+  });
 });
diff --git a/packages/mui-system/src/cssVars/createCssVarsProvider.js b/packages/mui-system/src/cssVars/createCssVarsProvider.js
@@ -48,6 +48,9 @@ export default function createCssVarsProvider(options) {
 
   const useColorScheme = () => React.useContext(ColorSchemeContext) || defaultContext;
 
+  const defaultColorSchemes = {};
+  const defaultComponents = {};
+
   function CssVarsProvider(props) {
     const {
       children,
@@ -75,12 +78,12 @@ export default function createCssVarsProvider(options) {
       return typeof defaultTheme === 'function' ? defaultTheme() : defaultTheme;
     }, [themeProp]);
     const scopedTheme = initialTheme[themeId];
+    const restThemeProp = scopedTheme || initialTheme;
     const {
-      colorSchemes = {},
-      components = {},
+      colorSchemes = defaultColorSchemes,
+      components = defaultComponents,
       cssVarPrefix,
-      ...restThemeProp
-    } = scopedTheme || initialTheme;
+    } = restThemeProp;
     const joinedColorSchemes = Object.keys(colorSchemes)
       .filter((k) => !!colorSchemes[k])
       .join(',');
@@ -126,42 +129,46 @@ export default function createCssVarsProvider(options) {
       colorScheme = ctx.colorScheme;
     }
 
-    // `colorScheme` is undefined on the server and hydration phase
-    const calculatedColorScheme = colorScheme || restThemeProp.defaultColorScheme;
+    const memoTheme = React.useMemo(() => {
+      // `colorScheme` is undefined on the server and hydration phase
+      const calculatedColorScheme = colorScheme || restThemeProp.defaultColorScheme;
 
-    // 2. get the `vars` object that refers to the CSS custom properties
-    const themeVars = restThemeProp.generateThemeVars?.() || restThemeProp.vars;
+      // 2. get the `vars` object that refers to the CSS custom properties
+      const themeVars = restThemeProp.generateThemeVars?.() || restThemeProp.vars;
 
-    // 3. Start composing the theme object
-    const theme = {
-      ...restThemeProp,
-      components,
-      colorSchemes,
-      cssVarPrefix,
-      vars: themeVars,
-    };
-    if (typeof theme.generateSpacing === 'function') {
-      theme.spacing = theme.generateSpacing();
-    }
+      // 3. Start composing the theme object
+      const theme = {
+        ...restThemeProp,
+        components,
+        colorSchemes,
+        cssVarPrefix,
+        vars: themeVars,
+      };
+      if (typeof theme.generateSpacing === 'function') {
+        theme.spacing = theme.generateSpacing();
+      }
 
-    // 4. Resolve the color scheme and merge it to the theme
-    if (calculatedColorScheme) {
-      const scheme = colorSchemes[calculatedColorScheme];
-      if (scheme && typeof scheme === 'object') {
-        // 4.1 Merge the selected color scheme to the theme
-        Object.keys(scheme).forEach((schemeKey) => {
-          if (scheme[schemeKey] && typeof scheme[schemeKey] === 'object') {
-            // shallow merge the 1st level structure of the theme.
-            theme[schemeKey] = {
-              ...theme[schemeKey],
-              ...scheme[schemeKey],
-            };
-          } else {
-            theme[schemeKey] = scheme[schemeKey];
-          }
-        });
+      // 4. Resolve the color scheme and merge it to the theme
+      if (calculatedColorScheme) {
+        const scheme = colorSchemes[calculatedColorScheme];
+        if (scheme && typeof scheme === 'object') {
+          // 4.1 Merge the selected color scheme to the theme
+          Object.keys(scheme).forEach((schemeKey) => {
+            if (scheme[schemeKey] && typeof scheme[schemeKey] === 'object') {
+              // shallow merge the 1st level structure of the theme.
+              theme[schemeKey] = {
+                ...theme[schemeKey],
+                ...scheme[schemeKey],
+              };
+            } else {
+              theme[schemeKey] = scheme[schemeKey];
+            }
+          });
+        }
       }
-    }
+
+      return resolveTheme ? resolveTheme(theme) : theme;
+    }, [restThemeProp, colorScheme, components, colorSchemes, cssVarPrefix]);
 
     // 5. Declaring effects
     // 5.1 Updates the selector value to use the current color scheme which tells CSS to use the proper stylesheet.
@@ -248,7 +255,7 @@ export default function createCssVarsProvider(options) {
           process.env.NODE_ENV === 'production'
             ? setMode
             : (newMode) => {
-                if (theme.colorSchemeSelector === 'media') {
+                if (memoTheme.colorSchemeSelector === 'media') {
                   console.error(
                     [
                       'MUI: The `setMode` function has no effect if `colorSchemeSelector` is `media` (`media` is the default value).',
@@ -270,7 +277,7 @@ export default function createCssVarsProvider(options) {
         setColorScheme,
         setMode,
         systemMode,
-        theme.colorSchemeSelector,
+        memoTheme.colorSchemeSelector,
       ],
     );
 
@@ -285,13 +292,12 @@ export default function createCssVarsProvider(options) {
 
     const element = (
       <React.Fragment>
-        <ThemeProvider
-          themeId={scopedTheme ? themeId : undefined}
-          theme={resolveTheme ? resolveTheme(theme) : theme}
-        >
+        <ThemeProvider themeId={scopedTheme ? themeId : undefined} theme={memoTheme}>
           {children}
         </ThemeProvider>
-        {shouldGenerateStyleSheet && <GlobalStyles styles={theme.generateStyleSheets?.() || []} />}
+        {shouldGenerateStyleSheet && (
+          <GlobalStyles styles={memoTheme.generateStyleSheets?.() || []} />
+        )}
       </React.Fragment>
     );