perf(react-tree): avoid processing hidden subtrees

This PR makes FlatTree support large data sets with good performance as
long as most items are not rendered at once, which should be the common
case.

This is done by optimizing two key functions.

The first deals with multiselect state, and simply batches all updates
instead of calling `ImmutableMap.set()` in a loop over the entire
subtree, which had a big performance impact due to each call copying
the entire Map.
When toggling the checkbox for a branch with a 2800 item subtree, the
delay went from an unusable 1800 ms to 8 ms.

The second fix is to avoid processing of invisible subtrees when
generating the visible tree based on open items.
This one requires a larger tree to be noticeable, but at 35k items, the
generator took 600 ms before and 10 ms after the fix.

An easy way to test is to modify the story "TreeSelection", changing the
`items` array so that it's generated with a loop.

Author: toerndev

change/@fluentui-react-tree-bc995a34-20ac-4060-84dd-dcfcf629b71d.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "performance optimization in react-tree",
+  "packageName": "@fluentui/react-tree",
+  "email": "maachin@gmail.com",
+  "dependentChangeType": "patch"
+}

packages/react-components/react-tree/library/src/components/FlatTree/useFlatControllableCheckedItems.ts

@@ -29,6 +29,7 @@ export function createNextFlatCheckedItems(
   if (data.selectionMode === 'single') {
     return ImmutableMap.from([[data.value, data.checked]]);
   }
+
   const treeItem = headlessTree.get(data.value);
   if (!treeItem) {
     if (process.env.NODE_ENV !== 'production') {
@@ -40,34 +41,50 @@ export function createNextFlatCheckedItems(
     }
     return previousCheckedItems;
   }
-  let nextCheckedItems = previousCheckedItems;
+
+  // Calling `ImmutableMap.set()` creates a new ImmutableMap - avoid this in loops.
+  // Instead write all updates to a native Map and create a new ImmutableMap at the end.
+  // Note that all descendants of the toggled item are processed even if they are collapsed,
+  // making the choice of algorithm more important.
+
+  const nextCheckedItemsMap = new Map(ImmutableMap.dangerouslyGetInternalMap(previousCheckedItems));
+
+  // The toggled item itself
+  nextCheckedItemsMap.set(data.value, data.checked);
+
+  // Descendant updates
   for (const children of headlessTree.subtree(data.value)) {
-    nextCheckedItems = nextCheckedItems.set(children.value, data.checked);
+    nextCheckedItemsMap.set(children.value, data.checked);
   }
-  nextCheckedItems = nextCheckedItems.set(data.value, data.checked);
 
+  // Ancestor updates - must be done after adding descendants and the toggle item.
+  // If any ancestor is mixed, all ancestors above it are mixed too.
   let isAncestorsMixed = false;
-  for (const parent of headlessTree.ancestors(treeItem.value)) {
-    // if one parent is mixed, all ancestors are mixed
+
+  for (const ancestor of headlessTree.ancestors(treeItem.value)) {
     if (isAncestorsMixed) {
-      nextCheckedItems = nextCheckedItems.set(parent.value, 'mixed');
+      nextCheckedItemsMap.set(ancestor.value, 'mixed');
       continue;
     }
-    let checkedChildrenAmount = 0;
-    for (const child of headlessTree.children(parent.value)) {
-      if ((nextCheckedItems.get(child.value) || false) === data.checked) {
-        checkedChildrenAmount++;
+
+    // For each ancestor, if all of its children now have the same checked state as the toggled item,
+    // set the ancestor to that checked state too. Otherwise it is 'mixed'.
+    let childrenWithSameState = 0;
+    for (const child of headlessTree.children(ancestor.value)) {
+      if ((nextCheckedItemsMap.get(child.value) || false) === data.checked) {
+        childrenWithSameState++;
       }
     }
-    // if all children are checked, parent is checked
-    if (checkedChildrenAmount === parent.childrenValues.length) {
-      nextCheckedItems = nextCheckedItems.set(parent.value, data.checked);
+
+    if (childrenWithSameState === ancestor.childrenValues.length) {
+      nextCheckedItemsMap.set(ancestor.value, data.checked);
     } else {
-      // if one parent is mixed, all ancestors are mixed
+      nextCheckedItemsMap.set(ancestor.value, 'mixed');
       isAncestorsMixed = true;
-      nextCheckedItems = nextCheckedItems.set(parent.value, 'mixed');
     }
   }
+
+  const nextCheckedItems = ImmutableMap.from(nextCheckedItemsMap);
   return nextCheckedItems;
 }
 

packages/react-components/react-tree/library/src/utils/createHeadlessTree.ts

@@ -261,10 +261,37 @@ function* HeadlessTreeVisibleItemsGenerator<Props extends HeadlessTreeItemProps>
   virtualTreeItems: HeadlessTree<Props>,
 ): Generator<HeadlessTreeItem<Props>, void, void> {
   let index = 0;
-  for (const item of HeadlessTreeSubtreeGenerator(virtualTreeItems.root.value, virtualTreeItems)) {
-    if (isItemVisible(item, openItems, virtualTreeItems)) {
-      item.index = index++;
-      yield item;
+  for (const item of recursiveVisibleItems(virtualTreeItems.root.value, openItems, virtualTreeItems)) {
+    item.index = index++;
+    yield item;
+  }
+}
+
+function* recursiveVisibleItems<Props extends HeadlessTreeItemProps>(
+  parentValue: TreeItemValue,
+  openItems: ImmutableSet<TreeItemValue>,
+  virtualTreeItems: HeadlessTree<Props>,
+): Generator<HeadlessTreeItem<Props>, void, void> {
+  const parent = virtualTreeItems.get(parentValue);
+  if (!parent || parent.childrenValues.length === 0) {
+    return;
+  }
+
+  for (const childValue of parent.childrenValues) {
+    const child = virtualTreeItems.get(childValue);
+    if (!child) {
+      continue;
+    }
+
+    if (isItemVisible(child, openItems, virtualTreeItems)) {
+      yield child;
+
+      // Process children only as long as their parents are open.
+      // This makes it possible to have large trees with good performance as
+      // long as most branches are not expanded.
+      if (openItems.has(childValue)) {
+        yield* recursiveVisibleItems(childValue, openItems, virtualTreeItems);
+      }
     }
   }
 }