@itwin/tree-widget-react 4.0.0-alpha.13 → 4.0.0-alpha.15

package/lib/esm/tree-widget-react/components/trees/categories-tree/internal/CategoriesTreeVisibilityHandler.js

@@ -2,16 +2,16 @@
  * Copyright (c) Bentley Systems, Incorporated. All rights reserved.
  * See LICENSE.md in the project root for license terms and full copyright notice.
  *--------------------------------------------------------------------------------------------*/
-import { concat, concatAll, defaultIfEmpty, defer, distinct, EMPTY, filter, firstValueFrom, forkJoin, from, fromEventPattern, map, merge, mergeAll, mergeMap, of, reduce, shareReplay, startWith, Subject, take, takeUntil, tap, toArray, } from "rxjs";
+import { concat, concatAll, defaultIfEmpty, defer, EMPTY, filter, firstValueFrom, forkJoin, from, fromEventPattern, map, merge, mergeMap, of, reduce, shareReplay, startWith, Subject, take, takeUntil, tap, } from "rxjs";
 import { assert, Id64 } from "@itwin/core-bentley";
 import { PerModelCategoryVisibility } from "@itwin/core-frontend";
 import { HierarchyNode, HierarchyNodeKey } from "@itwin/presentation-hierarchies";
 import { AlwaysAndNeverDrawnElementInfo } from "../../common/internal/AlwaysAndNeverDrawnElementInfo.js";
 import { toVoidPromise } from "../../common/internal/Rxjs.js";
 import { createVisibilityStatus } from "../../common/internal/Tooltip.js";
-import { getClassesByView, releaseMainThreadOnItemsCount, setDifference, setIntersection } from "../../common/internal/Utils.js";
+import { getClassesByView, getSetFromId64Arg, releaseMainThreadOnItemsCount, setDifference, setIntersection } from "../../common/internal/Utils.js";
 import { createVisibilityChangeEventListener } from "../../common/internal/VisibilityChangeEventListener.js";
-import { changeElementStateNoChildrenOperator, enableCategoryDisplay, enableSubCategoryDisplay, filterSubModeledElementIds, getElementOverriddenVisibility, getElementVisibility, getSubModeledElementsVisibilityStatus, getVisibilityFromAlwaysAndNeverDrawnElementsImpl, mergeVisibilityStatuses, } from "../../common/internal/VisibilityUtils.js";
+import { changeElementStateNoChildrenOperator, enableCategoryDisplay, enableSubCategoryDisplay, getVisibilityFromAlwaysAndNeverDrawnElementsImpl, mergeVisibilityStatuses, } from "../../common/internal/VisibilityUtils.js";
 import { createVisibilityHandlerResult } from "../../common/UseHierarchyVisibility.js";
 import { CategoriesTreeNode } from "./CategoriesTreeNode.js";
 import { createFilteredTree, parseCategoryKey, parseSubCategoryKey } from "./FilteredTree.js";
@@ -38,7 +38,7 @@ class CategoriesTreeVisibilityHandlerImpl {
             listeners: {
                 models: true,
                 categories: true,
-                elements: _props.hierarchyConfig.showElements,
+                elements: true,
                 displayStyle: true,
             },
         });
@@ -101,26 +101,26 @@ class CategoriesTreeVisibilityHandlerImpl {
             return this.getFilteredNodeVisibility({ node });
         }
         if (HierarchyNode.isClassGroupingNode(node)) {
-            return this.getClassGroupingNodeDisplayStatus(node);
+            const nodeInfo = this.getGroupingNodeInfo(node);
+            return this.getGroupedElementsVisibilityStatus({ categoryId: nodeInfo.categoryId, modelElementsMap: nodeInfo.modelElementsMap });
         }
         if (!HierarchyNode.isInstancesNode(node)) {
             return of(createVisibilityStatus("disabled"));
         }
         if (CategoriesTreeNode.isDefinitionContainerNode(node)) {
-            return this.getDefinitionContainerDisplayStatus({
+            return this.getDefinitionContainersVisibilityStatus({
                 definitionContainerIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
             });
         }
         if (CategoriesTreeNode.isModelNode(node)) {
-            return this.getModelVisibilityStatus({
+            return this.getModelsVisibilityStatus({
                 modelIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
             });
         }
         if (CategoriesTreeNode.isCategoryNode(node)) {
-            return this.getCategoryDisplayStatus({
+            return this.getCategoriesVisibilityStatus({
                 categoryIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
                 modelId: CategoriesTreeNode.getModelId(node),
-                ignoreSubCategories: node.extendedData?.isCategoryOfSubModel,
             });
         }
         const categoryId = CategoriesTreeNode.getCategoryId(node);
@@ -128,7 +128,7 @@ class CategoriesTreeVisibilityHandlerImpl {
             return of(createVisibilityStatus("disabled"));
         }
         if (CategoriesTreeNode.isSubCategoryNode(node)) {
-            return this.getSubCategoryDisplayStatus({
+            return this.getSubCategoriesVisibilityStatus({
                 categoryId,
                 subCategoryIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
             });
@@ -137,8 +137,8 @@ class CategoriesTreeVisibilityHandlerImpl {
         if (!modelId) {
             return of(createVisibilityStatus("disabled"));
         }
-        return this.getElementDisplayStatus({
-            elementId: node.key.instanceKeys[0].id,
+        return this.getElementsVisibilityStatus({
+            elementIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
             modelId,
             categoryId,
         });
@@ -147,272 +147,243 @@ class CategoriesTreeVisibilityHandlerImpl {
         return from(this.getVisibilityChangeTargets(props)).pipe(mergeMap(({ definitionContainerIds: definitionContainers, subCategories, modelIds: models, categories, elements }) => {
             const observables = new Array();
             if (definitionContainers?.size) {
-                observables.push(from(definitionContainers).pipe(mergeMap((definitionContainerId) => this.getDefinitionContainerDisplayStatus({ definitionContainerIds: [definitionContainerId] }))));
+                observables.push(from(definitionContainers).pipe(mergeMap((definitionContainerId) => this.getDefinitionContainersVisibilityStatus({ definitionContainerIds: definitionContainerId }))));
             }
             if (models?.size) {
-                observables.push(this.getModelVisibilityStatus({ modelIds: models }));
+                observables.push(this.getModelsVisibilityStatus({ modelIds: models }));
             }
             if (categories?.size) {
                 observables.push(from(categories).pipe(mergeMap((key) => {
                     const { modelId, categoryId } = parseCategoryKey(key);
-                    return this.getCategoryDisplayStatus({ ignoreSubCategories: !!modelId, modelId, categoryIds: [categoryId] });
+                    return this.getCategoriesVisibilityStatus({ modelId, categoryIds: categoryId });
                 })));
             }
             if (subCategories?.size) {
                 observables.push(from(subCategories).pipe(mergeMap((key) => {
                     const { subCategoryId, categoryId } = parseSubCategoryKey(key);
-                    return this.getSubCategoryDisplayStatus({ subCategoryIds: [subCategoryId], categoryId });
+                    return this.getSubCategoriesVisibilityStatus({ subCategoryIds: subCategoryId, categoryId });
                 })));
             }
             if (elements?.size) {
                 observables.push(from(elements).pipe(releaseMainThreadOnItemsCount(50), mergeMap(([categoryKey, elementIds]) => {
                     const { modelId, categoryId } = parseCategoryKey(categoryKey);
                     assert(modelId !== undefined);
-                    return from(elementIds).pipe(releaseMainThreadOnItemsCount(1000), mergeMap((elementId) => this.getElementDisplayStatus({ modelId, categoryId, elementId })));
+                    return from(elementIds).pipe(releaseMainThreadOnItemsCount(1000), mergeMap((elementId) => this.getElementsVisibilityStatus({ modelId, categoryId, elementIds: elementId })));
                 })));
             }
             return merge(...observables);
         }), mergeVisibilityStatuses);
     }
-    getModelVisibilityStatus({ modelIds }) {
+    getModelsVisibilityStatus({ modelIds }) {
         const result = defer(() => {
-            const viewport = this._props.viewport;
-            return from(modelIds).pipe(distinct(), mergeMap((modelId) => {
-                if (!viewport.view.viewsModel(modelId)) {
-                    return from(this._idsCache.getModelCategoryIds(modelId)).pipe(mergeMap((categoryIds) => from(this._idsCache.getCategoriesModeledElements(modelId, categoryIds))), getSubModeledElementsVisibilityStatus({
-                        parentNodeVisibilityStatus: createVisibilityStatus("hidden"),
-                        getModelVisibilityStatus: (modelProps) => this.getModelVisibilityStatus(modelProps),
+            return from(Id64.iterable(modelIds)).pipe(mergeMap((modelId) => {
+                // For hidden models we only need to check subModels
+                if (!this._props.viewport.view.viewsModel(modelId)) {
+                    return this.getSubModels({ modelIds: modelId }).pipe(mergeMap(({ subModels }) => {
+                        if (subModels && Id64.sizeOf(subModels) > 0) {
+                            return this.getModelsVisibilityStatus({ modelIds: subModels }).pipe(map((subModelsVisibilityStatus) => subModelsVisibilityStatus.state !== "hidden" ? createVisibilityStatus("partial") : createVisibilityStatus("hidden")));
+                        }
+                        return of(createVisibilityStatus("hidden"));
                     }));
                 }
-                return from(this._idsCache.getModelCategoryIds(modelId)).pipe(mergeMap((categoryIds) => this.getCategoryDisplayStatus({ modelId, categoryIds, ignoreSubCategories: true })), mergeVisibilityStatuses);
+                // For visible models we need to check all categories
+                return this.getCategories({ modelIds: modelId }).pipe(mergeMap(({ drawingCategories, spatialCategories }) => merge(drawingCategories ? of(drawingCategories).pipe(mergeMap((categoryIds) => this.getCategoriesVisibilityStatus({ modelId, categoryIds }))) : EMPTY, spatialCategories ? of(spatialCategories).pipe(mergeMap((categoryIds) => this.getCategoriesVisibilityStatus({ modelId, categoryIds }))) : EMPTY)), defaultIfEmpty(createVisibilityStatus("visible")));
             }), mergeVisibilityStatuses);
         });
-        return createVisibilityHandlerResult(this, { ids: modelIds }, result, undefined);
+        return createVisibilityHandlerResult(this, { modelIds }, result, undefined);
     }
-    getDefinitionContainerDisplayStatus(props) {
+    getDefinitionContainersVisibilityStatus(props) {
         const result = defer(() => {
-            return from(this._idsCache.getAllContainedCategories(props.definitionContainerIds)).pipe(mergeAll(), mergeMap((categoryId) => this.getCategoryDisplayStatus({ categoryIds: [categoryId] })), mergeVisibilityStatuses);
+            return from(this._idsCache.getAllContainedCategories(props.definitionContainerIds)).pipe(mergeMap((categoryIds) => this.getCategoriesVisibilityStatus({
+                categoryIds,
+                modelId: undefined,
+            })));
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    getSubCategoryDisplayStatus(props) {
-        const result = defer(async () => {
-            const { categoryId, subCategoryIds } = props;
-            let visibility = "unknown";
-            const categoryModels = [...(await this._idsCache.getCategoriesElementModels([categoryId], true)).values()].flat();
-            let nonDefaultModelDisplayStatesCount = 0;
-            for (const modelId of categoryModels) {
-                if (!this._props.viewport.view.viewsModel(modelId)) {
-                    if (visibility === "visible") {
-                        return createVisibilityStatus("partial");
+    getVisibleModelCategoriesVisibilityStatus({ modelId, categoryIds }) {
+        return merge(this.getVisibilityFromAlwaysAndNeverDrawnElements({
+            queryProps: { modelId, categoryIds },
+            defaultStatus: () => this.getVisibleModelDefaultCategoriesVisibilityStatus({ modelId, categoryIds }),
+        }), this.getSubModels({ modelId, categoryIds }).pipe(mergeMap(({ subModels }) => {
+            if (subModels && Id64.sizeOf(subModels) > 0) {
+                return this.getModelsVisibilityStatus({ modelIds: subModels });
+            }
+            return EMPTY;
+        }))).pipe(mergeVisibilityStatuses);
+    }
+    getVisibileCategorySubCategoriesVisibilityStatus(props) {
+        const { subCategoryIds } = props;
+        let subCategoryVisiblity = "unknown";
+        for (const subCategoryId of Id64.iterable(subCategoryIds)) {
+            const isSubCategoryVisible = this._props.viewport.isSubCategoryVisible(subCategoryId);
+            if (isSubCategoryVisible && subCategoryVisiblity === "hidden") {
+                return createVisibilityStatus("partial");
+            }
+            if (!isSubCategoryVisible && subCategoryVisiblity === "visible") {
+                return createVisibilityStatus("partial");
+            }
+            subCategoryVisiblity = isSubCategoryVisible ? "visible" : "hidden";
+        }
+        // If visibility is unknown, no subCategories were provided,
+        // Since category is visible we return visible
+        return createVisibilityStatus(subCategoryVisiblity === "unknown" ? "visible" : subCategoryVisiblity);
+    }
+    getSubCategoriesVisibilityStatus(props) {
+        const result = defer(() => {
+            return (props.modelId ? of({ id: props.categoryId, models: props.modelId }) : from(this.getModels({ categoryIds: props.categoryId }))).pipe(map(({ models }) => {
+                let visibility = "unknown";
+                let nonDefaultModelDisplayStatesCount = 0;
+                for (const modelId of Id64.iterable(models ?? [])) {
+                    if (!this._props.viewport.view.viewsModel(modelId)) {
+                        if (visibility === "visible") {
+                            return createVisibilityStatus("partial");
+                        }
+                        visibility = "hidden";
+                        ++nonDefaultModelDisplayStatesCount;
+                        continue;
                     }
-                    visibility = "hidden";
-                    ++nonDefaultModelDisplayStatesCount;
-                    continue;
-                }
-                const override = this._props.viewport.perModelCategoryVisibility.getOverride(modelId, categoryId);
-                if (override === PerModelCategoryVisibility.Override.Show) {
-                    if (visibility === "hidden") {
-                        return createVisibilityStatus("partial");
+                    const override = this._props.viewport.perModelCategoryVisibility.getOverride(modelId, props.categoryId);
+                    if (override === PerModelCategoryVisibility.Override.Show) {
+                        if (visibility === "hidden") {
+                            return createVisibilityStatus("partial");
+                        }
+                        visibility = "visible";
+                        ++nonDefaultModelDisplayStatesCount;
+                        continue;
                     }
-                    visibility = "visible";
-                    ++nonDefaultModelDisplayStatesCount;
-                    continue;
-                }
-                if (override === PerModelCategoryVisibility.Override.Hide) {
-                    if (visibility === "visible") {
-                        return createVisibilityStatus("partial");
+                    if (override === PerModelCategoryVisibility.Override.Hide) {
+                        if (visibility === "visible") {
+                            return createVisibilityStatus("partial");
+                        }
+                        visibility = "hidden";
+                        ++nonDefaultModelDisplayStatesCount;
+                        continue;
                     }
-                    visibility = "hidden";
-                    ++nonDefaultModelDisplayStatesCount;
-                    continue;
                 }
-            }
-            if (categoryModels.length > 0 && nonDefaultModelDisplayStatesCount === categoryModels.length) {
-                assert(visibility === "visible" || visibility === "hidden");
-                return createVisibilityStatus(visibility);
-            }
-            if (!this._props.viewport.view.viewsCategory(categoryId)) {
-                return createVisibilityStatus(visibility === "visible" ? "partial" : "hidden");
-            }
-            if (subCategoryIds.length === 0) {
-                if (visibility === "hidden") {
-                    return createVisibilityStatus("partial");
+                if (models && Id64.sizeOf(models) > 0 && nonDefaultModelDisplayStatesCount === Id64.sizeOf(models)) {
+                    assert(visibility === "visible" || visibility === "hidden");
+                    return createVisibilityStatus(visibility);
                 }
-                return createVisibilityStatus("visible");
-            }
-            for (const subCategoryId of subCategoryIds) {
-                const isSubCategoryVisible = this._props.viewport.isSubCategoryVisible(subCategoryId);
-                if (isSubCategoryVisible && visibility === "hidden") {
-                    return createVisibilityStatus("partial");
+                if (!this._props.viewport.view.viewsCategory(props.categoryId)) {
+                    return createVisibilityStatus(visibility === "visible" ? "partial" : "hidden");
                 }
-                if (!isSubCategoryVisible && visibility === "visible") {
-                    return createVisibilityStatus("partial");
+                if (Id64.sizeOf(props.subCategoryIds) === 0) {
+                    if (visibility === "hidden") {
+                        return createVisibilityStatus("partial");
+                    }
+                    return createVisibilityStatus("visible");
                 }
-                visibility = isSubCategoryVisible ? "visible" : "hidden";
-            }
-            assert(visibility === "visible" || visibility === "hidden");
-            return createVisibilityStatus(visibility);
+                const subCategoriesVisibility = this.getVisibileCategorySubCategoriesVisibilityStatus({ subCategoryIds: props.subCategoryIds });
+                return subCategoriesVisibility.state === visibility || visibility === "unknown" ? subCategoriesVisibility : createVisibilityStatus("partial");
+            }), mergeVisibilityStatuses);
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    getDefaultModelsCategoryVisibilityStatus({ modelId, categoryIds }) {
+    getVisibleModelDefaultCategoriesVisibilityStatus({ modelId, categoryIds }) {
         const viewport = this._props.viewport;
-        if (!viewport.view.viewsModel(modelId)) {
-            return createVisibilityStatus("hidden");
-        }
         let visibleCount = 0;
-        let hiddenCount = 0;
-        let visibleThroughCategorySelectorCount = 0;
-        for (const categoryId of categoryIds) {
-            if (viewport.view.viewsCategory(categoryId)) {
-                ++visibleThroughCategorySelectorCount;
-            }
+        for (const categoryId of Id64.iterable(categoryIds)) {
             const override = this._props.viewport.perModelCategoryVisibility.getOverride(modelId, categoryId);
-            if (override === PerModelCategoryVisibility.Override.Show) {
+            if (override === PerModelCategoryVisibility.Override.Show ||
+                (override === PerModelCategoryVisibility.Override.None && viewport.view.viewsCategory(categoryId))) {
                 ++visibleCount;
                 continue;
             }
-            if (override === PerModelCategoryVisibility.Override.Hide) {
-                ++hiddenCount;
-                continue;
-            }
-            if (visibleCount > 0 && hiddenCount > 0) {
+            if (visibleCount > 0) {
                 return createVisibilityStatus("partial");
             }
         }
-        if (hiddenCount + visibleCount > 0) {
-            return createVisibilityStatus(hiddenCount > 0 ? "hidden" : "visible");
-        }
-        return createVisibilityStatus(visibleThroughCategorySelectorCount > 0 ? "visible" : "hidden");
+        return visibleCount > 0 ? createVisibilityStatus("visible") : createVisibilityStatus("hidden");
     }
-    getDefaultCategoryVisibilityStatus({ categoryIds }) {
-        return from(this._idsCache.getSubCategories(categoryIds)).pipe(mergeMap((categoriesSubCategoriesMap) => {
-            return from(categoryIds).pipe(mergeMap((categoryId) => {
-                const subCategoryIds = categoriesSubCategoriesMap.get(categoryId);
-                return this.getSubCategoryDisplayStatus({ subCategoryIds: subCategoryIds ?? [], categoryId });
-            }), mergeVisibilityStatuses);
-        }));
-    }
-    getCategoryDisplayStatus(props) {
+    getCategoriesVisibilityStatus({ categoryIds, modelId: modelIdFromProps, }) {
         const result = defer(() => {
-            if (!this._props.hierarchyConfig.showElements) {
-                return from(this.getDefaultCategoryVisibilityStatus({ categoryIds: props.categoryIds }));
-            }
-            const modelsObservable = props.modelId
-                ? of(new Map(props.categoryIds.map((id) => [id, [props.modelId]])))
-                : from(this._idsCache.getCategoriesElementModels(props.categoryIds));
-            return merge(
-            // get visibility status from always and never drawn elements
-            this._props.hierarchyConfig.showElements
-                ? modelsObservable.pipe(mergeMap((categoryModelsMap) => {
-                    if (categoryModelsMap.size === 0) {
-                        return props.modelId
-                            ? of(this.getDefaultModelsCategoryVisibilityStatus({ modelId: props.modelId, categoryIds: props.categoryIds }))
-                            : from(this.getDefaultCategoryVisibilityStatus({ categoryIds: props.categoryIds }));
+            return (modelIdFromProps
+                ? from(Id64.iterable(categoryIds)).pipe(map((categoryId) => ({ id: categoryId, models: modelIdFromProps })))
+                : this.getModels({ categoryIds })).pipe(map(({ id, models }) => {
+                const acc = { categoryId: id, visibleModels: new Array(), hiddenModels: new Array() };
+                if (!models) {
+                    return acc;
+                }
+                for (const modelId of Id64.iterable(models)) {
+                    if (this._props.viewport.viewsModel(modelId)) {
+                        acc.visibleModels.push(modelId);
                     }
-                    return from(categoryModelsMap).pipe(mergeMap(([category, models]) => from(models).pipe(mergeMap((model) => {
-                        if (this._props.viewport.view.viewsModel(model)) {
-                            return this.getVisibilityFromAlwaysAndNeverDrawnElements({
-                                queryProps: props,
-                                defaultStatus: () => this.getDefaultModelsCategoryVisibilityStatus({ modelId: model, categoryIds: [category] }),
-                            }).pipe(mergeMap((visibilityStatusAlwaysAndNeverDraw) => {
-                                return from(this._idsCache.getCategoriesModeledElements(model, [category])).pipe(getSubModeledElementsVisibilityStatus({
-                                    parentNodeVisibilityStatus: visibilityStatusAlwaysAndNeverDraw,
-                                    getModelVisibilityStatus: (modelProps) => this.getModelVisibilityStatus(modelProps),
-                                }));
-                            }));
-                        }
-                        return from(this._idsCache.getCategoriesModeledElements(model, [category])).pipe(getSubModeledElementsVisibilityStatus({
-                            parentNodeVisibilityStatus: createVisibilityStatus("hidden"),
-                            getModelVisibilityStatus: (modelProps) => this.getModelVisibilityStatus(modelProps),
-                        }));
-                    }), mergeVisibilityStatuses)), mergeVisibilityStatuses);
-                }), map((visibilityStatus) => {
-                    return { visibilityStatus, type: 0 };
-                }))
-                : EMPTY, 
-            // get category status
-            (props.modelId
-                ? of(this.getDefaultModelsCategoryVisibilityStatus({ modelId: props.modelId, categoryIds: props.categoryIds }))
-                : from(this.getDefaultCategoryVisibilityStatus({ categoryIds: props.categoryIds }))).pipe(map((visibilityStatus) => {
-                return { visibilityStatus, type: 1 };
-            }))).pipe(toArray(), mergeMap(async (visibilityStatusesInfo) => {
-                let defaultStatus;
-                let alwaysNeverDrawStatus;
-                visibilityStatusesInfo.forEach((visibilityStatusInfo) => {
-                    switch (visibilityStatusInfo.type) {
-                        case 0:
-                            alwaysNeverDrawStatus = visibilityStatusInfo.visibilityStatus;
-                            break;
-                        case 1:
-                            defaultStatus = visibilityStatusInfo.visibilityStatus;
-                            break;
+                    else {
+                        acc.hiddenModels.push(modelId);
                     }
-                });
-                assert(defaultStatus !== undefined);
-                if (defaultStatus.state === "partial") {
-                    return defaultStatus;
-                }
-                // This can happen if:
-                // a) showElements is set to false
-                // b) root category does not have any elements (that dont have Parent)
-                // In both cases we don't need to look at modeled elements visibility
-                if (alwaysNeverDrawStatus === undefined) {
-                    return defaultStatus;
                 }
-                // In cases where
-                // a) SubCategories are hidden
-                // b) Category has model (it means that category is under hidden subModel)
-                // We dont need to look at default category status, it is already accounted for in always/never drawn visibility
-                if (this._props.hierarchyConfig.hideSubCategories || props.modelId) {
-                    return alwaysNeverDrawStatus;
-                }
-                if ((await this._idsCache.getSubCategories(props.categoryIds)).size === 0) {
-                    return alwaysNeverDrawStatus;
-                }
-                if (alwaysNeverDrawStatus.state === "partial" || alwaysNeverDrawStatus.state !== defaultStatus.state) {
-                    return createVisibilityStatus("partial");
-                }
-                return alwaysNeverDrawStatus;
-            }));
+                return acc;
+            }), mergeMap(({ categoryId, visibleModels, hiddenModels }) => {
+                return merge(
+                // For hidden models we only need to check subModels
+                hiddenModels.length > 0
+                    ? this.getSubModels({ modelIds: hiddenModels }).pipe(mergeMap(({ subModels }) => {
+                        if (subModels && Id64.sizeOf(subModels) > 0) {
+                            return this.getModelsVisibilityStatus({
+                                modelIds: subModels,
+                            }).pipe(map((subModelsVisibilityStatus) => subModelsVisibilityStatus.state !== "hidden" ? createVisibilityStatus("partial") : createVisibilityStatus("hidden")));
+                        }
+                        return of(createVisibilityStatus("hidden"));
+                    }))
+                    : EMPTY, 
+                // For visible models we need to check all categories
+                visibleModels.length > 0
+                    ? from(visibleModels).pipe(mergeMap((modelId) => this.getVisibleModelCategoriesVisibilityStatus({
+                        modelId,
+                        categoryIds: categoryId,
+                    })))
+                    : EMPTY, 
+                // We need to check subCategories as well
+                this.getSubCategories({ categoryIds: categoryId }).pipe(mergeMap(({ subCategories }) => {
+                    if (subCategories && Id64.sizeOf(subCategories) > 0) {
+                        return this.getSubCategoriesVisibilityStatus({ categoryId, modelId: modelIdFromProps, subCategoryIds: subCategories });
+                    }
+                    return EMPTY;
+                }))).pipe(defaultIfEmpty(createVisibilityStatus(this._props.viewport.view.viewsCategory(categoryId) ? "visible" : "hidden")));
+            }), mergeVisibilityStatuses);
         });
-        return createVisibilityHandlerResult(this, props, result, undefined);
+        return createVisibilityHandlerResult(this, { categoryIds, modelId: modelIdFromProps }, result, undefined);
     }
-    getClassGroupingNodeDisplayStatus(node) {
+    getGroupedElementsVisibilityStatus(props) {
         const result = defer(() => {
-            const info = this.getGroupingNodeInfo(node);
-            const { modelElementsMap, categoryId } = info;
-            return from(modelElementsMap).pipe(mergeMap(([modelId, elementIds]) => {
-                if (!this._props.viewport.view.viewsModel(modelId)) {
-                    return of([...elementIds]).pipe(filterSubModeledElementIds({ doesSubModelExist: async (id) => this._idsCache.hasSubModel(id) }), getSubModeledElementsVisibilityStatus({
-                        parentNodeVisibilityStatus: createVisibilityStatus("hidden"),
-                        getModelVisibilityStatus: (modelProps) => this.getModelVisibilityStatus(modelProps),
-                    }));
-                }
-                return this.getVisibilityFromAlwaysAndNeverDrawnElements({
-                    elements: elementIds,
-                    defaultStatus: () => this.getDefaultModelsCategoryVisibilityStatus({ categoryIds: [categoryId], modelId }),
-                }).pipe(mergeMap((visibilityStatusAlwaysAndNeverDraw) => {
-                    return of([...elementIds]).pipe(filterSubModeledElementIds({ doesSubModelExist: async (id) => this._idsCache.hasSubModel(id) }), getSubModeledElementsVisibilityStatus({
-                        parentNodeVisibilityStatus: visibilityStatusAlwaysAndNeverDraw,
-                        getModelVisibilityStatus: (modelProps) => this.getModelVisibilityStatus(modelProps),
-                    }));
-                }));
-            }), mergeVisibilityStatuses);
+            const { modelElementsMap, categoryId } = props;
+            return from(modelElementsMap).pipe(mergeMap(([modelId, elementIds]) => this.getElementsVisibilityStatus({ elementIds, modelId, categoryId })), mergeVisibilityStatuses);
         });
-        return createVisibilityHandlerResult(this, { node }, result, undefined);
+        return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    getElementDisplayStatus(props) {
+    getElementsVisibilityStatus(props) {
         const result = defer(() => {
-            const viewport = this._props.viewport;
-            const { elementId, modelId, categoryId } = props;
-            const viewsModel = viewport.view.viewsModel(modelId);
-            const elementStatus = getElementOverriddenVisibility({
-                elementId,
-                viewport,
-            });
-            return from(this._idsCache.hasSubModel(elementId)).pipe(mergeMap((hasSubModel) => (hasSubModel ? this.getModelVisibilityStatus({ modelIds: [elementId] }) : of(undefined))), map((subModelVisibilityStatus) => getElementVisibility(viewsModel, elementStatus, this.getDefaultModelsCategoryVisibilityStatus({ categoryIds: [categoryId], modelId }), subModelVisibilityStatus)));
+            const { elementIds, modelId, categoryId } = props;
+            // TODO: check child elements that are subModels
+            if (!this._props.viewport.view.viewsModel(modelId)) {
+                return from(elementIds).pipe(mergeMap((elementId) => from(this._idsCache.hasSubModel(elementId)).pipe(mergeMap((isSubModel) => {
+                    if (isSubModel) {
+                        return this.getModelsVisibilityStatus({
+                            modelIds: elementId,
+                        }).pipe(map((subModelVisibilityStatus) => subModelVisibilityStatus.state !== "hidden" ? createVisibilityStatus("partial") : createVisibilityStatus("hidden")));
+                    }
+                    return of(createVisibilityStatus("hidden"));
+                }))), mergeVisibilityStatuses);
+            }
+            // TODO: check child elements
+            // TODO: check child element categories
+            // TODO: check child elements that are subModels
+            return this.getVisibilityFromAlwaysAndNeverDrawnElements({
+                elements: elementIds,
+                defaultStatus: () => this.getVisibleModelDefaultCategoriesVisibilityStatus({ categoryIds: categoryId, modelId }),
+            }).pipe(mergeMap((visibilityStatusAlwaysAndNeverDraw) => {
+                return from(Id64.iterable(elementIds)).pipe(mergeMap((elementId) => from(this._idsCache.hasSubModel(elementId)).pipe(mergeMap((isSubModel) => {
+                    if (isSubModel) {
+                        return this.getModelsVisibilityStatus({
+                            modelIds: elementId,
+                        }).pipe(map((subModelVisibilityStatus) => subModelVisibilityStatus.state !== visibilityStatusAlwaysAndNeverDraw.state
+                            ? createVisibilityStatus("partial")
+                            : visibilityStatusAlwaysAndNeverDraw));
+                    }
+                    return of(visibilityStatusAlwaysAndNeverDraw);
+                }))), mergeVisibilityStatuses);
+            }));
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
@@ -422,29 +393,29 @@ class CategoriesTreeVisibilityHandlerImpl {
             return this.changeFilteredNodeVisibility({ node, on });
         }
         if (HierarchyNode.isClassGroupingNode(node)) {
-            return this.changeElementGroupingNodeState(node, on);
+            const nodeInfo = this.getGroupingNodeInfo(node);
+            return this.changeGroupedElementsVisibilityStatus({ categoryId: nodeInfo.categoryId, modelElementsMap: nodeInfo.modelElementsMap, on });
         }
         if (!HierarchyNode.isInstancesNode(node)) {
             return EMPTY;
         }
         if (CategoriesTreeNode.isDefinitionContainerNode(node)) {
-            return this.changeDefinitionContainerState({
+            return this.changeDefinitionContainersVisibilityStatus({
                 definitionContainerIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
                 on,
             });
         }
         if (CategoriesTreeNode.isModelNode(node)) {
-            return this.changeModelState({
-                ids: node.key.instanceKeys.map(({ id }) => id),
+            return this.changeModelsVisibilityStatus({
+                modelIds: node.key.instanceKeys.map(({ id }) => id),
                 on,
             });
         }
         if (CategoriesTreeNode.isCategoryNode(node)) {
-            return this.changeCategoryState({
+            return this.changeCategoriesVisibilityStatus({
                 categoryIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
                 modelId: CategoriesTreeNode.getModelId(node),
                 on,
-                ignoreSubCategories: node.extendedData?.isCategoryOfSubModel,
             });
         }
         const categoryId = CategoriesTreeNode.getCategoryId(node);
@@ -452,24 +423,18 @@ class CategoriesTreeVisibilityHandlerImpl {
             return EMPTY;
         }
         if (CategoriesTreeNode.isSubCategoryNode(node)) {
-            if (this._props.hierarchyConfig.hideSubCategories) {
-                return EMPTY;
-            }
-            return this.changeSubCategoryState({
+            return this.changeSubCategoriesVisibilityStatus({
                 categoryId,
                 subCategoryIds: node.key.instanceKeys.map((instanceKey) => instanceKey.id),
                 on,
             });
         }
-        if (!this._props.hierarchyConfig.showElements) {
-            return EMPTY;
-        }
         const modelId = CategoriesTreeNode.getModelId(node);
         if (!modelId) {
             return EMPTY;
         }
-        return this.changeElementsState({
-            elementIds: new Set([...node.key.instanceKeys.map(({ id }) => id)]),
+        return this.changeElementsVisibilityStatus({
+            elementIds: node.key.instanceKeys.map(({ id }) => id),
             modelId,
             categoryId,
             on,
@@ -483,63 +448,73 @@ class CategoriesTreeVisibilityHandlerImpl {
         return from(this.getVisibilityChangeTargets(props)).pipe(mergeMap(({ definitionContainerIds: definitionContainers, subCategories, modelIds: models, categories, elements }) => {
             const observables = new Array();
             if (definitionContainers?.size) {
-                observables.push(this.changeDefinitionContainerState({ definitionContainerIds: [...definitionContainers], on }));
+                observables.push(this.changeDefinitionContainersVisibilityStatus({ definitionContainerIds: definitionContainers, on }));
             }
             if (models?.size) {
-                observables.push(this.changeModelState({ ids: models, on }));
+                observables.push(this.changeModelsVisibilityStatus({ modelIds: models, on }));
             }
             if (categories?.size) {
                 observables.push(from(categories).pipe(mergeMap((key) => {
                     const { modelId, categoryId } = parseCategoryKey(key);
-                    return this.changeCategoryState({ modelId, categoryIds: [categoryId], ignoreSubCategories: false, on });
+                    return this.changeCategoriesVisibilityStatus({ modelId, categoryIds: categoryId, on });
                 })));
             }
             if (subCategories?.size) {
                 observables.push(from(subCategories).pipe(mergeMap((key) => {
                     const { subCategoryId, categoryId } = parseSubCategoryKey(key);
-                    return this.changeSubCategoryState({ categoryId, subCategoryIds: [subCategoryId], on });
+                    return this.changeSubCategoriesVisibilityStatus({ categoryId, subCategoryIds: subCategoryId, on });
                 })));
             }
             if (elements?.size) {
                 observables.push(from(elements).pipe(mergeMap(([categoryKey, elementIds]) => {
                     const { modelId, categoryId } = parseCategoryKey(categoryKey);
                     assert(modelId !== undefined);
-                    return this.changeElementsState({ modelId, categoryId, elementIds, on });
+                    return this.changeElementsVisibilityStatus({ modelId, categoryId, elementIds, on });
                 })));
             }
             return merge(...observables);
         }));
     }
-    changeModelState(props) {
-        const { ids, on } = props;
-        if (Id64.sizeOf(ids) === 0) {
+    changeModelsVisibilityStatus(props) {
+        const { modelIds, on } = props;
+        if (Id64.sizeOf(modelIds) === 0) {
             return EMPTY;
         }
         const result = defer(() => {
             const viewport = this._props.viewport;
-            viewport.perModelCategoryVisibility.clearOverrides(ids);
-            const idsObs = from(Id64.iterable(ids));
+            viewport.perModelCategoryVisibility.clearOverrides(modelIds);
             if (!on) {
-                viewport.changeModelDisplay(ids, false);
-                return idsObs.pipe(mergeMap(async (modelId) => ({ modelId, categoryIds: await this._idsCache.getModelCategoryIds(modelId) })), mergeMap(({ modelId, categoryIds }) => from(this._idsCache.getCategoriesModeledElements(modelId, categoryIds))), mergeMap((modeledElementIds) => this.changeModelState({ ids: modeledElementIds, on })));
+                viewport.changeModelDisplay(modelIds, false);
+                return this.getSubModels({ modelIds }).pipe(mergeMap(({ subModels }) => (subModels ? this.changeModelsVisibilityStatus({ modelIds: subModels, on }) : EMPTY)));
             }
-            return concat(from(viewport.addViewedModels(ids)), idsObs.pipe(mergeMap((modelId) => {
-                return from(this._idsCache.getModelCategoryIds(modelId)).pipe(mergeAll(), mergeMap((categoryId) => this.changeCategoryState({ categoryIds: [categoryId], modelId, on, ignoreSubCategories: true })));
+            return concat(from(viewport.addViewedModels(modelIds)), this.getCategories({ modelIds }).pipe(mergeMap(({ id, drawingCategories, spatialCategories }) => {
+                return merge(drawingCategories ? this.changeCategoriesVisibilityStatus({ categoryIds: drawingCategories, modelId: id, on }) : EMPTY, spatialCategories ? this.changeCategoriesVisibilityStatus({ categoryIds: spatialCategories, modelId: id, on }) : EMPTY);
             })));
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    showModelWithoutAnyCategoriesOrElements(modelId) {
+    showModelWithoutAnyCategoriesOrElements(modelId, categoriesToNotOverride) {
         const viewport = this._props.viewport;
         return forkJoin({
-            categories: this._idsCache.getModelCategoryIds(modelId),
-            alwaysDrawnElements: this._alwaysAndNeverDrawnElements.getAlwaysDrawnElements({ modelId }),
-        }).pipe(mergeMap(async ({ categories, alwaysDrawnElements }) => {
+            allModelCategories: this.getCategories({ modelIds: modelId }).pipe(reduce((acc, { drawingCategories, spatialCategories }) => {
+                for (const category of Id64.iterable(drawingCategories ?? [])) {
+                    acc.add(category);
+                }
+                for (const category of Id64.iterable(spatialCategories ?? [])) {
+                    acc.add(category);
+                }
+                return acc;
+            }, new Set())),
+            modelAlwaysDrawnElements: this._alwaysAndNeverDrawnElements.getAlwaysDrawnElements({ modelId }),
+        }).pipe(mergeMap(async ({ allModelCategories, modelAlwaysDrawnElements }) => {
             const alwaysDrawn = this._props.viewport.alwaysDrawn;
-            if (alwaysDrawn && alwaysDrawnElements) {
-                viewport.setAlwaysDrawn(setDifference(alwaysDrawn, alwaysDrawnElements));
+            if (alwaysDrawn && modelAlwaysDrawnElements) {
+                viewport.setAlwaysDrawn(setDifference(alwaysDrawn, modelAlwaysDrawnElements));
             }
-            categories.forEach((categoryId) => {
+            const categoriesToOverride = categoriesToNotOverride
+                ? setDifference(allModelCategories, getSetFromId64Arg(categoriesToNotOverride))
+                : allModelCategories;
+            categoriesToOverride.forEach((categoryId) => {
                 this.changeCategoryStateInViewportAccordingToModelVisibility(modelId, categoryId, false, false);
             });
             await viewport.addViewedModels(modelId);
@@ -560,90 +535,131 @@ class CategoriesTreeVisibilityHandlerImpl {
             viewport.changeCategoryDisplay(categoryId, true, changeSubCategories);
         }
     }
-    changeSubCategoryState(props) {
+    changeSubCategoriesVisibilityStatus(props) {
         const result = defer(() => {
             return concat(
             // make sure parent category and models are enabled
             props.on
-                ? concat(from(enableCategoryDisplay(this._props.viewport, [props.categoryId], props.on, false)), from(this.enableCategoriesElementModelsVisibility([props.categoryId])))
+                ? concat(from(enableCategoryDisplay(this._props.viewport, props.categoryId, props.on, false)), this.enableCategoriesElementModelsVisibility(props.categoryId))
                 : EMPTY, from(props.subCategoryIds).pipe(map((subCategoryId) => enableSubCategoryDisplay(this._props.viewport, subCategoryId, props.on))));
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    async enableCategoriesElementModelsVisibility(categoryIds) {
-        const categoriesModelsMap = await this._idsCache.getCategoriesElementModels(categoryIds, true);
-        const modelIds = [...categoriesModelsMap.values()].flat();
-        const hiddenModels = modelIds.filter((modelId) => !this._props.viewport.view.viewsModel(modelId));
-        if (hiddenModels.length > 0) {
-            this._props.viewport.changeModelDisplay(hiddenModels, true);
-        }
+    enableCategoriesElementModelsVisibility(categoryIds) {
+        return from(this._idsCache.getCategoriesElementModels(categoryIds, true)).pipe(mergeMap((categoriesModelsMap) => categoriesModelsMap.values()), reduce((acc, modelIds) => {
+            modelIds.forEach((modelId) => {
+                if (!this._props.viewport.view.viewsModel(modelId)) {
+                    acc.add(modelId);
+                }
+            });
+            return acc;
+        }, new Set()), map((hiddenModels) => {
+            if (hiddenModels.size > 0) {
+                this._props.viewport.changeModelDisplay(hiddenModels, true);
+            }
+        }));
     }
-    changeDefinitionContainerState(props) {
+    changeDefinitionContainersVisibilityStatus(props) {
         const result = defer(() => {
-            return from(this._idsCache.getAllContainedCategories(props.definitionContainerIds)).pipe(mergeMap((categoryIds) => {
-                return this.changeCategoryState({ categoryIds, on: props.on });
-            }));
+            return from(this._idsCache.getAllContainedCategories(props.definitionContainerIds)).pipe(mergeMap((categoryIds) => this.changeCategoriesVisibilityStatus({ categoryIds, modelId: undefined, on: props.on })));
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    changeCategoryState(props) {
+    changeCategoriesVisibilityStatus(props) {
         const result = defer(() => {
-            const { modelId, categoryIds, on, ignoreSubCategories } = props;
+            const { modelId: modelIdFromProps, categoryIds, on } = props;
             const viewport = this._props.viewport;
-            if (!this._props.hierarchyConfig.showElements) {
-                return concat(from(enableCategoryDisplay(viewport, categoryIds, on, on)), on ? from(this.enableCategoriesElementModelsVisibility(categoryIds)) : EMPTY);
-            }
-            const modelIdsObservable = modelId
-                ? of(new Map(categoryIds.map((id) => [id, [modelId]])))
-                : from(this._idsCache.getCategoriesElementModels(categoryIds, true));
-            return concat(modelId ? EMPTY : from(enableCategoryDisplay(viewport, categoryIds, on, ignoreSubCategories ? false : on)), on
-                ? modelIdsObservable.pipe(mergeMap((categoriesMap) => from(categoriesMap.values())), mergeAll(), filter((modelIdToCheck) => !viewport.view.viewsModel(modelIdToCheck)), mergeMap((modelIdToCheck) => this.showModelWithoutAnyCategoriesOrElements(modelIdToCheck)))
-                : EMPTY, modelIdsObservable.pipe(mergeMap((categoriesMap) => from(categoriesMap.entries())), mergeMap(([categoryId, modelIds]) => {
-                return from(modelIds).pipe(mergeMap((modelOfCategory) => {
-                    this.changeCategoryStateInViewportAccordingToModelVisibility(modelOfCategory, categoryId, on, !!ignoreSubCategories);
-                    return this._alwaysAndNeverDrawnElements.clearAlwaysAndNeverDrawnElements({ modelId: modelOfCategory, categoryIds: [categoryId] });
-                }));
-            })), modelIdsObservable.pipe(mergeMap((categoriesMap) => from(categoriesMap.entries())), mergeMap(([categoryId, modelIds]) => {
-                return from(modelIds).pipe(mergeMap((modelOfCategory) => from(this._idsCache.getCategoriesModeledElements(modelOfCategory, [categoryId])).pipe(mergeMap((modeledElementIds) => this.changeModelState({ ids: modeledElementIds, on })))));
-            })));
+            const modelIdsObservable = (modelIdFromProps
+                ? of(new Map([[modelIdFromProps, getSetFromId64Arg(categoryIds)]]))
+                : this.getModels({ categoryIds }).pipe(reduce((acc, { id, models }) => {
+                    if (!models) {
+                        return acc;
+                    }
+                    for (const modelId of Id64.iterable(models)) {
+                        let entry = acc.get(modelId);
+                        if (!entry) {
+                            entry = new Set();
+                            acc.set(modelId, entry);
+                        }
+                        entry.add(id);
+                    }
+                    return acc;
+                }, new Map()))).pipe(mergeMap((modelCategoriesMap) => modelCategoriesMap.entries()), shareReplay());
+            return concat(
+            // If modelId was provided: add override
+            // If modelId was not provided: change categoryDisplay and remove categories per model overrides
+            modelIdFromProps
+                ? of(viewport.perModelCategoryVisibility.setOverride(modelIdFromProps, categoryIds, on ? PerModelCategoryVisibility.Override.Show : PerModelCategoryVisibility.Override.Hide))
+                : concat(from(enableCategoryDisplay(viewport, categoryIds, on, on)), modelIdsObservable.pipe(map(([modelId, modelCategories]) => {
+                    viewport.perModelCategoryVisibility.setOverride(modelId, modelCategories, PerModelCategoryVisibility.Override.None);
+                }))), 
+            // If categories visibility needs to be turned on, we need to turn on models without turning on unrelated elements or categories for that model
+            on
+                ? modelIdsObservable.pipe(mergeMap(([modelId, categories]) => {
+                    if (!viewport.view.viewsModel(modelId)) {
+                        return this.showModelWithoutAnyCategoriesOrElements(modelId, categories);
+                    }
+                    return EMPTY;
+                }))
+                : EMPTY, this._alwaysAndNeverDrawnElements.clearAlwaysAndNeverDrawnElements({ categoryIds, modelId: modelIdFromProps }), this.getSubModels({ categoryIds, modelId: modelIdFromProps }).pipe(mergeMap(({ subModels }) => (subModels ? this.changeModelsVisibilityStatus({ modelIds: subModels, on }) : EMPTY))));
         });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
-    doChangeElementsState(props) {
-        return defer(() => {
-            const { modelId, categoryId, elementIds, on } = props;
-            const viewport = this._props.viewport;
-            return concat(on && !viewport.view.viewsModel(modelId) ? this.showModelWithoutAnyCategoriesOrElements(modelId) : EMPTY, defer(() => {
-                const categoryVisibility = this.getDefaultModelsCategoryVisibilityStatus({ categoryIds: [categoryId], modelId });
-                const isDisplayedByDefault = categoryVisibility.state === "visible";
-                return this.queueElementsVisibilityChange(elementIds, on, isDisplayedByDefault);
-            }), from(elementIds).pipe(mergeMap(async (elementId) => ({ elementId, isSubModel: await this._idsCache.hasSubModel(elementId) })), filter(({ isSubModel }) => isSubModel), map(({ elementId }) => elementId), toArray(), mergeMap((subModelIds) => this.changeModelState({ ids: subModelIds, on }))));
-        });
-    }
     /**
      * Updates visibility of all grouping node's elements.
      * @see `changeElementState`
      */
-    changeElementGroupingNodeState(node, on) {
-        const info = this.getGroupingNodeInfo(node);
-        const result = from(info.modelElementsMap).pipe(mergeMap(([modelId, elementIds]) => {
-            return this.doChangeElementsState({ modelId, elementIds, categoryId: info.categoryId, on });
-        }));
-        return createVisibilityHandlerResult(this, { node, on }, result, undefined);
+    changeGroupedElementsVisibilityStatus(props) {
+        const result = defer(() => {
+            return from(props.modelElementsMap).pipe(mergeMap(([modelId, elementIds]) => {
+                return this.changeElementsVisibilityStatus({ modelId, elementIds, categoryId: props.categoryId, on: props.on });
+            }));
+        });
+        return createVisibilityHandlerResult(this, props, result, undefined);
     }
     /**
      * Updates visibility of an element and all its child elements by adding them to the always/never drawn list.
      * @note If element is to be enabled and model is hidden, it will be enabled.
      */
-    changeElementsState(props) {
-        const result = this.doChangeElementsState(props);
+    changeElementsVisibilityStatus(props) {
+        const result = defer(() => {
+            const { modelId, categoryId, elementIds, on } = props;
+            const viewport = this._props.viewport;
+            // TODO: change child elements
+            // TODO: change child element categories
+            // TODO: change child subModels
+            return concat(
+            // Change elements state
+            defer(() => {
+                if (!viewport.view.viewsModel(modelId)) {
+                    if (!on) {
+                        return this.queueElementsVisibilityChange(elementIds, on, false);
+                    }
+                    return this.showModelWithoutAnyCategoriesOrElements(modelId).pipe(mergeMap(() => {
+                        const defaultVisibility = this.getVisibleModelDefaultCategoriesVisibilityStatus({ categoryIds: categoryId, modelId });
+                        const displayedByDefault = defaultVisibility.state === "visible";
+                        return this.queueElementsVisibilityChange(elementIds, on, displayedByDefault);
+                    }));
+                }
+                const categoryVisibility = this.getVisibleModelDefaultCategoriesVisibilityStatus({ categoryIds: categoryId, modelId });
+                const isDisplayedByDefault = categoryVisibility.state === "visible";
+                return this.queueElementsVisibilityChange(elementIds, on, isDisplayedByDefault);
+            }), 
+            // Change visibility of elements that are models
+            from(Id64.iterable(elementIds)).pipe(mergeMap((elementId) => from(this._idsCache.hasSubModel(elementId)).pipe(mergeMap((isSubModel) => {
+                if (isSubModel) {
+                    return this.changeModelsVisibilityStatus({ modelIds: elementId, on });
+                }
+                return EMPTY;
+            })))));
+        });
         return createVisibilityHandlerResult(this, props, result, undefined);
     }
     queueElementsVisibilityChange(elementIds, on, visibleByDefault) {
         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), changeElementStateNoChildrenOperator({ on, isDisplayedByDefault: visibleByDefault, viewport: this._props.viewport }), tap({
             next: () => {
@@ -674,18 +690,16 @@ class CategoriesTreeVisibilityHandlerImpl {
         if ("elements" in props) {
             return of(getVisibilityFromAlwaysAndNeverDrawnElementsImpl({
                 ...props,
-                alwaysDrawn: viewport.alwaysDrawn?.size ? setIntersection(props.elements, viewport.alwaysDrawn) : undefined,
-                neverDrawn: viewport.neverDrawn?.size ? setIntersection(props.elements, viewport.neverDrawn) : undefined,
-                totalCount: props.elements.size,
+                alwaysDrawn: viewport.alwaysDrawn?.size ? setIntersection(Id64.iterable(props.elements), viewport.alwaysDrawn) : undefined,
+                neverDrawn: viewport.neverDrawn?.size ? setIntersection(Id64.iterable(props.elements), viewport.neverDrawn) : undefined,
+                totalCount: Id64.sizeOf(props.elements),
                 viewport,
             }));
         }
         const { modelId, categoryIds } = props.queryProps;
-        const totalCount = (modelId ? of(new Map(categoryIds.map((categoryId) => [categoryId, [modelId]]))) : from(this._idsCache.getCategoriesElementModels(categoryIds))).pipe(mergeMap((categoriesMap) => from(categoriesMap)), mergeMap(([categoryId, modelIds]) => {
-            return from(modelIds).pipe(mergeMap((modelOfCategory) => from(this._idsCache.getCategoryElementsCount(modelOfCategory, categoryId))), reduce((acc, specificModelCategoryCount) => {
-                return acc + specificModelCategoryCount;
-            }, 0));
-        }));
+        const totalCount = from(Id64.iterable(categoryIds)).pipe(mergeMap((categoryId) => this.getElementsCount({ modelId, categoryId })), reduce((acc, specificModelCategoryCount) => {
+            return acc + specificModelCategoryCount;
+        }, 0));
         return forkJoin({
             totalCount,
             alwaysDrawn: this._alwaysAndNeverDrawnElements.getAlwaysDrawnElements(props.queryProps),
@@ -698,6 +712,33 @@ class CategoriesTreeVisibilityHandlerImpl {
             });
         }));
     }
+    getCategories(props) {
+        return from(Id64.iterable(props.modelIds)).pipe(mergeMap((modelId) => from(this._props.idsCache.getModelCategoryIds(modelId)).pipe(map((categories) => ({ id: modelId, ...(this._props.viewport.view.is2d() ? { drawingCategories: categories } : { spatialCategories: categories }) })))));
+    }
+    getElementsCount(props) {
+        return from(this._props.idsCache.getCategoryElementsCount(props.modelId, props.categoryId));
+    }
+    getModels(props) {
+        return from(Id64.iterable(props.categoryIds)).pipe(mergeMap((categoryId) => from(this._props.idsCache.getCategoriesElementModels(categoryId, true)).pipe(mergeMap((categoryModelsMap) => (categoryModelsMap.size > 0 ? categoryModelsMap.values() : of(undefined))), map((categoryModels) => ({ id: categoryId, models: categoryModels })))));
+    }
+    getSubCategories(props) {
+        return from(Id64.iterable(props.categoryIds)).pipe(mergeMap((categoryId) => from(this._props.idsCache.getSubCategories(categoryId)).pipe(mergeMap((categorySubCategoriesMap) => (categorySubCategoriesMap.size > 0 ? categorySubCategoriesMap.values() : of(undefined))), map((subCategories) => ({ id: categoryId, subCategories })))));
+    }
+    getSubModels(props) {
+        if ("modelIds" in props) {
+            return from(Id64.iterable(props.modelIds)).pipe(mergeMap((modelId) => from(this._props.idsCache.getModelCategoryIds(modelId)).pipe(mergeMap((categoryIds) => from(this._props.idsCache.getCategoriesModeledElements(modelId, categoryIds))), map((subModels) => ({ id: modelId, subModels })))));
+        }
+        if (props.modelId) {
+            return from(Id64.iterable(props.categoryIds)).pipe(mergeMap((categoryId) => from(this._props.idsCache.getCategoriesModeledElements(props.modelId, categoryId)).pipe(map((subModels) => ({ id: categoryId, subModels })))));
+        }
+        return from(Id64.iterable(props.categoryIds)).pipe(mergeMap((categoryId) => from(this._props.idsCache.getCategoriesElementModels(categoryId)).pipe(mergeMap((categoryModelsMap) => {
+            const models = categoryModelsMap.get(categoryId);
+            if (!models) {
+                return of({ id: categoryId, subModels: undefined });
+            }
+            return from(models).pipe(mergeMap((modelId) => from(this._props.idsCache.getCategoriesModeledElements(modelId, categoryId))), map((subModels) => ({ id: categoryId, subModels })));
+        }))));
+    }
     getGroupingNodeInfo(node) {
         const modelElementsMap = node.extendedData?.modelElementsMap;
         const categoryId = node.extendedData?.categoryId;