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

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;
+  }
 }
 
 /**
@@ -224,134 +233,62 @@ 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> {
+  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 && (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;
-        }
+export 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) {
+      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];
 }

packages/@react-aria/collections/src/CollectionBuilder.tsx

@@ -10,12 +10,12 @@
  * governing permissions and limitations under the License.
  */
 
-import {BaseCollection} from './BaseCollection';
+import {BaseCollection, CollectionNode} from './BaseCollection';
 import {BaseNode, Document, ElementNode} from './Document';
 import {CachedChildrenOptions, useCachedChildren} from './useCachedChildren';
 import {createPortal} from 'react-dom';
 import {FocusableContext} from '@react-aria/interactions';
-import {forwardRefType, Node} from '@react-types/shared';
+import {forwardRefType, Key, Node} from '@react-types/shared';
 import {Hidden} from './Hidden';
 import React, {createContext, ForwardedRef, forwardRef, JSX, ReactElement, ReactNode, useCallback, useContext, useMemo, useRef, useState} from 'react';
 import {useIsSSR} from '@react-aria/ssr';
@@ -127,21 +127,29 @@ function useCollectionDocument<T extends object, C extends BaseCollection<T>>(cr
 
 const SSRContext = createContext<BaseNode<any> | null>(null);
 
-function useSSRCollectionNode<T extends Element>(Type: string, props: object, ref: ForwardedRef<T>, rendered?: any, children?: ReactNode, render?: (node: Node<T>) => ReactElement) {
+export type CollectionNodeClass<T> = {
+  new (key: Key): CollectionNode<T>,
+  readonly type: string
+};
+
+// TODO: discuss the former Type arg, renamed to CollectionNodeClass
+function useSSRCollectionNode<T extends Element>(CollectionNodeClass: CollectionNodeClass<T>, props: object, ref: ForwardedRef<T>, rendered?: any, children?: ReactNode, render?: (node: Node<any>) => ReactElement) {
   // During SSR, portals are not supported, so the collection children will be wrapped in an SSRContext.
   // Since SSR occurs only once, we assume that the elements are rendered in order and never re-render.
   // Therefore we can create elements in our collection document during render so that they are in the
   // collection by the time we need to use the collection to render to the real DOM.
   // After hydration, we switch to client rendering using the portal.
   let itemRef = useCallback((element: ElementNode<any> | null) => {
-    element?.setProps(props, ref, rendered, render);
-  }, [props, ref, rendered, render]);
+    // TODO: check setProps api
+    element?.setProps(props, ref, rendered, render, CollectionNodeClass);
+  }, [props, ref, rendered, render, CollectionNodeClass]);
   let parentNode = useContext(SSRContext);
   if (parentNode) {
     // Guard against double rendering in strict mode.
     let element = parentNode.ownerDocument.nodesByProps.get(props);
     if (!element) {
-      element = parentNode.ownerDocument.createElement(Type);
+      // TODO: check this, maybe should just pass the CollectionNodeClass as a whole?
+      element = parentNode.ownerDocument.createElement(CollectionNodeClass.type);
       element.setProps(props, ref, rendered, render);
       parentNode.appendChild(element);
       parentNode.ownerDocument.updateCollection();
@@ -153,14 +161,16 @@ function useSSRCollectionNode<T extends Element>(Type: string, props: object, re
       : null;
   }
 
+  // console.log('type', CollectionNodeClass, CollectionNodeClass.type)
   // @ts-ignore
-  return <Type ref={itemRef}>{children}</Type>;
+  // TODO: could just make this a div perhaps, but keep it in line with how it used to work
+  return <CollectionNodeClass.type ref={itemRef}>{children}</CollectionNodeClass.type>;
 }
 
-// eslint-disable-next-line @typescript-eslint/no-unused-vars
-export function createLeafComponent<T extends object, P extends object, E extends Element>(type: string, render: (props: P, ref: ForwardedRef<E>) => ReactElement | null): (props: P & React.RefAttributes<E>) => ReactElement | null;
-export function createLeafComponent<T extends object, P extends object, E extends Element>(type: string, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null): (props: P & React.RefAttributes<E>) => ReactElement | null;
-export function createLeafComponent<P extends object, E extends Element>(type: string, render: (props: P, ref: ForwardedRef<E>, node?: any) => ReactElement | null): (props: P & React.RefAttributes<any>) => ReactElement | null {
+// TODO: check the signature of the CollectionNodeClass here and other places (aka useSSRCollectionNode and branchCompoennt). If I use the generic it complains. Perhaps it should be unknown? Or maybe the definitions in Listbox and stuff shouldn't use a generic?
+export function createLeafComponent<T extends object, P extends object, E extends Element>(CollectionNodeClass: CollectionNodeClass<any>, render: (props: P, ref: ForwardedRef<E>) => ReactElement | null): (props: P & React.RefAttributes<T>) => ReactElement | null;
+export function createLeafComponent<T extends object, P extends object, E extends Element>(CollectionNodeClass: CollectionNodeClass<any>, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null): (props: P & React.RefAttributes<T>) => ReactElement | null;
+export function createLeafComponent<P extends object, E extends Element>(CollectionNodeClass: CollectionNodeClass<any>, render: (props: P, ref: ForwardedRef<E>, node?: any) => ReactElement | null): (props: P & React.RefAttributes<any>) => ReactElement | null {
   let Component = ({node}) => render(node.props, node.props.ref, node);
   let Result = (forwardRef as forwardRefType)((props: P, ref: ForwardedRef<E>) => {
     let focusableProps = useContext(FocusableContext);
@@ -173,7 +183,7 @@ export function createLeafComponent<P extends object, E extends Element>(type: s
     }
 
     return useSSRCollectionNode(
-      type,
+      CollectionNodeClass,
       props,
       ref,
       'children' in props ? props.children : null,
@@ -191,11 +201,12 @@ export function createLeafComponent<P extends object, E extends Element>(type: s
   return Result;
 }
 
-export function createBranchComponent<T extends object, P extends {children?: any}, E extends Element>(type: string, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null, useChildren: (props: P) => ReactNode = useCollectionChildren): (props: P & React.RefAttributes<E>) => ReactElement | null {
+// TODO: check the signature of this too
+export function createBranchComponent<T extends object, P extends {children?: any}, E extends Element>(CollectionNodeClass: CollectionNodeClass<any>, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null, useChildren: (props: P) => ReactNode = useCollectionChildren): (props: P & React.RefAttributes<E>) => ReactElement | null {
   let Component = ({node}) => render(node.props, node.props.ref, node);
   let Result = (forwardRef as forwardRefType)((props: P, ref: ForwardedRef<E>) => {
     let children = useChildren(props);
-    return useSSRCollectionNode(type, props, ref, null, children, node => <Component node={node} />) ?? <></>;
+    return useSSRCollectionNode(CollectionNodeClass, props, ref, null, children, node => <Component node={node} />) ?? <></>;
   });
   // @ts-ignore
   Result.displayName = render.name;