ThemeContext: Fix useStyles memoization (#26200)

packages/grafana-ui/src/themes/ThemeContext.test.tsx

 import React from 'react';
 import { config } from '@grafana/runtime';
+import { renderHook } from '@testing-library/react-hooks';
 import { css } from 'emotion';
 import { mount } from 'enzyme';
-import { useStyles } from './ThemeContext';
+import { memoizedStyleCreators, mockThemeContext, useStyles } from './ThemeContext';
 
 describe('useStyles', () => {
-  it('passes in theme and returns style object', () => {
+  it('memoizes the passed in function correctly', () => {
+    const stylesCreator = () => ({});
+    const { rerender, result } = renderHook(() => useStyles(stylesCreator));
+    const storedReference = result.current;
+
+    rerender();
+    expect(storedReference).toBe(result.current);
+  });
+
+  it('does not memoize if the passed in function changes every time', () => {
+    const { rerender, result } = renderHook(() => useStyles(() => ({})));
+    const storedReference = result.current;
+    rerender();
+    expect(storedReference).not.toBe(result.current);
+  });
+
+  it('updates the memoized function when the theme changes', () => {
+    const stylesCreator = () => ({});
+    const { rerender, result } = renderHook(() => useStyles(stylesCreator));
+    const storedReference = result.current;
+
+    const restoreThemeContext = mockThemeContext({});
+    rerender();
+    expect(storedReference).not.toBe(result.current);
+    restoreThemeContext();
+  });
+
+  it('cleans up memoized functions whenever a new one comes along or the component unmounts', () => {
+    const styleCreators: Function[] = [];
+    const { rerender, unmount } = renderHook(() => {
+      const styleCreator = () => ({});
+      styleCreators.push(styleCreator);
+      return useStyles(styleCreator);
+    });
+
+    expect(typeof memoizedStyleCreators.get(styleCreators[0])).toBe('function');
+    rerender();
+    expect(memoizedStyleCreators.get(styleCreators[0])).toBeUndefined();
+    expect(typeof memoizedStyleCreators.get(styleCreators[1])).toBe('function');
+    unmount();
+    expect(memoizedStyleCreators.get(styleCreators[0])).toBeUndefined();
+    expect(memoizedStyleCreators.get(styleCreators[1])).toBeUndefined();
+  });
+
+  it('passes in theme and returns style object', done => {
     const Dummy: React.FC = function() {
       const styles = useStyles(theme => {
         expect(theme).toEqual(config.theme);
 
         return {
           someStyle: css`
-            color: ${theme?.palette.critical};
+            color: ${theme.palette.critical};
           `,
         };
       });
 
       expect(typeof styles.someStyle).toBe('string');
+      done();
 
       return <div>dummy</div>;
     };

packages/grafana-ui/src/themes/ThemeContext.tsx

-import React, { useContext } from 'react';
+import { GrafanaTheme, GrafanaThemeType } from '@grafana/data';
 import hoistNonReactStatics from 'hoist-non-react-statics';
-
-import { getTheme } from './getTheme';
+import React, { useContext, useEffect } from 'react';
 import { Themeable } from '../types/theme';
-import { GrafanaTheme, GrafanaThemeType } from '@grafana/data';
+import { getTheme } from './getTheme';
 import { stylesFactory } from './stylesFactory';
 
 type Omit<T, K> = Pick<T, Exclude<keyof T, K>>;
@@ -14,6 +13,9 @@ type Subtract<T, K> = Omit<T, keyof K>;
  */
 let ThemeContextMock: React.Context<GrafanaTheme> | null = null;
 
+// Used by useStyles()
+export const memoizedStyleCreators = new WeakMap();
+
 // Use Grafana Dark theme by default
 export const ThemeContext = React.createContext(getTheme(GrafanaThemeType.Dark));
 ThemeContext.displayName = 'ThemeContext';
@@ -39,9 +41,29 @@ export function useTheme(): GrafanaTheme {
   return useContext(ThemeContextMock || ThemeContext);
 }
 
-/** Hook for using memoized styles with access to the theme. */
+/**
+ * Hook for using memoized styles with access to the theme.
+ *
+ * NOTE: For memoization to work, you need to ensure that the function
+ * you pass in doesn't change, or only if it needs to. (i.e. declare
+ * your style creator outside of a function component or use `useCallback()`.)
+ * */
 export function useStyles<T>(getStyles: (theme: GrafanaTheme) => T) {
-  return stylesFactory(getStyles)(useTheme());
+  const theme = useTheme();
+
+  let memoizedStyleCreator = memoizedStyleCreators.get(getStyles);
+  if (!memoizedStyleCreator) {
+    memoizedStyleCreator = stylesFactory(getStyles);
+    memoizedStyleCreators.set(getStyles, memoizedStyleCreator);
+  }
+
+  useEffect(() => {
+    return () => {
+      memoizedStyleCreators.delete(getStyles);
+    };
+  }, [getStyles]);
+
+  return memoizedStyleCreator(theme);
 }
 
 /**