[Box] Cache makeStyles result per props JSON

packages/material-ui-styles/src/getStylesCreator/index.js

@@ -1 +1,2 @@
 export { default } from './getStylesCreator';
+export { default as noopTheme } from './noopTheme'; // private

packages/material-ui/src/Box/Box.js

@@ -1,3 +1,5 @@
+import React, { useEffect, useMemo } from 'react';
+import PropTypes from 'prop-types';
 import {
   borders,
   compose,
@@ -11,7 +13,11 @@ import {
   typography,
   css,
 } from '@material-ui/system';
-import styled from '../styles/styled';
+import { noopTheme } from '@material-ui/styles/getStylesCreator';
+import clsx from 'clsx';
+import { useTheme, makeStyles } from '@material-ui/styles';
+import { chainPropTypes } from '@material-ui/utils';
+import warning from 'warning';
 
 export const styleFunction = css(
   compose(
@@ -27,9 +33,125 @@ export const styleFunction = css(
   ),
 );
 
+function pickAndOmit(input, fields) {
+  const picked = {};
+  const omitted = {};
+
+  Object.keys(input).forEach(prop => {
+    if (fields.indexOf(prop) === -1) {
+      omitted[prop] = input[prop];
+    } else {
+      picked[prop] = input[prop];
+    }
+  });
+
+  return [picked, omitted];
+}
+
+/* eslint-disable react/forbid-foreign-prop-types */
+const { filterProps, propTypes = {} } = styleFunction;
+/* eslint-enable react/forbid-foreign-prop-types */
+
+// Same as multiKeyStore
+const stylesCacheStore = new Map();
+
 /**
  * @ignore - do not document.
  */
-const Box = styled('div')(styleFunction, { name: 'MuiBox' });
+const Box = React.forwardRef(function Box(props, ref) {
+  const { children, className: classNameProp, clone, component: ComponentProp, ...other } = props;
+  const theme = useTheme() || noopTheme;
+
+  const [styleProps, spread] = pickAndOmit(other, filterProps);
+
+  let themeStylesCache = stylesCacheStore.get(theme);
+  if (!themeStylesCache) {
+    themeStylesCache = new Map();
+    stylesCacheStore.set(theme, themeStylesCache);
+  }
+
+  let hasFunc = false;
+  const cacheKey = JSON.stringify(styleProps, (_key, value) => {
+    if (typeof value !== 'function') return value;
+    if (!hasFunc) {
+      warning(false, 'Material-UI: You can not pass a function as Box style attribute.');
+      hasFunc = true;
+    }
+    return 'invalid';
+  });
+
+  const cacheEntry = useMemo(() => {
+    let entry = hasFunc ? undefined : themeStylesCache.get(cacheKey);
+    if (!entry) {
+      const style = styleFunction({ theme, ...other });
+      const useStyles = makeStyles({ root: style }, { classNamePrefix: 'MuiBox' });
+      entry = { useStyles, refs: 0 };
+      if (!hasFunc) themeStylesCache.set(cacheKey, entry);
+    }
+    entry.refs += 1;
+    return entry;
+    // cacheKey depends on 'other' and 'hasFunc'.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [themeStylesCache, cacheKey]);
+
+  useEffect(() => {
+    return () => {
+      cacheEntry.refs -= 1;
+      if (cacheEntry.refs <= 0) {
+        themeStylesCache.delete(cacheKey);
+      }
+    };
+    // Using same deps as cacheEntry.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [themeStylesCache, cacheKey]);
+
+  const classes = cacheEntry.useStyles();
+  const className = clsx(classes.root, classNameProp);
+
+  if (clone) {
+    return React.cloneElement(children, {
+      className: clsx(children.props.className, className),
+    });
+  }
+
+  if (typeof children === 'function') {
+    return children({ className, ...spread });
+  }
+
+  const FinalComponent = ComponentProp || 'div';
+
+  return (
+    <FinalComponent ref={ref} className={className} {...spread}>
+      {children}
+    </FinalComponent>
+  );
+});
+
+Box.propTypes = {
+  /**
+   * A render function or node.
+   */
+  children: PropTypes.oneOfType([PropTypes.node, PropTypes.func]),
+  /**
+   * @ignore
+   */
+  className: PropTypes.string,
+  /**
+   * If `true`, the component will recycle it's children DOM element.
+   * It's using `React.cloneElement` internally.
+   */
+  clone: chainPropTypes(PropTypes.bool, props => {
+    if (props.clone && props.component) {
+      return new Error('You can not use the clone and component prop at the same time.');
+    }
+    return null;
+  }),
+  /**
+   * The component used for the root node.
+   * Either a string to use a DOM element or a component.
+   */
+  component: PropTypes.elementType,
+  ...propTypes,
+};
 
 export default Box;

packages/material-ui/src/Box/Box.test.js

@@ -3,6 +3,7 @@ import { assert } from 'chai';
 import { createMount } from '@material-ui/core/test-utils';
 import describeConformance from '@material-ui/core/test-utils/describeConformance';
 import Box from './Box';
+import consoleErrorMock from 'test/utils/consoleErrorMock';
 
 describe('<Box />', () => {
   let mount;
@@ -53,4 +54,26 @@ describe('<Box />', () => {
     assert.strictEqual(element.getAttribute('font-family'), null);
     assert.strictEqual(element.getAttribute('font-size'), null);
   });
+
+  describe('warnings', () => {
+    beforeEach(() => {
+      consoleErrorMock.spy();
+    });
+
+    afterEach(() => {
+      consoleErrorMock.reset();
+    });
+
+    it('should warn when function is specified for style attribute', () => {
+      mount(
+        <React.Fragment>
+          <Box fontSize={{ xs: 'h6.fontSize', sm: () => 'h4.fontSize', md: () => 'h3.fontSize' }} />
+        </React.Fragment>,
+      );
+      assert.include(
+        consoleErrorMock.args()[0][0],
+        'You can not pass a function as Box style attribute',
+      );
+    });
+  });
 });