POC of Combobox functioning with custom rendered options without significant performance loss

packages/grafana-ui/src/components/Combobox/Combobox.story.tsx

@@ -1,12 +1,13 @@
 import { action } from '@storybook/addon-actions';
 import { useArgs } from '@storybook/preview-api';
 import { Meta, StoryFn, StoryObj } from '@storybook/react';
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 
 import { Field } from '../Forms/Field';
 
 import { Combobox, ComboboxProps } from './Combobox';
 import mdx from './Combobox.mdx';
+import { MENU_ITEM_FONT_SIZE, MENU_ITEM_LINE_HEIGHT, MENU_OPTION_HEIGHT } from './getComboboxStyles';
 import { fakeSearchAPI, generateGroupingOptions, generateOptions } from './storyUtils';
 import { ComboboxOption } from './types';
 
@@ -171,6 +172,65 @@ export const ManyOptions: Story = {
   },
 };
 
+export const NodeOptions: Story = {
+  args: {
+    numberOfOptions: 350,
+    options: undefined,
+    value: undefined,
+  },
+  parameters: {
+    docs: {
+      source: {
+        // necessary to keep storybook from choking on the option generation code
+        type: 'code',
+      },
+    },
+  },
+  render: ({ numberOfOptions, ...args }: PropsAndCustomArgs) => {
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    const [dynamicArgs, setArgs] = useArgs();
+
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    const options: ComboboxOption[] = useMemo(() => {
+      return Array.from({ length: numberOfOptions }, (_, index) => {
+        const label =
+          index % 2 === 0
+            ? {
+                node: (
+                  <>
+                    {`Option ${index}`}
+                    <br />
+                    (Even)
+                  </>
+                ),
+                text: `Option ${index}`,
+                size: MENU_ITEM_FONT_SIZE * MENU_ITEM_LINE_HEIGHT + MENU_OPTION_HEIGHT,
+              }
+            : `Option ${index}`;
+        return {
+          label,
+          value: String(index),
+        };
+      });
+    }, [numberOfOptions]);
+
+    const { onChange, ...rest } = args;
+    return (
+      <Field label="Test input" description="Input with some custom rendered options">
+        <Combobox
+          {...rest}
+          {...dynamicArgs}
+          options={options}
+          onChange={(value: ComboboxOption | null) => {
+            setArgs({ value: value?.value || null });
+            onChangeAction(value);
+          }}
+        />
+      </Field>
+    );
+  },
+};
+
 function loadOptionsWithLabels(inputValue: string) {
   loadOptionsAction(inputValue);
   return fakeSearchAPI(`http://example.com/search?errorOnQuery=break&query=${inputValue}`);

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

@@ -1,7 +1,7 @@
 import { cx } from '@emotion/css';
-import { useVirtualizer, type Range } from '@tanstack/react-virtual';
 import { useCombobox } from 'downshift';
-import React, { ComponentProps, useCallback, useId, useMemo } from 'react';
+import React, { ComponentProps, useId, useMemo } from 'react';
+import { Subject } from 'rxjs';
 
 import { t } from '@grafana/i18n';
 
@@ -14,11 +14,10 @@ import { Portal } from '../Portal/Portal';
 import { ComboboxList } from './ComboboxList';
 import { SuffixIcon } from './SuffixIcon';
 import { itemToString } from './filter';
-import { getComboboxStyles, MENU_OPTION_HEIGHT, MENU_OPTION_HEIGHT_DESCRIPTION } from './getComboboxStyles';
+import { getComboboxStyles } from './getComboboxStyles';
 import { ComboboxOption } from './types';
 import { useComboboxFloat } from './useComboboxFloat';
 import { useOptions } from './useOptions';
-import { isNewGroup } from './utils';
 
 // TODO: It would be great if ComboboxOption["label"] was more generic so that if consumers do pass it in (for async),
 // then the onChange handler emits ComboboxOption with the label as non-undefined.
@@ -154,7 +153,6 @@ export const Combobox = <T extends string | number>(props: ComboboxProps<T>) =>
 
   const {
     options: filteredOptions,
-    groupStartIndices,
     updateOptions,
     asyncLoading,
     asyncError,
@@ -195,49 +193,34 @@ export const Combobox = <T extends string | number>(props: ComboboxProps<T>) =>
 
   const styles = useStyles2(getComboboxStyles);
 
+  // Maybe need to move this to ComboboxList? Is this actually necessary? it's only expanding the returned
+  // array of indexes, and that's already expanded by the overscan setting
   // Injects the group header for the first rendered item into the range to render.
   // Accepts the range that useVirtualizer wants to render, and then returns indexes
   // to actually render.
-  const rangeExtractor = useCallback(
+  // const rangeExtractor = useCallback(
-    (range: Range) => {
+  //   (range: Range) => {
-      const startIndex = Math.max(0, range.startIndex - range.overscan);
+  //     const startIndex = Math.max(0, range.startIndex - range.overscan);
-      const endIndex = Math.min(filteredOptions.length - 1, range.endIndex + range.overscan);
+  //     const endIndex = Math.min(filteredOptions.length - 1, range.endIndex + range.overscan);
-      const rangeToReturn = Array.from({ length: endIndex - startIndex + 1 }, (_, i) => startIndex + i);
+  //     const rangeToReturn = Array.from({ length: endIndex - startIndex + 1 }, (_, i) => startIndex + i);
 
-      // If the first item doesn't have a group, no need to find a header for it
+  //     // If the first item doesn't have a group, no need to find a header for it
-      const firstDisplayedOption = filteredOptions[rangeToReturn[0]];
+  //     const firstDisplayedOption = filteredOptions[rangeToReturn[0]];
-      if (firstDisplayedOption?.group) {
+  //     if (firstDisplayedOption?.group) {
-        const groupStartIndex = groupStartIndices.get(firstDisplayedOption.group);
+  //       const groupStartIndex = groupStartIndices.get(firstDisplayedOption.group);
-        if (groupStartIndex !== undefined && groupStartIndex < rangeToReturn[0]) {
+  //       if (groupStartIndex !== undefined && groupStartIndex < rangeToReturn[0]) {
-          rangeToReturn.unshift(groupStartIndex);
+  //         rangeToReturn.unshift(groupStartIndex);
-        }
+  //       }
-      }
+  //     }
 
-      return rangeToReturn;
+  //     return rangeToReturn;
-    },
+  //   },
-    [filteredOptions, groupStartIndices]
+  //   [filteredOptions, groupStartIndices]
-  );
+  // );
 
-  const rowVirtualizer = useVirtualizer({
+  const scrollToIndexObservable = useMemo(() => {
-    count: filteredOptions.length,
+    return new Subject<number>();
-    getScrollElement: () => scrollRef.current,
+  }, []);
-    estimateSize: (index: number) => {
-      const firstGroupItem = isNewGroup(filteredOptions[index], index > 0 ? filteredOptions[index - 1] : undefined);
-      const hasDescription = 'description' in filteredOptions[index];
-      const hasGroup = 'group' in filteredOptions[index];
-
-      let itemHeight = MENU_OPTION_HEIGHT;
-      if (hasDescription) {
-        itemHeight = MENU_OPTION_HEIGHT_DESCRIPTION;
-      }
-      if (firstGroupItem && hasGroup) {
-        itemHeight += MENU_OPTION_HEIGHT;
-      }
-      return itemHeight;
-    },
-    overscan: VIRTUAL_OVERSCAN_ITEMS,
-    rangeExtractor,
-  });
 
   const {
     isOpen,
@@ -290,7 +273,7 @@ export const Combobox = <T extends string | number>(props: ComboboxProps<T>) =>
 
     onHighlightedIndexChange: ({ highlightedIndex, type }) => {
       if (type !== useCombobox.stateChangeTypes.MenuMouseLeave) {
-        rowVirtualizer.scrollToIndex(highlightedIndex);
+        scrollToIndexObservable.next(highlightedIndex);
       }
     },
     onStateChange: ({ inputValue: newInputValue, type, selectedItem: newSelectedItem }) => {
@@ -418,6 +401,7 @@ export const Combobox = <T extends string | number>(props: ComboboxProps<T>) =>
               scrollRef={scrollRef}
               getItemProps={getItemProps}
               error={asyncError}
+              scrollToIndexObservable={scrollToIndexObservable}
             />
           )}
         </div>

packages/grafana-ui/src/components/Combobox/ComboboxList.tsx

@@ -1,14 +1,15 @@
 import { cx } from '@emotion/css';
 import { useVirtualizer } from '@tanstack/react-virtual';
 import type { UseComboboxPropGetters } from 'downshift';
-import { useCallback } from 'react';
+import { useCallback, useEffect } from 'react';
+import { Subject } from 'rxjs';
 
 import { useStyles2 } from '../../themes/ThemeContext';
 import { Checkbox } from '../Forms/Checkbox';
 import { ScrollContainer } from '../ScrollContainer/ScrollContainer';
 
 import { AsyncError, LoadingOptions, NotFoundError } from './MessageRows';
-import { getComboboxStyles, MENU_OPTION_HEIGHT, MENU_OPTION_HEIGHT_DESCRIPTION } from './getComboboxStyles';
+import { DESCRIPTION_HEIGHT, getComboboxStyles, MENU_OPTION_HEIGHT } from './getComboboxStyles';
 import { ALL_OPTION_VALUE, ComboboxOption } from './types';
 import { isNewGroup } from './utils';
 
@@ -24,6 +25,7 @@ interface ComboboxListProps<T extends string | number> {
   isMultiSelect?: boolean;
   error?: boolean;
   loading?: boolean;
+  scrollToIndexObservable?: Subject<number>;
 }
 
 export const ComboboxList = <T extends string | number>({
@@ -36,6 +38,7 @@ export const ComboboxList = <T extends string | number>({
   isMultiSelect = false,
   error = false,
   loading = false,
+  scrollToIndexObservable,
 }: ComboboxListProps<T>) => {
   const styles = useStyles2(getComboboxStyles);
 
@@ -46,8 +49,11 @@ export const ComboboxList = <T extends string | number>({
       const hasGroup = 'group' in options[index];
 
       let itemHeight = MENU_OPTION_HEIGHT;
+      if (typeof options[index].label === 'object') {
+        itemHeight = options[index].label.size;
+      }
       if (hasDescription) {
-        itemHeight = MENU_OPTION_HEIGHT_DESCRIPTION;
+        itemHeight += DESCRIPTION_HEIGHT;
       }
       if (firstGroupItem && hasGroup) {
         itemHeight += MENU_OPTION_HEIGHT;
@@ -64,6 +70,19 @@ export const ComboboxList = <T extends string | number>({
     overscan: VIRTUAL_OVERSCAN_ITEMS,
   });
 
+  useEffect(() => {
+    if (scrollToIndexObservable) {
+      const subscription = scrollToIndexObservable.subscribe((index) => {
+        rowVirtualizer.scrollToIndex(index);
+      });
+
+      return () => {
+        subscription.unsubscribe();
+      };
+    }
+    return undefined;
+  }, [scrollToIndexObservable, rowVirtualizer]);
+
   const isOptionSelected = useCallback(
     (item: ComboboxOption<T>) => selectedItems.some((opt) => opt.value === item.value),
     [selectedItems]
@@ -90,6 +109,8 @@ export const ComboboxList = <T extends string | number>({
           // the option for aria-describedby.
           const groupHeaderId = groupHeaderItem ? `combobox-option-group-${groupHeaderItem.value}` : undefined;
 
+          const label = typeof item.label === 'object' ? item.label.node : (item.label ?? item.value);
+
           return (
             // Wrapping div should have no styling other than virtual list positioning.
             // It's children (header and option) should appear as flat list items.
@@ -151,7 +172,7 @@ export const ComboboxList = <T extends string | number>({
                 )}
 
                 <div className={styles.optionBody}>
-                  <div className={styles.optionLabel}>{item.label ?? item.value}</div>
+                  <div className={styles.optionLabel}>{label}</div>
 
                   {item.description && <div className={styles.optionDescription}>{item.description}</div>}
                 </div>

packages/grafana-ui/src/components/Combobox/filter.ts

@@ -6,6 +6,11 @@ export function itemToString<T extends string | number>(item?: ComboboxOption<T>
   if (item == null) {
     return '';
   }
+
+  if (typeof item.label === 'object') {
+    return item.label.text;
+  }
+
   return item.label ?? item.value.toString();
 }
 

packages/grafana-ui/src/components/Combobox/getComboboxStyles.ts

@@ -13,8 +13,8 @@ export const MENU_ITEM_LINE_HEIGHT = 1.5;
 
 // Used with Downshift to get the height of each item
 export const MENU_OPTION_HEIGHT = MENU_ITEM_GAP + MENU_ITEM_PADDING * 2 + MENU_ITEM_FONT_SIZE * MENU_ITEM_LINE_HEIGHT;
-export const MENU_OPTION_HEIGHT_DESCRIPTION =
+export const DESCRIPTION_HEIGHT = MENU_ITEM_DESCRIPTION_FONT_SIZE * MENU_ITEM_LINE_HEIGHT;
-  MENU_OPTION_HEIGHT + MENU_ITEM_DESCRIPTION_FONT_SIZE * MENU_ITEM_LINE_HEIGHT;
+export const MENU_OPTION_HEIGHT_DESCRIPTION = MENU_OPTION_HEIGHT + DESCRIPTION_HEIGHT;
 export const POPOVER_MAX_HEIGHT = MENU_OPTION_HEIGHT * 8.5;
 
 export const getComboboxStyles = (theme: GrafanaTheme2) => {

packages/grafana-ui/src/components/Combobox/storyUtils.ts

@@ -1,7 +1,7 @@
-import { ComboboxOption } from './types';
+import { ComboboxStringOption } from './types';
 
-let fakeApiOptions: Array<ComboboxOption<string>>;
+let fakeApiOptions: Array<ComboboxStringOption<string>>;
-export async function fakeSearchAPI(urlString: string): Promise<Array<ComboboxOption<string>>> {
+export async function fakeSearchAPI(urlString: string): Promise<Array<ComboboxStringOption<string>>> {
   const searchParams = new URL(urlString).searchParams;
 
   const errorOnQuery = searchParams.get('errorOnQuery')?.toLowerCase();
@@ -26,19 +26,19 @@ export async function fakeSearchAPI(urlString: string): Promise<Array<ComboboxOp
 
   const delay = searchQuery.length % 2 === 0 ? 200 : 1000;
 
-  return new Promise<Array<ComboboxOption<string>>>((resolve) => {
+  return new Promise<Array<ComboboxStringOption<string>>>((resolve) => {
     setTimeout(() => resolve(filteredOptions), delay);
   });
 }
 
-export async function generateOptions(amount: number): Promise<ComboboxOption[]> {
+export async function generateOptions(amount: number): Promise<ComboboxStringOption[]> {
   return Array.from({ length: amount }, (_, index) => ({
     label: 'Option ' + index,
     value: index.toString(),
   }));
 }
 
-export async function generateGroupingOptions(amount: number): Promise<ComboboxOption[]> {
+export async function generateGroupingOptions(amount: number): Promise<ComboboxStringOption[]> {
   return Array.from({ length: amount }, (_, index) => ({
     label: 'Option ' + index,
     value: index.toString(),

packages/grafana-ui/src/components/Combobox/types.ts

@@ -1,9 +1,28 @@
+import { ReactNode } from 'react';
+
 export const ALL_OPTION_VALUE = '__GRAFANA_INTERNAL_MULTICOMBOBOX_ALL_OPTION__';
 
+interface NodeOption {
+  node: ReactNode;
+  /**
+   * text property much exactly match the text in the `node` ReactNode property
+   */
+  text: string;
+  size: number;
+}
+
 export type ComboboxOption<T extends string | number = string> = {
-  label?: string;
+  label?: string | NodeOption;
   value: T;
   description?: string;
   group?: string;
   infoOption?: boolean;
 };
+
+export interface ComboboxStringOption<T extends string | number = string> {
+  label?: string;
+  value: T;
+  description?: string;
+  group?: string;
+  infoOption?: boolean;
+}

packages/grafana-ui/src/components/Combobox/useComboboxFloat.ts

@@ -67,7 +67,8 @@ export const useComboboxFloat = (items: Array<ComboboxOption<string | number>>,
     const itemsToLookAt = Math.min(items.length, WIDTH_CALCULATION_LIMIT_ITEMS);
 
     for (let i = 0; i < itemsToLookAt; i++) {
-      const itemLabel = items[i].label ?? items[i].value.toString();
+      const label = items[i].label;
+      const itemLabel = (typeof label === 'object' ? label.text : label) ?? items[i].value.toString();
       longestItem = itemLabel.length > longestItem.length ? itemLabel : longestItem;
     }
 

packages/grafana-ui/src/components/Combobox/useMeasureMulti.ts

@@ -30,9 +30,9 @@ export function useMeasureMulti<T extends string | number>(
     let currWidth = 0;
     for (let i = 0; i < selectedItems.length; i++) {
       // Measure text width and add size of padding, separator and close button
-      currWidth +=
+      const item = selectedItems[i];
-        measureText(selectedItems[i].label || '', FONT_SIZE).width +
+      const text = typeof item.label === 'object' ? item.label.text : item.label || '';
-        (disabled ? EXTRA_PILL_DISABLED_SIZE : EXTRA_PILL_SIZE);
+      currWidth += measureText(text, FONT_SIZE).width + (disabled ? EXTRA_PILL_DISABLED_SIZE : EXTRA_PILL_SIZE);
       if (currWidth > maxWidth) {
         // If there is no space for that item, show the current number of items,
         // but always show at least 1 item. Cap at maximum number of items.