fix+feat(select, listbox): bug on dataset with "sections", add support for scrollshadow (#4462)

* fix: add custom function to calculate rowHeight for dataset with sections

* fix: scroll shadow is now working in virtualized components

* chore: add changeset

* fix: to pass test cases use function call instead of function component

Author: vinroger

.changeset/quiet-geese-lay.md

@@ -0,0 +1,7 @@
+---
+"@nextui-org/autocomplete": patch
+"@nextui-org/listbox": patch
+"@nextui-org/select": patch
+---
+
+add support for dataset with section, add support for scrollshadow

packages/components/autocomplete/src/use-autocomplete.ts

@@ -461,6 +461,7 @@ export function useAutocomplete<T extends object>(originalProps: UseAutocomplete
             itemHeight,
           }
         : undefined,
+      scrollShadowProps: slotsProps.scrollShadowProps,
       ...mergeProps(slotsProps.listboxProps, listBoxProps, {
         shouldHighlightOnFocus: true,
       }),

packages/components/listbox/src/virtualized-listbox.tsx

@@ -1,7 +1,10 @@
-import {useRef} from "react";
+import {useMemo, useRef, useState} from "react";
 import {mergeProps} from "@react-aria/utils";
-import {useVirtualizer} from "@tanstack/react-virtual";
+import {useVirtualizer, VirtualItem} from "@tanstack/react-virtual";
 import {isEmpty} from "@nextui-org/shared-utils";
+import {Node} from "@react-types/shared";
+import {ScrollShadowProps, useScrollShadow} from "@nextui-org/scroll-shadow";
+import {filterDOMProps} from "@nextui-org/react-utils";
 
 import ListboxItem from "./listbox-item";
 import ListboxSection from "./listbox-section";
@@ -11,8 +14,50 @@ import {UseListboxReturn} from "./use-listbox";
 interface Props extends UseListboxReturn {
   isVirtualized?: boolean;
   virtualization?: VirtualizationProps;
+  /* Here in virtualized listbox, scroll shadow needs custom implementation. Hence this is the only way to pass props to scroll shadow */
+  scrollShadowProps?: Partial<ScrollShadowProps>;
 }
 
+const getItemSizesForCollection = (collection: Node<object>[], itemHeight: number) => {
+  const sizes: number[] = [];
+
+  for (const item of collection) {
+    if (item.type === "section") {
+      /* +1 for the section header */
+      sizes.push(([...item.childNodes].length + 1) * itemHeight);
+    } else {
+      sizes.push(itemHeight);
+    }
+  }
+
+  return sizes;
+};
+
+const getScrollState = (element: HTMLDivElement | null) => {
+  if (
+    !element ||
+    element.scrollTop === undefined ||
+    element.clientHeight === undefined ||
+    element.scrollHeight === undefined
+  ) {
+    return {
+      isTop: false,
+      isBottom: false,
+      isMiddle: false,
+    };
+  }
+
+  const isAtTop = element.scrollTop === 0;
+  const isAtBottom = Math.ceil(element.scrollTop + element.clientHeight) >= element.scrollHeight;
+  const isInMiddle = !isAtTop && !isAtBottom;
+
+  return {
+    isTop: isAtTop,
+    isBottom: isAtBottom,
+    isMiddle: isInMiddle,
+  };
+};
+
 const VirtualizedListbox = (props: Props) => {
   const {
     Component,
@@ -29,6 +74,7 @@ const VirtualizedListbox = (props: Props) => {
     disableAnimation,
     getEmptyContentProps,
     getListProps,
+    scrollShadowProps,
   } = props;
 
   const {virtualization} = props;
@@ -45,24 +91,29 @@ const VirtualizedListbox = (props: Props) => {
 
   const listHeight = Math.min(maxListboxHeight, itemHeight * state.collection.size);
 
-  const parentRef = useRef(null);
+  const parentRef = useRef<HTMLDivElement>(null);
+  const itemSizes = useMemo(
+    () => getItemSizesForCollection([...state.collection], itemHeight),
+    [state.collection, itemHeight],
+  );
 
   const rowVirtualizer = useVirtualizer({
-    count: state.collection.size,
+    count: [...state.collection].length,
     getScrollElement: () => parentRef.current,
-    estimateSize: () => itemHeight,
+    estimateSize: (i) => itemSizes[i],
   });
 
   const virtualItems = rowVirtualizer.getVirtualItems();
 
-  const renderRow = ({
-    index,
-    style: virtualizerStyle,
-  }: {
-    index: number;
-    style: React.CSSProperties;
-  }) => {
-    const item = [...state.collection][index];
+  /* Here we need the base props for scroll shadow, contains the className (scrollbar-hide and scrollshadow config based on the user inputs on select props) */
+  const {getBaseProps: getBasePropsScrollShadow} = useScrollShadow({...scrollShadowProps});
+
+  const renderRow = (virtualItem: VirtualItem) => {
+    const item = [...state.collection][virtualItem.index];
+
+    if (!item) {
+      return null;
+    }
 
     const itemProps = {
       color,
@@ -74,6 +125,15 @@ const VirtualizedListbox = (props: Props) => {
       ...item.props,
     };
 
+    const virtualizerStyle = {
+      position: "absolute" as const,
+      top: 0,
+      left: 0,
+      width: "100%",
+      height: `${virtualItem.size}px`,
+      transform: `translateY(${virtualItem.start}px)`,
+    };
+
     if (item.type === "section") {
       return (
         <ListboxSection
@@ -102,6 +162,12 @@ const VirtualizedListbox = (props: Props) => {
     return listboxItem;
   };
 
+  const [scrollState, setScrollState] = useState({
+    isTop: false,
+    isBottom: true,
+    isMiddle: false,
+  });
+
   const content = (
     <Component {...getListProps()}>
       {!state.collection.size && !hideEmptyContent && (
@@ -110,11 +176,18 @@ const VirtualizedListbox = (props: Props) => {
         </li>
       )}
       <div
+        {...filterDOMProps(getBasePropsScrollShadow())}
         ref={parentRef}
+        data-bottom-scroll={scrollState.isTop}
+        data-top-bottom-scroll={scrollState.isMiddle}
+        data-top-scroll={scrollState.isBottom}
         style={{
           height: maxListboxHeight,
           overflow: "auto",
         }}
+        onScroll={(e) => {
+          setScrollState(getScrollState(e.target as HTMLDivElement));
+        }}
       >
         {listHeight > 0 && itemHeight > 0 && (
           <div
@@ -124,19 +197,7 @@ const VirtualizedListbox = (props: Props) => {
               position: "relative",
             }}
           >
-            {virtualItems.map((virtualItem) =>
-              renderRow({
-                index: virtualItem.index,
-                style: {
-                  position: "absolute",
-                  top: 0,
-                  left: 0,
-                  width: "100%",
-                  height: `${virtualItem.size}px`,
-                  transform: `translateY(${virtualItem.start}px)`,
-                },
-              }),
-            )}
+            {virtualItems.map((virtualItem) => renderRow(virtualItem))}
           </div>
         )}
       </div>

packages/components/select/src/use-select.ts

@@ -161,7 +161,9 @@ export type UseSelectProps<T> = Omit<
   SelectVariantProps & {
     /**
      * The height of each item in the listbox.
+     * For dataset with sections, the itemHeight must be the height of each item (including padding, border, margin).
      * This is required for virtualized listboxes to calculate the height of each item.
+     * @default 36
      */
     itemHeight?: number;
     /**
@@ -208,7 +210,7 @@ export function useSelect<T extends object>(originalProps: UseSelectProps<T>) {
     onSelectionChange,
     placeholder,
     isVirtualized,
-    itemHeight = 32,
+    itemHeight = 36,
     maxListboxHeight = 256,
     children,
     disallowEmptySelection = false,
@@ -564,6 +566,7 @@ export function useSelect<T extends object>(originalProps: UseSelectProps<T>) {
       className: slots.listbox({
         class: clsx(classNames?.listbox, props?.className),
       }),
+      scrollShadowProps: slotsProps.scrollShadowProps,
       ...mergeProps(slotsProps.listboxProps, props, menuProps),
     } as ListboxProps;
   };

packages/components/select/stories/select.stories.tsx

@@ -1391,6 +1391,104 @@ export const CustomItemHeight = {
   },
 };
 
+const AVATAR_DECORATIONS: {[key: string]: string[]} = {
+  arcane: ["jinx", "atlas-gauntlets", "flame-chompers", "fishbones", "hexcore", "shimmer"],
+  anime: ["cat-ears", "heart-bloom", "in-love", "in-tears", "soul-leaving-body", "starry-eyed"],
+  "lofi-vibes": ["chromawave", "cozy-cat", "cozy-headphones", "doodling", "rainy-mood"],
+  valorant: [
+    "a-hint-of-clove",
+    "blade-storm",
+    "cypher",
+    "frag-out",
+    "omen-cowl",
+    "reyna-leer",
+    "vct-supernova",
+    "viper",
+    "yoru",
+    "carnalito2",
+    "a-hint-of-clove2",
+    "blade-storm2",
+    "cypher2",
+    "frag-out2",
+    "omen-cowl2",
+    "reyna-leer2",
+    "vct-supernova2",
+    "viper2",
+    "yoru2",
+    "carnalito3",
+    "a-hint-of-clove3",
+    "blade-storm3",
+    "cypher3",
+    "frag-out3",
+    "omen-cowl3",
+    "reyna-leer3",
+    "vct-supernova3",
+    "viper3",
+    "yoru3",
+    "carnalito4",
+    "a-hint-of-clove4",
+    "blade-storm4",
+    "cypher4",
+    "frag-out4",
+    "omen-cowl4",
+    "reyna-leer4",
+    "vct-supernova4",
+    "viper4",
+    "yoru4",
+  ],
+  spongebob: [
+    "flower-clouds",
+    "gary-the-snail",
+    "imagination",
+    "musclebob",
+    "sandy-cheeks",
+    "spongebob",
+  ],
+  arcade: ["clyde-invaders", "hot-shot", "joystick", "mallow-jump", "pipedream", "snake"],
+  "street-fighter": ["akuma", "cammy", "chun-li", "guile", "juri", "ken", "m.bison", "ryu"],
+};
+
+export const NonVirtualizedVsVirtualizedWithSections = {
+  render: () => {
+    const SelectComponent = ({isVirtualized}: {isVirtualized: boolean}) => (
+      <Select
+        disallowEmptySelection
+        className="max-w-xs"
+        color="secondary"
+        defaultSelectedKeys={["jinx"]}
+        isVirtualized={isVirtualized}
+        label={`Avatar Decoration ${isVirtualized ? "(Virtualized)" : "(Non-virtualized)"}`}
+        selectedKeys={["jinx"]}
+        selectionMode="single"
+        variant="bordered"
+      >
+        {Object.keys(AVATAR_DECORATIONS).map((key) => (
+          <SelectSection
+            key={key}
+            classNames={{
+              heading: "uppercase text-secondary",
+            }}
+            title={key}
+          >
+            {AVATAR_DECORATIONS[key].map((item) => (
+              <SelectItem key={item} className="capitalize" color="secondary" variant="bordered">
+                {item.replace(/-/g, " ")}
+              </SelectItem>
+            ))}
+          </SelectSection>
+        ))}
+      </Select>
+    );
+
+    return (
+      <div className="flex gap-4 w-full">
+        <SelectComponent isVirtualized={false} />
+        <SelectComponent isVirtualized={true} />
+      </div>
+    );
+  },
+};
+
 export const ValidationBehaviorAria = {
   render: ValidationBehaviorAriaTemplate,
   args: {