fix bugs with subdialog filtering, arrow nav, dividers, etc

Author: LFDanLu

packages/@react-aria/collections/src/BaseCollection.ts

@@ -65,8 +65,17 @@ export class CollectionNode<T> implements Node<T> {
     node.render = this.render;
     node.colSpan = this.colSpan;
     node.colIndex = this.colIndex;
+    node.filter = this.filter;
     return node;
   }
+
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: (textValue: string) => boolean): CollectionNode<T> | null {
+    let [firstKey, lastKey] = filterChildren(collection, newCollection, this.firstChildKey, filterFn);
+    let newNode: Mutable<CollectionNode<T>> = this.clone();
+    newNode.firstChildKey = firstKey;
+    newNode.lastChildKey = lastKey;
+    return newNode;
+  }
 }
 
 /**
@@ -213,141 +222,61 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
     this.frozen = !isSSR;
   }
 
-  // TODO: this is pretty specific to menu, will need to check if it is generic enough
-  // Will need to handle varying levels I assume but will revisit after I get searchable menu working for base menu
-  // TODO: an alternative is to simply walk the collection and add all item nodes that match the filter and any sections/separators we encounter
-  // to an array, then walk that new array and fix all the next/Prev keys while adding them to the new collection
-  UNSTABLE_filter(filterFn: (nodeValue: string) => boolean): BaseCollection<T> {
+  UNSTABLE_filter(filterFn: (textValue: string) => boolean): BaseCollection<T> {
     let newCollection = new BaseCollection<T>();
-    // This tracks the absolute last node we've visited in the collection when filtering, used for setting up the filteredCollection's lastKey and
-    // for updating the next/prevKey for every non-filtered node.
-    let lastNode: Mutable<CollectionNode<T>> | null = null;
-
-    for (let node of this) {
-      if (node.type === 'section' && node.hasChildNodes) {
-        let clonedSection: Mutable<CollectionNode<T>> = (node as CollectionNode<T>).clone();
-        let lastChildInSection: Mutable<CollectionNode<T>> | null = null;
-        for (let child of this.getChildren(node.key)) {
-          if (shouldKeepNode(child, filterFn, this, newCollection)) {
-            let clonedChild: Mutable<CollectionNode<T>> = (child as CollectionNode<T>).clone();
-            // eslint-disable-next-line max-depth
-            if (lastChildInSection == null) {
-              clonedSection.firstChildKey = clonedChild.key;
-            }
-
-            // eslint-disable-next-line max-depth
-            if (newCollection.firstKey == null) {
-              newCollection.firstKey = clonedSection.key;
-            }
-
-            // eslint-disable-next-line max-depth
-            if (lastChildInSection && lastChildInSection.parentKey === clonedChild.parentKey) {
-              lastChildInSection.nextKey = clonedChild.key;
-              clonedChild.prevKey = lastChildInSection.key;
-            } else {
-              clonedChild.prevKey = null;
-            }
-
-            clonedChild.nextKey = null;
-            newCollection.addNode(clonedChild);
-            lastChildInSection = clonedChild;
-          }
-        }
+    let [firstKey, lastKey] = filterChildren(this, newCollection, this.firstKey, filterFn);
+    newCollection.firstKey = firstKey;
+    newCollection.lastKey = lastKey;
+    return newCollection;
+  }
+}
 
-        // Add newly filtered section to collection if it has any valid child nodes, otherwise remove it and its header if any
-        if (lastChildInSection) {
-          if (lastChildInSection.type !== 'header') {
-            clonedSection.lastChildKey = lastChildInSection.key;
-
-            // If the old prev section was filtered out, will need to attach to whatever came before
-            // eslint-disable-next-line max-depth
-            if (lastNode == null) {
-              clonedSection.prevKey = null;
-            } else if (lastNode.type === 'section' || lastNode.type === 'separator') {
-              lastNode.nextKey = clonedSection.key;
-              clonedSection.prevKey = lastNode.key;
-            }
-            clonedSection.nextKey = null;
-            lastNode = clonedSection;
-            newCollection.addNode(clonedSection);
-          } else {
-            if (newCollection.firstKey === clonedSection.key) {
-              newCollection.firstKey = null;
-            }
-            newCollection.removeNode(lastChildInSection.key);
-          }
-        }
-      } else if (node.type === 'separator') {
-        // will need to check if previous section key exists, if it does then we add the separator to the collection.
-        // After the full collection is created we'll need to remove it it is the last node in the section (aka no following section after the separator)
-        let clonedSeparator: Mutable<CollectionNode<T>> = (node as CollectionNode<T>).clone();
-        clonedSeparator.nextKey = null;
-        if (lastNode?.type === 'section') {
-          lastNode.nextKey = clonedSeparator.key;
-          clonedSeparator.prevKey = lastNode.key;
-          lastNode = clonedSeparator;
-          newCollection.addNode(clonedSeparator);
-        }
-      } else {
-        // At this point, the node is either a subdialogtrigger node or a standard row/item
-        let clonedNode: Mutable<CollectionNode<T>> = (node as CollectionNode<T>).clone();
-        if (shouldKeepNode(clonedNode, filterFn, this, newCollection)) {
-          if (newCollection.firstKey == null) {
-            newCollection.firstKey = clonedNode.key;
-          }
-
-          if (lastNode != null) {
-            if (
-              (lastNode.type !== 'section' && lastNode.type !== 'separator' && lastNode.parentKey === clonedNode.parentKey) ||
-              (clonedNode.type === 'loader')
-            ) {
-              lastNode.nextKey = clonedNode.key;
-              clonedNode.prevKey = lastNode.key;
-            } else {
-              clonedNode.prevKey = null;
-            }
-          }
-
-          clonedNode.nextKey = null;
-          newCollection.addNode(clonedNode);
-          lastNode = clonedNode;
-        }
+function filterChildren<T>(collection: BaseCollection<T>, newCollection: BaseCollection<T>, firstChildKey: Key | null, filterFn: (textValue: string) => boolean): [Key | null, Key | null] {
+  // loop over the siblings for firstChildKey
+  // create new nodes based on calling node.filter for each child
+  // if it returns null then don't include it, otherwise update its prev/next keys
+  // add them to the newCollection
+  if (firstChildKey == null) {
+    return [null, null];
+  }
+
+  let firstNode: Node<T> | null = null;
+  let lastNode: Node<T> | null = null;
+  let currentNode = collection.getItem(firstChildKey);
+
+  while (currentNode != null) {
+    let newNode: Mutable<CollectionNode<T>> | null = (currentNode as CollectionNode<T>).filter(collection, newCollection, filterFn);
+    if (newNode != null) {
+      if (lastNode) {
+        newNode.prevKey = lastNode.key;
+        lastNode.nextKey = newNode.key;
       }
-    }
 
-    if (lastNode?.type === 'separator' && lastNode.nextKey === null) {
-      let lastSection;
-      if (lastNode.prevKey != null) {
-        lastSection = newCollection.getItem(lastNode.prevKey) as Mutable<CollectionNode<T>>;
-        lastSection.nextKey = null;
+      if (firstNode == null) {
+        firstNode = newNode;
       }
-      newCollection.removeNode(lastNode.key);
-      lastNode = lastSection;
-    }
 
-    newCollection.lastKey = lastNode?.key || null;
+      newCollection.addNode(newNode);
+      lastNode = newNode;
+    }
 
-    return newCollection;
+    currentNode = currentNode.nextKey ? collection.getItem(currentNode.nextKey) : null;
   }
-}
 
-function shouldKeepNode<T>(node: Node<T>, filterFn: (nodeValue: string) => boolean, oldCollection: BaseCollection<T>, newCollection: BaseCollection<T>): boolean {
-  if (node.type === 'subdialogtrigger' || node.type === 'submenutrigger') {
-    // Subdialog wrapper should only have one child, if it passes the filter add it to the new collection since we don't need to
-    // do any extra handling for its first/next key
-    let triggerChild = [...oldCollection.getChildren(node.key)][0];
-    if (triggerChild && filterFn(triggerChild.textValue)) {
-      let clonedChild: Mutable<CollectionNode<T>> = (triggerChild as CollectionNode<T>).clone();
-      newCollection.addNode(clonedChild);
-      return true;
+  // TODO: this is pretty specific to dividers but doesn't feel like there is a good way to get around it since we only can know
+  // to filter the last separator in a collection only after performing a filter for the rest of the contents after it
+  // Its gross that it needs to live here, might be nice if somehow we could have this live in the separator code
+  if (lastNode && lastNode.type === 'separator') {
+    let prevKey = lastNode.prevKey;
+    newCollection.removeNode(lastNode.key);
+
+    if (prevKey) {
+      lastNode = newCollection.getItem(prevKey) as Mutable<CollectionNode<T>>;
+      lastNode.nextKey = null;
     } else {
-      return false;
+      lastNode = null;
     }
-  } else if (node.type === 'header' || node.type === 'loader') {
-    // TODO what about tree multiple loaders? Should a loader still be preserved if its parent row is filtered out
-    // Actually how should a tree structure be filtered? Do levels no longer matter or does it filter at each level?
-    return true;
-  } else {
-    return filterFn(node.textValue);
   }
+
+  return [firstNode?.key ?? null, lastNode?.key ?? null];
 }

packages/react-aria-components/src/Header.tsx

@@ -21,6 +21,10 @@ class HeaderNode extends CollectionNode<any> {
   constructor(key: Key) {
     super('header', key);
   }
+
+  filter(): CollectionNode<any> {
+    return this.clone();
+  }
 }
 
 export const Header = /*#__PURE__*/ createLeafComponent(HeaderNode, function Header(props: HTMLAttributes<HTMLElement>, ref: ForwardedRef<HTMLElement>) {

packages/react-aria-components/src/ListBox.tsx

@@ -11,7 +11,7 @@
  */
 
 import {AriaListBoxOptions, AriaListBoxProps, DraggableItemResult, DragPreviewRenderer, DroppableCollectionResult, DroppableItemResult, FocusScope, ListKeyboardDelegate, mergeProps, useCollator, useFocusRing, useHover, useListBox, useListBoxSection, useLocale, useOption} from 'react-aria';
-import {Collection, CollectionBuilder, CollectionNode, createBranchComponent, createLeafComponent} from '@react-aria/collections';
+import {BaseCollection, Collection, CollectionBuilder, CollectionNode, createBranchComponent, createLeafComponent} from '@react-aria/collections';
 import {CollectionProps, CollectionRendererContext, ItemRenderProps, SectionContext, SectionProps} from './Collection';
 import {ContextValue, DEFAULT_SLOT, Provider, RenderProps, ScrollableProps, SlotProps, StyleProps, StyleRenderProps, useContextProps, useRenderProps, useSlot} from './utils';
 import {DragAndDropContext, DropIndicatorContext, DropIndicatorProps, useDndPersistedKeys, useRenderDropIndicator} from './DragAndDrop';
@@ -307,10 +307,24 @@ function ListBoxSectionInner<T extends object>(props: ListBoxSectionProps<T>, re
 
 // todo make a class here
 
-class SectionNode extends CollectionNode<any> {
+class SectionNode<T> extends CollectionNode<T> {
   constructor(key: Key) {
     super('section', key);
   }
+
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: (textValue: string) => boolean): CollectionNode<T> | null {
+    let filteredSection = super.filter(collection, newCollection, filterFn);
+    if (filteredSection) {
+      if (filteredSection.lastChildKey !== null) {
+        let lastChild = collection.getItem(filteredSection.lastChildKey);
+        if (lastChild && lastChild.type !== 'header') {
+          return filteredSection;
+        }
+      }
+    }
+
+    return null;
+  }
 }
 
 /**
@@ -341,10 +355,18 @@ export interface ListBoxItemProps<T = object> extends RenderProps<ListBoxItemRen
 
 // TODO create item type here
 
-class ItemNode extends CollectionNode<any> {
+class ItemNode<T> extends CollectionNode<T> {
   constructor(key: Key) {
     super('item', key);
   }
+
+  filter(_, __, filterFn: (textValue: string) => boolean): CollectionNode<T> | null {
+    if (filterFn(this.textValue)) {
+      return this.clone();
+    }
+
+    return null;
+  }
 }
 
 /**
@@ -489,6 +511,10 @@ class LoaderNode extends CollectionNode<any> {
   constructor(key: Key) {
     super('loader', key);
   }
+
+  filter(): CollectionNode<any> | null {
+    return this.clone();
+  }
 }
 
 const ListBoxDropIndicatorForwardRef = forwardRef(ListBoxDropIndicator);

packages/react-aria-components/src/Menu.tsx

@@ -107,11 +107,22 @@ export interface SubmenuTriggerProps {
 
 const SubmenuTriggerContext = createContext<{parentMenuRef: RefObject<HTMLElement | null>, shouldUseVirtualFocus?: boolean} | null>(null);
 
-// todo: what logic should this have?
-class SubMenuTriggerNode extends CollectionNode<any> {
+class SubMenuTriggerNode<T> extends CollectionNode<T> {
   constructor(key: Key) {
     super('submenutrigger', key);
   }
+
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: (textValue: string) => boolean): CollectionNode<T> | null {
+    let triggerNode = collection.getItem(this.firstChildKey!);
+    if (triggerNode && filterFn(triggerNode.textValue)) {
+      // TODO: perhaps should call super.filter for correctness, but basically add the menu item child of the submenutrigger
+      // to the keymap so it renders
+      newCollection.addNode(triggerNode as CollectionNode<T>);
+      return this.clone();
+    }
+
+    return null;
+  }
 }
 
 /**
@@ -326,10 +337,24 @@ function MenuSectionInner<T extends object>(props: MenuSectionProps<T>, ref: For
 }
 
 // todo can probably reuse the SectionNode from ListBox?
-class SectionNode extends CollectionNode<any> {
+class SectionNode<T> extends CollectionNode<T> {
   constructor(key: Key) {
     super('section', key);
   }
+
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: (textValue: string) => boolean): CollectionNode<T> | null {
+    let filteredSection = super.filter(collection, newCollection, filterFn);
+    if (filteredSection) {
+      if (filteredSection.lastChildKey !== null) {
+        let lastChild = collection.getItem(filteredSection.lastChildKey);
+        if (lastChild && lastChild.type !== 'header') {
+          return filteredSection;
+        }
+      }
+    }
+
+    return null;
+  }
 }
 
 /**
@@ -370,10 +395,18 @@ export interface MenuItemProps<T = object> extends RenderProps<MenuItemRenderPro
 const MenuItemContext = createContext<ContextValue<MenuItemProps, HTMLDivElement>>(null);
 
 // TODO maybe this needs to a separate node type?
-class MenuItemNode extends CollectionNode<any> {
+class MenuItemNode<T> extends CollectionNode<T> {
   constructor(key: Key) {
     super('item', key);
   }
+
+  filter(_, __, filterFn: (textValue: string) => boolean): CollectionNode<T> | null {
+    if (filterFn(this.textValue)) {
+      return this.clone();
+    }
+
+    return null;
+  }
 }
 
 /**

packages/react-aria-components/src/Separator.tsx

@@ -11,7 +11,7 @@
  */
 
 import {SeparatorProps as AriaSeparatorProps, useSeparator} from 'react-aria';
-import {CollectionNode, createLeafComponent} from '@react-aria/collections';
+import {BaseCollection, CollectionNode, createLeafComponent} from '@react-aria/collections';
 import {ContextValue, SlotProps, StyleProps, useContextProps} from './utils';
 import {filterDOMProps} from '@react-aria/utils';
 import {Key} from '@react-types/shared';
@@ -25,6 +25,14 @@ class SeparatorNode extends CollectionNode<any> {
   constructor(key: Key) {
     super('separator', key);
   }
+
+  filter(_, newCollection: BaseCollection<any>): CollectionNode<any> | null {
+    if (newCollection.getItem(this.prevKey!)) {
+      return this.clone();
+    }
+
+    return null;
+  }
 }
 
 export const Separator = /*#__PURE__*/ createLeafComponent(SeparatorNode, function Separator(props: SeparatorProps, ref: ForwardedRef<HTMLElement>) {

packages/react-aria-components/stories/Autocomplete.stories.tsx

@@ -13,11 +13,11 @@
 import {action} from '@storybook/addon-actions';
 import {Autocomplete, Button, Collection, DialogTrigger, Header, Input, Keyboard, Label, ListBox, ListBoxSection, ListLayout, Menu, MenuItem, MenuSection, MenuTrigger, Popover, SearchField, Select, SelectValue, Separator, SubmenuTrigger, Text, TextField, Virtualizer} from 'react-aria-components';
 import {MyListBoxItem, MyMenuItem} from './utils';
+import {MyListBoxLoaderIndicator, renderEmptyState} from './ListBox.stories';
 import React from 'react';
 import styles from '../example/index.css';
 import {useAsyncList, useListData, useTreeData} from 'react-stately';
 import {useFilter} from 'react-aria';
-import { MyListBoxLoaderIndicator, renderEmptyState } from './ListBox.stories';
 
 export default {
   title: 'React Aria Components',