[styles] Use memoized context in StylesProvider (mui#34637)

Co-authored-by: Marija Najdova <[email protected]>
jesrodri · Dec 6, 2022 · d8d5585 · d8d5585
1 parent 6ba0574
commit d8d5585
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 68 deletions.
diff --git a/packages/mui-styles/src/StylesProvider/StylesProvider.js b/packages/mui-styles/src/StylesProvider/StylesProvider.js
@@ -6,23 +6,24 @@ import createGenerateClassName from '../createGenerateClassName';
 import jssPreset from '../jssPreset';
 
 // Default JSS instance.
-const jss = create(jssPreset());
+const defaultJSS = create(jssPreset());
 
 // Use a singleton or the provided one by the context.
 //
 // The counter-based approach doesn't tolerate any mistake.
 // It's much safer to use the same counter everywhere.
-const generateClassName = createGenerateClassName();
+const defaultGenerateClassName = createGenerateClassName();
 
+const defaultSheetsManager = new Map();
 // Exported for test purposes
-export const sheetsManager = new Map();
+export { defaultSheetsManager as sheetsManager };
 
 const defaultOptions = {
   disableGeneration: false,
-  generateClassName,
-  jss,
+  generateClassName: defaultGenerateClassName,
+  jss: defaultJSS,
   sheetsCache: null,
-  sheetsManager,
+  sheetsManager: defaultSheetsManager,
   sheetsRegistry: null,
 };
 
@@ -38,39 +39,71 @@ export default function StylesProvider(props) {
   const { children, injectFirst = false, disableGeneration = false, ...localOptions } = props;
 
   const outerOptions = React.useContext(StylesContext);
-  const context = { ...outerOptions, disableGeneration, ...localOptions };
+  const {
+    generateClassName,
+    jss,
+    serverGenerateClassName,
+    sheetsCache,
+    sheetsManager,
+    sheetsRegistry,
+  } = { ...outerOptions, ...localOptions };
 
   if (process.env.NODE_ENV !== 'production') {
-    if (typeof window === 'undefined' && !context.sheetsManager) {
-      console.error('MUI: You need to use the ServerStyleSheets API when rendering on the server.');
+    if (injectFirst && localOptions.jss) {
+      console.error('MUI: You cannot use the jss and injectFirst props at the same time.');
     }
   }
 
-  if (process.env.NODE_ENV !== 'production') {
-    if (context.jss.options.insertionPoint && injectFirst) {
-      console.error(
-        'MUI: You cannot use a custom insertionPoint and <StylesContext injectFirst> at the same time.',
-      );
+  const value = React.useMemo(() => {
+    const context = {
+      disableGeneration,
+      generateClassName,
+      jss,
+      serverGenerateClassName,
+      sheetsCache,
+      sheetsManager,
+      sheetsRegistry,
+    };
+
+    if (process.env.NODE_ENV !== 'production') {
+      if (typeof window === 'undefined' && !context.sheetsManager) {
+        console.error(
+          'MUI: You need to use the ServerStyleSheets API when rendering on the server.',
+        );
+      }
     }
-  }
 
-  if (process.env.NODE_ENV !== 'production') {
-    if (injectFirst && localOptions.jss) {
-      console.error('MUI: You cannot use the jss and injectFirst props at the same time.');
+    if (process.env.NODE_ENV !== 'production') {
+      if (context.jss.options.insertionPoint && injectFirst) {
+        console.error(
+          'MUI: You cannot use a custom insertionPoint and <StylesContext injectFirst> at the same time.',
+        );
+      }
     }
-  }
 
-  if (!context.jss.options.insertionPoint && injectFirst && typeof window !== 'undefined') {
-    if (!injectFirstNode) {
-      const head = document.head;
-      injectFirstNode = document.createComment('mui-inject-first');
-      head.insertBefore(injectFirstNode, head.firstChild);
-    }
+    if (!context.jss.options.insertionPoint && injectFirst && typeof window !== 'undefined') {
+      if (!injectFirstNode) {
+        const head = document.head;
+        injectFirstNode = document.createComment('mui-inject-first');
+        head.insertBefore(injectFirstNode, head.firstChild);
+      }
 
-    context.jss = create({ plugins: jssPreset().plugins, insertionPoint: injectFirstNode });
-  }
+      context.jss = create({ plugins: jssPreset().plugins, insertionPoint: injectFirstNode });
+    }
 
-  return <StylesContext.Provider value={context}>{children}</StylesContext.Provider>;
+    return context;
+  }, [
+    injectFirst,
+    disableGeneration,
+    generateClassName,
+    jss,
+    serverGenerateClassName,
+    sheetsCache,
+    sheetsManager,
+    sheetsRegistry,
+  ]);
+
+  return <StylesContext.Provider value={value}>{children}</StylesContext.Provider>;
 }
 
 StylesProvider.propTypes = {

diff --git a/packages/mui-styles/src/makeStyles/makeStyles.test.js b/packages/mui-styles/src/makeStyles/makeStyles.test.js
@@ -415,46 +415,6 @@ describe('makeStyles', () => {
     });
   });
 
-  describe('react-hot-loader', () => {
-    it('should take the new stylesCreator into account', () => {
-      const useStyles1 = makeStyles({ root: { padding: 8 } });
-      const useStyles2 = makeStyles({ root: { padding: 4 } });
-
-      let hmr = false;
-
-      const StyledComponent = () => {
-        // Simulate react-hot-loader behavior
-        /* eslint-disable react-hooks/rules-of-hooks */
-        if (hmr) {
-          useStyles2();
-        } else {
-          useStyles1();
-        }
-        /* eslint-enable react-hooks/rules-of-hooks */
-
-        return <div />;
-      };
-
-      const sheetsRegistry = new SheetsRegistry();
-      const wrapper = mount(
-        <StylesProvider sheetsRegistry={sheetsRegistry} sheetsCache={new Map()}>
-          <StyledComponent />
-        </StylesProvider>,
-      );
-      expect(sheetsRegistry.registry.length).to.equal(1);
-      expect(sheetsRegistry.registry[0].rules.raw).to.deep.equal({
-        root: { padding: 8 },
-      });
-
-      hmr = true;
-      wrapper.setProps({});
-      expect(sheetsRegistry.registry.length).to.equal(1);
-      expect(sheetsRegistry.registry[0].rules.raw).to.deep.equal({
-        root: { padding: 4 },
-      });
-    });
-  });
-
   describe('classname quality', () => {
     let sheetsRegistry;