@itwin/tree-widget-react 3.17.0 → 3.17.1

package/lib/esm/tree-widget-react/components/trees/models-tree/internal/ModelsTreeVisibilityHandler.js

@@ -7,10 +7,11 @@ import { assert, Id64 } from "@itwin/core-bentley";
 import { PerModelCategoryVisibility } from "@itwin/core-frontend";
 import { HierarchyNode, HierarchyNodeKey } from "@itwin/presentation-hierarchies";
 import { enableCategoryDisplay, loadCategoriesFromViewport } from "../../common/CategoriesVisibilityUtils.js";
-import { reduceWhile, toVoidPromise } from "../../common/Rxjs.js";
+import { getTooltipOptions, mergeVisibilityStatuses, releaseMainThreadOnItemsCount, setDifference, setIntersection } from "../../common/internal/Utils.js";
+import { toVoidPromise } from "../../common/Rxjs.js";
 import { createVisibilityStatus } from "../../common/Tooltip.js";
 import { createVisibilityHandlerResult } from "../../common/UseHierarchyVisibility.js";
-import { getIdsFromChildrenTree, releaseMainThreadOnItemsCount } from "../Utils.js";
+import { getIdsFromChildrenTree } from "../Utils.js";
 import { AlwaysAndNeverDrawnElementInfo } from "./AlwaysAndNeverDrawnElementInfo.js";
 import { createFilteredTree, parseCategoryKey } from "./FilteredTree.js";
 import { ModelsTreeNode } from "./ModelsTreeNode.js";
@@ -88,7 +89,9 @@ class ModelsTreeVisibilityHandlerImpl {
         if (!HierarchyNode.isInstancesNode(node)) {
             return of(createVisibilityStatus("disabled"));
         }
-        if (node.filtering?.filteredChildrenIdentifierPaths?.length && !node.filtering.isFilterTarget) {
+        // Only call getFilteredNodeVisibility when node is not a filter target, is not a child of filter target and has filtered children.
+        // Otherwise, it can be handled normally.
+        if (node.filtering?.filteredChildrenIdentifierPaths?.length && !node.filtering.isFilterTarget && !node.filtering.hasFilterTargetAncestor) {
             return this.getFilteredNodeVisibility({ node });
         }
         if (ModelsTreeNode.isSubjectNode(node)) {
@@ -109,13 +112,18 @@ class ModelsTreeVisibilityHandlerImpl {
             });
         }
         const categoryId = ModelsTreeNode.getCategoryId(node);
+        const rootCategoryIds = getRootCategoryIds({ parentKeys: node.parentKeys, modelId });
+        assert(rootCategoryIds !== undefined);
         if (!categoryId) {
             return of(createVisibilityStatus("disabled"));
         }
         return this.getElementDisplayStatus({
             elementIds: node.key.instanceKeys.map(({ id }) => id),
+            parentKeys: node.parentKeys,
+            rootCategoryIds,
             modelId,
             categoryId,
+            totalChildrenCount: node.extendedData?.childrenCount ?? 0,
         });
     }
     getFilteredNodeVisibility(props) {
@@ -135,14 +143,53 @@ class ModelsTreeVisibilityHandlerImpl {
                 })));
             }
             if (elements?.size) {
-                observables.push(from(elements).pipe(releaseMainThreadOnItemsCount(50), mergeMap(([categoryKey, elementsMap]) => {
+                const filterTargetElements = new Array();
+                elements.forEach((elementsMap) => elementsMap.forEach(({ isFilterTarget }, elementId) => {
+                    if (isFilterTarget) {
+                        filterTargetElements.push(elementId);
+                    }
+                }));
+                let childrenCountMapObs;
+                if (HierarchyNode.isClassGroupingNode(props.node)) {
+                    const groupingNodesFilterTargets = props.node.extendedData?.filterTargets;
+                    const nestedFilterTargetElements = filterTargetElements.filter((filterTarget) => !groupingNodesFilterTargets?.has(filterTarget));
+                    // Only need to request children count for indirect children filter targets.
+                    // Direct children filter targets already have children count stored in grouping nodes extended data.
+                    childrenCountMapObs = this.#idsCache.getAllChildrenCount({ elementIds: nestedFilterTargetElements });
+                }
+                else {
+                    childrenCountMapObs = this.#idsCache.getAllChildrenCount({ elementIds: filterTargetElements });
+                }
+                observables.push(childrenCountMapObs.pipe(
+                // Get children count for all filter target elements.
+                map((elementCountMap) => {
+                    if (!HierarchyNode.isClassGroupingNode(props.node) || !props.node.extendedData?.filterTargets) {
+                        return elementCountMap;
+                    }
+                    // Need to add children count (stored in grouping nodes' extended data) of direct children filter targets.
+                    const groupingNodesFilterTargets = props.node.extendedData.filterTargets;
+                    groupingNodesFilterTargets.forEach((value, key) => elementCountMap.set(key, value.childrenCount));
+                    return elementCountMap;
+                }), mergeMap((elementCountMap) => from(elements).pipe(releaseMainThreadOnItemsCount(50), mergeMap(([categoryKey, elementsMap]) => {
                     const { modelId, categoryId } = parseCategoryKey(categoryKey);
+                    let totalChildrenCount = 0;
+                    elementsMap.forEach((_, elementId) => {
+                        const childCount = elementCountMap.get(elementId);
+                        if (childCount) {
+                            totalChildrenCount += childCount;
+                        }
+                    });
+                    const parentElementIdsPath = getParentElementsIdsPath({ parentKeys: props.node.parentKeys, modelId });
+                    const rootCategoryIds = getRootCategoryIds({ parentKeys: props.node.parentKeys, modelId });
                     return this.getElementsDisplayStatus({
                         elementIds: [...elementsMap.keys()],
+                        rootCategoryIds: rootCategoryIds ?? categoryId,
+                        parentElementIdsPath,
                         modelId,
                         categoryId,
+                        childrenCount: totalChildrenCount,
                     });
-                })));
+                })))));
             }
             return merge(...observables);
         }), mergeVisibilityStatuses());
@@ -264,25 +311,33 @@ class ModelsTreeVisibilityHandlerImpl {
         return createVisibilityHandlerResult(this, props, result, this.#props.overrides?.getCategoryDisplayStatus);
     }
     getClassGroupingNodeDisplayStatus(node) {
-        const { elementIds, modelId, categoryId } = this.getGroupingNodeInfo(node);
+        const { elementIds, modelId, categoryId, childrenCount } = this.getGroupingNodeInfo(node);
+        const rootCategoryIds = getRootCategoryIds({ parentKeys: node.parentKeys, modelId });
+        assert(rootCategoryIds !== undefined);
         const result = this.getElementsDisplayStatus({
             elementIds,
             modelId,
             categoryId,
+            rootCategoryIds,
+            parentElementIdsPath: getParentElementsIdsPath({ parentKeys: node.parentKeys, modelId }),
+            childrenCount: childrenCount ?? 0,
         });
         return createVisibilityHandlerResult(this, { node }, result, this.#props.overrides?.getElementGroupingNodeDisplayStatus);
     }
-    getElementDisplayStatus({ elementIds, modelId, categoryId, }) {
+    getElementDisplayStatus({ parentKeys, totalChildrenCount, elementIds, modelId, rootCategoryIds, categoryId, }) {
         const result = this.getElementsDisplayStatus({
+            rootCategoryIds,
             elementIds,
             modelId,
             categoryId,
+            parentElementIdsPath: getParentElementsIdsPath({ parentKeys, modelId }),
+            childrenCount: totalChildrenCount,
         });
         return createVisibilityHandlerResult(this, { elementId: elementIds[0], modelId, categoryId }, result, this.#props.overrides?.getElementDisplayStatus);
     }
     getElementsDisplayStatus(props) {
         return defer(() => {
-            const { modelId, categoryId, elementIds } = props;
+            const { modelId, categoryId, elementIds, parentElementIdsPath, childrenCount, rootCategoryIds } = props;
             if (!this.#props.viewport.view.viewsModel(modelId)) {
                 return of(elementIds).pipe(this.getSubModeledElementsVisibilityStatus({
                     tooltips: {
@@ -296,6 +351,10 @@ class ModelsTreeVisibilityHandlerImpl {
             }
             return this.getVisibilityFromAlwaysAndNeverDrawnElements({
                 elements: elementIds,
+                parentElementIdsPath,
+                modelIds: modelId,
+                rootCategoryIds,
+                childrenCount,
                 defaultStatus: () => {
                     const status = this.getDefaultCategoryVisibilityStatus({ categoryIds: categoryId, modelId, ignoreTooltip: true });
                     return createVisibilityStatus(status.state, getTooltipOptions(`modelsTree.groupingNode.${status.state}ThroughCategory`));
@@ -331,7 +390,9 @@ class ModelsTreeVisibilityHandlerImpl {
             if (!HierarchyNode.isInstancesNode(node)) {
                 return EMPTY;
             }
-            if (node.filtering?.filteredChildrenIdentifierPaths?.length && !node.filtering.isFilterTarget) {
+            // Only call changeFilteredNodeVisibility when node is not a filter target, is not a child of filter target and has filtered children.
+            // Otherwise, it can be handled normally.
+            if (node.filtering?.filteredChildrenIdentifierPaths?.length && !node.filtering.isFilterTarget && !node.filtering.hasFilterTargetAncestor) {
                 return this.changeFilteredNodeVisibility({ node, on });
             }
             if (ModelsTreeNode.isSubjectNode(node)) {
@@ -355,12 +416,18 @@ class ModelsTreeVisibilityHandlerImpl {
             if (!categoryId) {
                 return EMPTY;
             }
-            return this.changeElementsState({
-                elementIds: new Set([...node.key.instanceKeys.map(({ id }) => id)]),
+            const elementIds = new Set(node.key.instanceKeys.map(({ id }) => id));
+            return this.#idsCache.getChildrenTree({ elementIds }).pipe(map((childrenTree) => {
+                // Children tree contains provided elementIds, they are at the root of this tree.
+                // We want to skip them and only get ids of children.
+                return getIdsFromChildrenTree({ tree: childrenTree, predicate: ({ depth }) => depth > 0 });
+            }), mergeMap((children) => this.changeElementsState({
+                elementIds,
                 modelId,
+                children: children.size > 0 ? children : undefined,
                 categoryId,
                 on,
-            });
+            })));
         });
         if (this.#props.viewport.isAlwaysDrawnExclusive) {
             return concat(this.removeAlwaysDrawnExclusive(), changeObs);
@@ -369,7 +436,7 @@ class ModelsTreeVisibilityHandlerImpl {
     }
     changeFilteredNodeVisibility({ on, ...props }) {
         assert(this.#filteredTree !== undefined);
-        return from(this.#filteredTree).pipe(map((filteredTree) => filteredTree.getVisibilityChangeTargets(props.node)), mergeMap(({ subjects, models, categories, elements }) => {
+        return from(this.#filteredTree).pipe(map((filteredTree) => filteredTree.getVisibilityChangeTargets(props.node)), mergeMap(({ subjects, categories, elements, models }) => {
             const observables = new Array();
             if (subjects?.size) {
                 observables.push(this.changeSubjectNodeState([...subjects], on));
@@ -384,10 +451,24 @@ class ModelsTreeVisibilityHandlerImpl {
                 })));
             }
             if (elements?.size) {
-                observables.push(from(elements).pipe(mergeMap(([categoryKey, elementsMap]) => {
+                const filterTargetElements = new Array();
+                const elementsSet = new Set();
+                elements.forEach((elementsMap) => elementsMap.forEach(({ isFilterTarget }, elementId) => {
+                    elementsSet.add(elementId);
+                    if (isFilterTarget) {
+                        filterTargetElements.push(elementId);
+                    }
+                }));
+                observables.push(this.#idsCache.getChildrenTree({ elementIds: filterTargetElements }).pipe(map((childrenTree) => getIdsFromChildrenTree({ tree: childrenTree })), map((childrenIds) => setDifference(childrenIds, elementsSet)), mergeMap((childElementIds) => from(elements).pipe(mergeMap(([categoryKey, elementsMap]) => {
                     const { modelId, categoryId } = parseCategoryKey(categoryKey);
-                    return this.changeElementsState({ modelId, categoryId, elementIds: new Set([...elementsMap.keys()]), on });
-                })));
+                    return this.changeElementsState({
+                        modelId,
+                        categoryId,
+                        elementIds: new Set([...elementsMap.keys()]),
+                        on,
+                        children: childElementIds.size > 0 ? childElementIds : undefined,
+                    });
+                })))));
             }
             return merge(...observables);
         }));
@@ -422,9 +503,7 @@ class ModelsTreeVisibilityHandlerImpl {
             const idsObs = from(Id64.iterable(ids));
             if (!on) {
                 viewport.changeModelDisplay(ids, false);
-                return idsObs.pipe(mergeMap((modelId) => this.#idsCache.getModelCategories(modelId).pipe(map((categoryIds) => {
-                    return { modelId, categoryIds };
-                }))), mergeMap(({ modelId, categoryIds }) => this.#idsCache.getCategoriesModeledElements(modelId, categoryIds)), mergeMap((modeledElementIds) => this.changeModelState({ ids: modeledElementIds, on })));
+                return idsObs.pipe(mergeMap((modelId) => forkJoin({ categoryIds: this.#idsCache.getModelCategories(modelId), modelId: of(modelId) })), mergeMap(({ modelId, categoryIds }) => this.#idsCache.getCategoriesModeledElements(modelId, categoryIds)), mergeMap((modeledElementIds) => this.changeModelState({ ids: modeledElementIds, on })));
             }
             return concat(defer(() => {
                 viewport.perModelCategoryVisibility.clearOverrides(ids);
@@ -440,15 +519,12 @@ class ModelsTreeVisibilityHandlerImpl {
         return forkJoin({
             categories: this.#idsCache.getModelCategories(modelId),
             alwaysDrawnElements: this.getAlwaysOrNeverDrawnElements({ modelIds: modelId, setType: "always" }),
-        }).pipe(mergeMap(async ({ categories, alwaysDrawnElements }) => {
+        }).pipe(mergeMap(({ categories, alwaysDrawnElements }) => {
             const alwaysDrawn = this.#props.viewport.alwaysDrawn;
             if (alwaysDrawn && alwaysDrawnElements) {
                 viewport.setAlwaysDrawn(setDifference(alwaysDrawn, alwaysDrawnElements));
             }
-            categories.forEach((categoryId) => {
-                this.changeCategoryStateInViewportAccordingToModelVisibility(modelId, categoryId, false);
-            });
-            await viewport.addViewedModels(modelId);
+            return concat(from(Id64.iterable(categories)).pipe(releaseMainThreadOnItemsCount(300), map((categoryId) => this.changeCategoryStateInViewportAccordingToModelVisibility(modelId, categoryId, false))), defer(async () => viewport.addViewedModels(modelId)));
         }));
     }
     changeCategoryStateInViewportAccordingToModelVisibility(modelId, categoryId, on) {
@@ -470,12 +546,7 @@ class ModelsTreeVisibilityHandlerImpl {
         const result = defer(() => {
             const viewport = this.#props.viewport;
             const { modelId, categoryIds, on } = props;
-            return concat(props.on && !viewport.view.viewsModel(modelId) ? this.showModelWithoutAnyCategoriesOrElements(modelId) : EMPTY, defer(() => {
-                for (const categoryId of Id64.iterable(categoryIds)) {
-                    this.changeCategoryStateInViewportAccordingToModelVisibility(modelId, categoryId, on);
-                }
-                return this.clearAlwaysAndNeverDrawnElements(props);
-            }), this.#idsCache
+            return concat(props.on && !viewport.view.viewsModel(modelId) ? this.showModelWithoutAnyCategoriesOrElements(modelId) : EMPTY, from(Id64.iterable(categoryIds)).pipe(releaseMainThreadOnItemsCount(300), map((categoryId) => this.changeCategoryStateInViewportAccordingToModelVisibility(modelId, categoryId, on))), this.clearAlwaysAndNeverDrawnElements(props), this.#idsCache
                 .getCategoriesModeledElements(modelId, categoryIds)
                 .pipe(mergeMap((modeledElementIds) => this.changeModelState({ ids: modeledElementIds, on }))));
         });
@@ -483,12 +554,25 @@ class ModelsTreeVisibilityHandlerImpl {
     }
     doChangeElementsState(props) {
         return defer(() => {
-            const { modelId, categoryId, elementIds, on } = props;
+            const { modelId, categoryId, elementIds, on, children } = props;
             const viewport = this.#props.viewport;
             return concat(on && !viewport.view.viewsModel(modelId) ? this.showModelWithoutAnyCategoriesOrElements(modelId) : EMPTY, defer(() => {
                 const categoryVisibility = this.getDefaultCategoryVisibilityStatus({ categoryIds: categoryId, modelId, ignoreTooltip: true });
-                const isDisplayedByDefault = categoryVisibility.state === "visible";
-                return this.queueElementsVisibilityChange(elementIds, on, isDisplayedByDefault);
+                const isCategoryVisible = categoryVisibility.state === "visible";
+                // Make sure to add all children to always/never drawn list.
+                const isDisplayedByDefault = 
+                // When category is visible and elements need to be turned off, or when category is hidden and elements need to be turned on,
+                // We can set isDisplayedByDefault to isCategoryVisible. This allows to not check if each element is in the elementIds list or not.
+                isCategoryVisible === !on
+                    ? () => isCategoryVisible
+                    : (elementId) => {
+                        if (elementIds.has(elementId)) {
+                            return isCategoryVisible;
+                        }
+                        return !on;
+                    };
+                const elementsToChange = children ? [...elementIds, ...(typeof children === "string" ? [children] : children)] : elementIds;
+                return this.queueElementsVisibilityChange(elementsToChange, on, isDisplayedByDefault);
             }), from(elementIds).pipe(mergeMap((elementId) => forkJoin({
                 elementId: of(elementId),
                 isSubModel: this.#idsCache.hasSubModel(elementId),
@@ -500,7 +584,14 @@ class ModelsTreeVisibilityHandlerImpl {
      * @see `changeElementState`
      */
     changeElementGroupingNodeState(node, on) {
-        const result = this.doChangeElementsState({ ...this.getGroupingNodeInfo(node), on });
+        const { modelId, categoryId, elementIds } = this.getGroupingNodeInfo(node);
+        const result = this.#idsCache.getChildrenTree({ elementIds }).pipe(map((childrenTree) => getIdsFromChildrenTree({ tree: childrenTree, predicate: ({ depth }) => depth > 0 })), mergeMap((children) => this.doChangeElementsState({
+            modelId,
+            categoryId,
+            elementIds,
+            children: children.size > 0 ? children : undefined,
+            on,
+        })));
         return createVisibilityHandlerResult(this, { node, on }, result, this.#props.overrides?.changeElementGroupingNodeState);
     }
     /**
@@ -515,7 +606,7 @@ class ModelsTreeVisibilityHandlerImpl {
         const finishedSubject = new Subject();
         // observable to track if visibility change is finished/cancelled
         const changeFinished = finishedSubject.pipe(startWith(false), shareReplay(1), filter((finished) => finished));
-        const changeObservable = from(elementIds).pipe(
+        const changeObservable = from(Id64.iterable(elementIds)).pipe(
         // check if visibility change is not finished (cancelled) due to change overall change request being cancelled
         takeUntil(changeFinished), this.changeElementStateNoChildrenOperator({ on, isDisplayedByDefault: visibleByDefault }), tap({
             next: () => {
@@ -546,7 +637,7 @@ class ModelsTreeVisibilityHandlerImpl {
                     const wasRemoved = acc.neverDrawn.delete(elementId);
                     acc.changedNeverDrawn ||= wasRemoved;
                     // If exclusive mode is enabled, we must add the element to the always drawn list.
-                    if ((!isDisplayedByDefault || isAlwaysDrawnExclusive) && !acc.alwaysDrawn.has(elementId)) {
+                    if ((!isDisplayedByDefault(elementId) || isAlwaysDrawnExclusive) && !acc.alwaysDrawn.has(elementId)) {
                         acc.alwaysDrawn.add(elementId);
                         acc.changedAlwaysDrawn = true;
                     }
@@ -555,7 +646,7 @@ class ModelsTreeVisibilityHandlerImpl {
                     const wasRemoved = acc.alwaysDrawn.delete(elementId);
                     acc.changedAlwaysDrawn ||= wasRemoved;
                     // If exclusive mode is not enabled, we have to add the element to the never drawn list.
-                    if (isDisplayedByDefault && !isAlwaysDrawnExclusive && !acc.neverDrawn.has(elementId)) {
+                    if (isDisplayedByDefault(elementId) && !isAlwaysDrawnExclusive && !acc.neverDrawn.has(elementId)) {
                         acc.neverDrawn.add(elementId);
                         acc.changedNeverDrawn = true;
                     }
@@ -577,20 +668,22 @@ class ModelsTreeVisibilityHandlerImpl {
         if (totalCount === 0) {
             return props.defaultStatus();
         }
-        if (neverDrawn?.size === totalCount) {
+        const neverDrawnSize = neverDrawn ? Id64.sizeOf(neverDrawn) : undefined;
+        const alwaysDrawnSize = alwaysDrawn ? Id64.sizeOf(alwaysDrawn) : undefined;
+        if (neverDrawnSize === totalCount) {
             return createVisibilityStatus("hidden", getTooltipOptions(props.tooltips.allElementsInNeverDrawnList, ignoreTooltip));
         }
-        if (alwaysDrawn?.size === totalCount) {
+        if (alwaysDrawnSize === totalCount) {
             return createVisibilityStatus("visible", getTooltipOptions(props.tooltips.allElementsInAlwaysDrawnList, ignoreTooltip));
         }
         const viewport = this.#props.viewport;
         if (viewport.isAlwaysDrawnExclusive && viewport.alwaysDrawn?.size) {
-            return alwaysDrawn?.size
+            return alwaysDrawnSize
                 ? createVisibilityStatus("partial", getTooltipOptions(props.tooltips.elementsInBothAlwaysAndNeverDrawn, ignoreTooltip))
                 : createVisibilityStatus("hidden", getTooltipOptions(props.tooltips.noElementsInExclusiveAlwaysDrawnList, ignoreTooltip));
         }
         const status = props.defaultStatus();
-        if ((status.state === "visible" && neverDrawn?.size) || (status.state === "hidden" && alwaysDrawn?.size)) {
+        if ((status.state === "visible" && neverDrawnSize) || (status.state === "hidden" && alwaysDrawnSize)) {
             return createVisibilityStatus("partial", getTooltipOptions(undefined, ignoreTooltip));
         }
         return status;
@@ -606,21 +699,41 @@ class ModelsTreeVisibilityHandlerImpl {
             return of(props.defaultStatus());
         }
         if ("elements" in props) {
-            return of(this.getVisibilityFromAlwaysAndNeverDrawnElementsImpl({
-                ...props,
-                alwaysDrawn: viewport.alwaysDrawn?.size ? setIntersection(props.elements, viewport.alwaysDrawn) : undefined,
-                neverDrawn: viewport.neverDrawn?.size ? setIntersection(props.elements, viewport.neverDrawn) : undefined,
-                totalCount: Id64.sizeOf(props.elements),
-                ignoreTooltip,
+            const parentElementIdsPath = [...props.parentElementIdsPath, props.elements];
+            return forkJoin({
+                childAlwaysDrawn: this.getAlwaysOrNeverDrawnElements({
+                    modelIds: props.modelIds,
+                    categoryIds: props.rootCategoryIds,
+                    parentElementIdsPath,
+                    setType: "always",
+                }),
+                childNeverDrawn: this.getAlwaysOrNeverDrawnElements({
+                    modelIds: props.modelIds,
+                    categoryIds: props.rootCategoryIds,
+                    parentElementIdsPath,
+                    setType: "never",
+                }),
+            }).pipe(map(({ childAlwaysDrawn, childNeverDrawn }) => {
+                const alwaysDrawn = new Set([...childAlwaysDrawn, ...(viewport.alwaysDrawn?.size ? setIntersection(props.elements, viewport.alwaysDrawn) : [])]);
+                const neverDrawn = new Set([...childNeverDrawn, ...(viewport.neverDrawn?.size ? setIntersection(props.elements, viewport.neverDrawn) : [])]);
+                return this.getVisibilityFromAlwaysAndNeverDrawnElementsImpl({
+                    ...props,
+                    alwaysDrawn: alwaysDrawn.size > 0 ? alwaysDrawn : undefined,
+                    neverDrawn: neverDrawn.size > 0 ? neverDrawn : undefined,
+                    totalCount: props.childrenCount + Id64.sizeOf(props.elements),
+                    ignoreTooltip,
+                });
             }));
         }
         const { modelId, categoryIds } = props.categoryProps;
-        return from(Id64.iterable(categoryIds)).pipe(mergeMap((categoryId) => forkJoin({
-            categoryId: of(categoryId),
-            totalCount: this.#idsCache.getCategoryElementsCount(modelId, categoryId),
-            alwaysDrawn: this.getAlwaysOrNeverDrawnElements({ modelIds: modelId, categoryIds: categoryId, setType: "always" }),
-            neverDrawn: this.getAlwaysOrNeverDrawnElements({ modelIds: modelId, categoryIds: categoryId, setType: "never" }),
-        })), 
+        return from(Id64.iterable(categoryIds)).pipe(releaseMainThreadOnItemsCount(100), mergeMap((categoryId) => {
+            return forkJoin({
+                categoryId: of(categoryId),
+                totalCount: this.#idsCache.getCategoryElementsCount(modelId, categoryId),
+                alwaysDrawn: this.getAlwaysOrNeverDrawnElements({ modelIds: modelId, categoryIds: categoryId, setType: "always" }),
+                neverDrawn: this.getAlwaysOrNeverDrawnElements({ modelIds: modelId, categoryIds: categoryId, setType: "never" }),
+            });
+        }), 
         // There is a known bug:
         // Categories that don't have root elements will make visibility result incorrect
         // E.g.:
@@ -639,9 +752,31 @@ class ModelsTreeVisibilityHandlerImpl {
         }), mergeVisibilityStatuses());
     }
     getAlwaysOrNeverDrawnElements(props) {
-        return this.#alwaysAndNeverDrawnElements
-            .getElementsTree(props)
-            .pipe(map((childrenTree) => getIdsFromChildrenTree({ tree: childrenTree, predicate: ({ treeEntry }) => treeEntry.isInAlwaysOrNeverDrawnSet })));
+        if (!this.#filteredTree) {
+            return this.#alwaysAndNeverDrawnElements
+                .getElementsTree(props)
+                .pipe(map((childrenTree) => getIdsFromChildrenTree({ tree: childrenTree, predicate: ({ treeEntry }) => treeEntry.isInAlwaysOrNeverDrawnSet })));
+        }
+        // When filtered tree is present, children tree retrieved from alwaysAndNeverDrawnElements may include elements that are not present in filtered tree.
+        // Need to filter out such elements.
+        return forkJoin({
+            filteredTree: from(this.#filteredTree),
+            childrenTree: this.#alwaysAndNeverDrawnElements.getElementsTree(props),
+        }).pipe(map(({ filteredTree, childrenTree }) => {
+            const parentIdsPath = [props.modelIds];
+            if ("categoryIds" in props) {
+                parentIdsPath.push(props.categoryIds);
+                if ("parentElementIdsPath" in props) {
+                    props.parentElementIdsPath.forEach((parentElementIds) => parentIdsPath.push(parentElementIds));
+                }
+            }
+            // Get all ids that exist in filtered tree.
+            const elements = filteredTree.getElementsFromUnfilteredChildrenTree({ parentIdsPath, childrenTree });
+            return elements
+                ? // Need to filter elements that are in always/never drawn set.
+                    setIntersection(elements, getIdsFromChildrenTree({ tree: childrenTree, predicate: ({ treeEntry }) => treeEntry.isInAlwaysOrNeverDrawnSet }))
+                : new Set();
+        }));
     }
     clearAlwaysAndNeverDrawnElements(props) {
         return forkJoin({
@@ -662,7 +797,9 @@ class ModelsTreeVisibilityHandlerImpl {
         const categoryId = ModelsTreeNode.getCategoryId(node);
         assert(!!modelId && !!categoryId);
         const elementIds = new Set(node.groupedInstanceKeys.map((key) => key.id));
-        return { modelId, categoryId, elementIds };
+        const childrenCount = node.extendedData?.childrenCount;
+        const isFiltered = node.extendedData?.isFiltered;
+        return { modelId, categoryId, elementIds, childrenCount, isFiltered };
     }
     getSubModeledElementsVisibilityStatus({ parentNodeVisibilityStatus, haveSubModel, tooltips, ignoreTooltips, }) {
         return (obs) => {
@@ -672,7 +809,7 @@ class ModelsTreeVisibilityHandlerImpl {
                 if (haveSubModel === "yes") {
                     return of(modeledElementIds);
                 }
-                return from(Id64.iterable(modeledElementIds)).pipe(mergeMap((elementId) => forkJoin({
+                return from(Id64.iterable(modeledElementIds)).pipe(releaseMainThreadOnItemsCount(100), mergeMap((elementId) => forkJoin({
                     hasSubModel: this.#idsCache.hasSubModel(elementId),
                     elementId: of(elementId),
                 })), filter(({ hasSubModel }) => hasSubModel), map(({ elementId }) => elementId), toArray());
@@ -682,51 +819,11 @@ class ModelsTreeVisibilityHandlerImpl {
                 if (Id64.sizeOf(modeledElementIds) === 0) {
                     return of(parentNodeVisibilityStatus);
                 }
-                return this.getModelVisibilityStatus({ modelIds: modeledElementIds }).pipe(startWith(parentNodeVisibilityStatus), mergeVisibilityStatuses(tooltips, ignoreTooltips));
+                return this.getModelVisibilityStatus({ modelIds: typeof modeledElementIds === "string" ? [modeledElementIds] : modeledElementIds }).pipe(startWith(parentNodeVisibilityStatus), mergeVisibilityStatuses(tooltips, ignoreTooltips));
             }));
         };
     }
 }
-function mergeVisibilities(obs) {
-    return obs.pipe(reduceWhile((x) => x.allVisible || x.allHidden, (acc, val) => {
-        acc.allVisible &&= val === "visible";
-        acc.allHidden &&= val === "hidden";
-        return acc;
-    }, { allVisible: true, allHidden: true }), map((x) => {
-        if (!x) {
-            return "empty";
-        }
-        return x.allVisible ? "visible" : x.allHidden ? "hidden" : "partial";
-    }));
-}
-function mergeVisibilityStatuses(tooltipMap, ignoreTooltip) {
-    return (obs) => {
-        return obs.pipe(map((visibilityStatus) => visibilityStatus.state), mergeVisibilities, map((visibility) => {
-            if (visibility === "empty") {
-                visibility = "visible";
-            }
-            return createVisibilityStatus(visibility, getTooltipOptions(tooltipMap?.[visibility], ignoreTooltip));
-        }));
-    };
-}
-function setDifference(lhs, rhs) {
-    const result = new Set();
-    for (const x of lhs) {
-        if (!rhs.has(x)) {
-            result.add(x);
-        }
-    }
-    return result;
-}
-function setIntersection(lhs, rhs) {
-    const result = new Set();
-    for (const x of lhs) {
-        if (rhs.has(x)) {
-            result.add(x);
-        }
-    }
-    return result;
-}
 /**
  * Enables display of all given models. Also enables display of all categories and clears always and
  * never drawn lists in the viewport.
@@ -793,9 +890,33 @@ export async function toggleModels(models, enable, viewport) {
 export function areAllModelsVisible(models, viewport) {
     return models.length !== 0 ? models.every((id) => viewport.viewsModel(id)) : false;
 }
-function getTooltipOptions(key, ignoreTooltip) {
-    return {
-        useTooltip: ignoreTooltip ? false : key,
-    };
+function getParentElementsIdsPath(props) {
+    const parentElementIdsPath = new Array();
+    const modelIndex = props.parentKeys.findIndex((parentKey) => HierarchyNodeKey.isInstances(parentKey) && parentKey.instanceKeys.some((instanceKey) => instanceKey.id === props.modelId));
+    if (modelIndex === -1) {
+        return parentElementIdsPath;
+    }
+    // Hierarchy is model -> category -> root elements.
+    // Root elements index is 2 more than models'.
+    const firstParentIndex = modelIndex + 2;
+    for (let i = firstParentIndex; i < props.parentKeys.length; ++i) {
+        const parentKey = props.parentKeys[i];
+        if (!HierarchyNodeKey.isInstances(parentKey)) {
+            continue;
+        }
+        parentElementIdsPath.push(parentKey.instanceKeys.map(({ id }) => id));
+    }
+    return parentElementIdsPath;
+}
+function getRootCategoryIds(props) {
+    const modelIndex = props.parentKeys.findIndex((parentKey) => HierarchyNodeKey.isInstances(parentKey) && parentKey.instanceKeys.some((instanceKey) => instanceKey.id === props.modelId));
+    if (modelIndex === -1 || modelIndex + 1 === props.parentKeys.length) {
+        return undefined;
+    }
+    const rootCategoryKeys = props.parentKeys[modelIndex + 1];
+    if (!HierarchyNodeKey.isInstances(rootCategoryKeys)) {
+        return undefined;
+    }
+    return rootCategoryKeys.instanceKeys.map(({ id }) => id);
 }
 //# sourceMappingURL=ModelsTreeVisibilityHandler.js.map