feat: (React Aria) Implement filtering on a per CollectionNode basis

packages/@react-aria/autocomplete/src/useAutocomplete.ts

@@ -10,10 +10,10 @@
  * governing permissions and limitations under the License.
  */
 
-import {AriaLabelingProps, BaseEvent, DOMProps, RefObject} from '@react-types/shared';
+import {AriaLabelingProps, BaseEvent, DOMProps, Node, RefObject} from '@react-types/shared';
 import {AriaTextFieldProps} from '@react-aria/textfield';
 import {AutocompleteProps, AutocompleteState} from '@react-stately/autocomplete';
-import {CLEAR_FOCUS_EVENT, FOCUS_EVENT, getActiveElement, getOwnerDocument, isCtrlKeyPressed, mergeProps, mergeRefs, useEffectEvent, useEvent, useId, useLabels, useObjectRef} from '@react-aria/utils';
+import {CLEAR_FOCUS_EVENT, FOCUS_EVENT, getActiveElement, getOwnerDocument, isCtrlKeyPressed, mergeProps, mergeRefs, useEffectEvent, useEvent, useLabels, useObjectRef, useSlotId} from '@react-aria/utils';
 import {dispatchVirtualBlur, dispatchVirtualFocus, getVirtuallyFocusedElement, moveVirtualFocus} from '@react-aria/focus';
 import {getInteractionModality} from '@react-aria/interactions';
 // @ts-ignore
@@ -27,12 +27,14 @@ export interface CollectionOptions extends DOMProps, AriaLabelingProps {
   /** Whether typeahead is disabled. */
   disallowTypeAhead: boolean
 }
-export interface AriaAutocompleteProps extends AutocompleteProps {
+
+// TODO; For now go with Node here, but maybe pare it down to just the essentials? Value, key, and maybe type?
+export interface AriaAutocompleteProps<T> extends AutocompleteProps {
   /**
    * An optional filter function used to determine if a option should be included in the autocomplete list.
    * Include this if the items you are providing to your wrapped collection aren't filtered by default.
    */
-  filter?: (textValue: string, inputValue: string) => boolean,
+  filter?: (textValue: string, inputValue: string, node: Node<T>) => boolean,
 
   /**
    * Whether or not to focus the first item in the collection after a filter is performed.
@@ -41,22 +43,22 @@ export interface AriaAutocompleteProps extends AutocompleteProps {
   disableAutoFocusFirst?: boolean
 }
 
-export interface AriaAutocompleteOptions extends Omit<AriaAutocompleteProps, 'children'> {
+export interface AriaAutocompleteOptions<T> extends Omit<AriaAutocompleteProps<T>, 'children'> {
   /** The ref for the wrapped collection element. */
   inputRef: RefObject<HTMLInputElement | null>,
   /** The ref for the wrapped collection element. */
   collectionRef: RefObject<HTMLElement | null>
 }
 
-export interface AutocompleteAria {
+export interface AutocompleteAria<T> {
   /** Props for the autocomplete textfield/searchfield element. These should be passed to the textfield/searchfield aria hooks respectively. */
   textFieldProps: AriaTextFieldProps,
   /** Props for the collection, to be passed to collection's respective aria hook (e.g. useMenu). */
   collectionProps: CollectionOptions,
   /** Ref to attach to the wrapped collection. */
   collectionRef: RefObject<HTMLElement | null>,
   /** A filter function that returns if the provided collection node should be filtered out of the collection. */
-  filter?: (nodeTextValue: string) => boolean
+  filter?: (nodeTextValue: string, node: Node<T>) => boolean
 }
 
 /**
@@ -65,15 +67,15 @@ export interface AutocompleteAria {
  * @param props - Props for the autocomplete.
  * @param state - State for the autocomplete, as returned by `useAutocompleteState`.
  */
-export function useAutocomplete(props: AriaAutocompleteOptions, state: AutocompleteState): AutocompleteAria {
+export function useAutocomplete<T>(props: AriaAutocompleteOptions<T>, state: AutocompleteState): AutocompleteAria<T> {
   let {
     inputRef,
     collectionRef,
     filter,
     disableAutoFocusFirst = false
   } = props;
 
-  let collectionId = useId();
+  let collectionId = useSlotId();
   let timeout = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
   let delayNextActiveDescendant = useRef(false);
   let queuedActiveDescendant = useRef<string | null>(null);
@@ -316,9 +318,9 @@ export function useAutocomplete(props: AriaAutocompleteOptions, state: Autocompl
     'aria-label': stringFormatter.format('collectionLabel')
   });
 
-  let filterFn = useCallback((nodeTextValue: string) => {
+  let filterFn = useCallback((nodeTextValue: string, node: Node<T>) => {
     if (filter) {
-      return filter(nodeTextValue, state.inputValue);
+      return filter(nodeTextValue, state.inputValue, node);
     }
 
     return true;
@@ -352,13 +354,19 @@ export function useAutocomplete(props: AriaAutocompleteOptions, state: Autocompl
     }
   };
 
-  return {
-    textFieldProps: {
-      value: state.inputValue,
-      onChange,
+  // Only apply the autocomplete specific behaviors if the collection component wrapped by it is actually
+  // being filtered/allows filtering by the Autocomplete.
+  let textFieldProps = {
+    value: state.inputValue,
+    onChange
+  } as AriaTextFieldProps<HTMLInputElement>;
+
+  if (collectionId) {
+    textFieldProps = {
+      ...textFieldProps,
       onKeyDown,
       autoComplete: 'off',
-      'aria-haspopup': 'listbox',
+      'aria-haspopup': collectionId ? 'listbox' : undefined,
       'aria-controls': collectionId,
       // TODO: readd proper logic for completionMode = complete (aria-autocomplete: both)
       'aria-autocomplete': 'list',
@@ -370,7 +378,11 @@ export function useAutocomplete(props: AriaAutocompleteOptions, state: Autocompl
       enterKeyHint: 'go',
       onBlur,
       onFocus
-    },
+    };
+  }
+
+  return {
+    textFieldProps,
     collectionProps: mergeProps(collectionProps, {
       shouldUseVirtualFocus,
       disallowTypeAhead: true

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

@@ -17,6 +17,8 @@ export type Mutable<T> = {
   -readonly[P in keyof T]: T[P]
 }
 
+type FilterFn<T> = (textValue: string, node: Node<T>) => boolean;
+
 /** An immutable object representing a Node in a Collection. */
 export class CollectionNode<T> implements Node<T> {
   readonly type: string;
@@ -65,8 +67,64 @@ 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: FilterFn<T>): 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;
+  }
+}
+
+// TODO: naming, but essentially these nodes shouldn't be affected by filtering (BaseNode)?
+// Perhaps this filter logic should be in CollectionNode instead and the current logic of CollectionNode's filter should move to Table
+export class FilterLessNode<T> extends CollectionNode<T> {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: FilterFn<T>): FilterLessNode<T> | null {
+    return this.clone();
+  }
+}
+
+export class ItemNode<T> extends CollectionNode<T> {
+  static readonly type = 'item';
+
+  constructor(key: Key) {
+    super(ItemNode.type, key);
+  }
+
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: FilterFn<T>): ItemNode<T> | null {
+    if (filterFn(this.textValue, this)) {
+      return this.clone();
+    }
+
+    return null;
+  }
+}
+
+export class SectionNode<T> extends CollectionNode<T> {
+  static readonly type = 'section';
+
+  constructor(key: Key) {
+    super(SectionNode.type, key);
+  }
+
+  filter(collection: BaseCollection<T>, newCollection: BaseCollection<T>, filterFn: FilterFn<T>): SectionNode<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;
+  }
 }
 
 /**
@@ -224,134 +282,65 @@ 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> {
-    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;
-          }
-        }
+  filter(filterFn: FilterFn<T>, newCollection?: BaseCollection<T>): BaseCollection<T> {
+    if (newCollection == null) {
+      newCollection = new BaseCollection<T>();
+    }
 
-        // 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 && (lastNode.type !== 'section' && lastNode.type !== 'separator') && lastNode.parentKey === clonedNode.parentKey) {
-            lastNode.nextKey = clonedNode.key;
-            clonedNode.prevKey = lastNode.key;
-          } else {
-            clonedNode.prevKey = null;
-          }
-
-          clonedNode.nextKey = null;
-          newCollection.addNode(clonedNode);
-          lastNode = clonedNode;
-        }
+    let [firstKey, lastKey] = filterChildren(this, newCollection, this.firstKey, filterFn);
+    newCollection.firstKey = firstKey;
+    newCollection.lastKey = lastKey;
+    return newCollection;
+  }
+}
+
+function filterChildren<T>(collection: BaseCollection<T>, newCollection: BaseCollection<T>, firstChildKey: Key | null, filterFn: FilterFn<T>): [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) {
+      newNode.nextKey = 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') {
-    return true;
-  } else {
-    return filterFn(node.textValue);
   }
+
+  return [firstNode?.key ?? null, lastNode?.key ?? null];
 }