GrafanaUI: Support memoization of useStyles additional arguments (#75000)

* GrafanaUI: Support memoisation of useStyles additional arguments

* remove spooky any

packages/grafana-ui/package.json

@@ -77,6 +77,7 @@
     "jquery": "3.7.0",
     "lodash": "4.17.21",
     "memoize-one": "6.0.0",
+    "micro-memoize": "^4.1.2",
     "moment": "2.29.4",
     "monaco-editor": "0.34.0",
     "ol": "7.4.0",

packages/grafana-ui/src/components/Badge/Badge.tsx

@@ -1,5 +1,5 @@
 import { css, cx } from '@emotion/css';
-import React, { HTMLAttributes, useCallback } from 'react';
+import React, { HTMLAttributes } from 'react';
 import tinycolor from 'tinycolor2';
 
 import { GrafanaTheme2 } from '@grafana/data';
@@ -19,7 +19,7 @@ export interface BadgeProps extends HTMLAttributes<HTMLDivElement> {
 }
 
 export const Badge = React.memo<BadgeProps>(({ icon, color, text, tooltip, className, ...otherProps }) => {
-  const styles = useStyles2(useCallback((theme) => getStyles(theme, color), [color]));
+  const styles = useStyles2(getStyles, color);
   const badge = (
     <div className={cx(styles.wrapper, className)} {...otherProps}>
       {icon && <Icon name={icon} size="sm" />}

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

@@ -2,16 +2,36 @@ import { css } from '@emotion/css';
 import { render, renderHook } from '@testing-library/react';
 import React from 'react';
 
+import { GrafanaTheme2 } from '@grafana/data';
+
 import { mockThemeContext, useStyles2 } from './ThemeContext';
 
 describe('useStyles', () => {
   it('memoizes the passed in function correctly', () => {
-    const stylesCreator = () => ({});
-    const { rerender, result } = renderHook(() => useStyles2(stylesCreator));
-    const storedReference = result.current;
+    // implementation has extra arguments to implicitly test the typescript definition of useStyles2
+    const getStyles = jest.fn((theme: GrafanaTheme2, isOdd: boolean) => ({ row: 'row-class-name' }));
 
-    rerender();
-    expect(storedReference).toBe(result.current);
+    function Row({ isOdd }: { isOdd: boolean }) {
+      const styles = useStyles2(getStyles, isOdd);
+      return <div className={styles.row} />;
+    }
+
+    function TestUseStyles() {
+      return (
+        <>
+          <Row isOdd={true} />
+          <Row isOdd={false} />
+          <Row isOdd={true} />
+          <Row isOdd={false} />
+          <Row isOdd={true} />
+          <Row isOdd={false} />
+        </>
+      );
+    }
+
+    render(<TestUseStyles />);
+
+    expect(getStyles).toHaveBeenCalledTimes(2);
   });
 
   it('does not memoize if the passed in function changes every time', () => {

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

@@ -1,4 +1,5 @@
 import hoistNonReactStatics from 'hoist-non-react-statics';
+import memoize from 'micro-memoize';
 import React, { useContext } from 'react';
 
 import { createTheme, GrafanaTheme, GrafanaTheme2 } from '@grafana/data';
@@ -89,6 +90,7 @@ export function useStyles<T>(getStyles: (theme: GrafanaTheme) => T) {
   const theme = useTheme();
 
   let memoizedStyleCreator: typeof getStyles = memoizedStyleCreators.get(getStyles);
+
   if (!memoizedStyleCreator) {
     memoizedStyleCreator = stylesFactory(getStyles);
     memoizedStyleCreators.set(getStyles, memoizedStyleCreator);
@@ -98,27 +100,38 @@ export function useStyles<T>(getStyles: (theme: GrafanaTheme) => T) {
 }
 
 /**
- * Hook for using memoized styles with access to the theme.
+ * Hook for using memoized styles with access to the theme. Pass additional
+ * arguments to the getStyles function as additional arguments to this hook.
  *
- * 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()`.)
+ * Prefer using primitive values (boolean, number, string, etc) for
+ * additional arguments for better performance
+ *
+ * const getStyles = (theme, isDisabled, isOdd) => {css(...)}
+ * [...]
+ * const styles = useStyles2(getStyles, true, Boolean(index % 2))
+ *
+ * NOTE: For memoization to work, ensure that all arguments don't change
+ * across renders (or only change if they need to)
  * */
 /** @public */
-export function useStyles2<T>(getStyles: (theme: GrafanaTheme2) => T) {
+export function useStyles2<T extends unknown[], CSSReturnValue>(
+  getStyles: (theme: GrafanaTheme2, ...args: T) => CSSReturnValue,
+  ...additionalArguments: T
+): CSSReturnValue {
   const theme = useTheme2();
 
   let memoizedStyleCreator: typeof getStyles = memoizedStyleCreators.get(getStyles);
+
   if (!memoizedStyleCreator) {
-    memoizedStyleCreator = stylesFactory(getStyles);
+    memoizedStyleCreator = memoize(getStyles, { maxSize: 10 }); // each getStyles function will memoize 10 different sets of props
     memoizedStyleCreators.set(getStyles, memoizedStyleCreator);
   }
 
-  return memoizedStyleCreator(theme);
+  return memoizedStyleCreator(theme, ...additionalArguments);
 }
 
 /**
- * Enables theme context  mocking
+ * Enables theme context mocking
  */
 /** @public */
 export const mockThemeContext = (theme: Partial<GrafanaTheme2>) => {

yarn.lock

@@ -4151,6 +4151,7 @@ __metadata:
     jquery: 3.7.0
     lodash: 4.17.21
     memoize-one: 6.0.0
+    micro-memoize: ^4.1.2
     mock-raf: 1.0.1
     moment: 2.29.4
     monaco-editor: 0.34.0
@@ -24412,6 +24413,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"micro-memoize@npm:^4.1.2":
+  version: 4.1.2
+  resolution: "micro-memoize@npm:4.1.2"
+  checksum: 4b02750622d44b5ab31573c629b5d91927dd0c2727743ff75e790c223ab6cd02c48cc3bddea69da0dffb688091a0a71a17944947dd165f8ba9e03728bc30a76d
+  languageName: node
+  linkType: hard
+
 "micromatch@npm:^4.0.0, micromatch@npm:^4.0.2, micromatch@npm:^4.0.4":
   version: 4.0.4
   resolution: "micromatch@npm:4.0.4"