@itwin/tree-widget-react 4.0.0-alpha.2 → 4.0.0-alpha.20

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

@@ -2,338 +2,350 @@
  * Copyright (c) Bentley Systems, Incorporated. All rights reserved.
  * See LICENSE.md in the project root for license terms and full copyright notice.
  *--------------------------------------------------------------------------------------------*/
-import { bufferTime, filter, firstValueFrom, mergeAll, mergeMap, ReplaySubject, Subject } from "rxjs";
-import { assert } from "@itwin/core-bentley";
+import { defer, filter, forkJoin, from, map, mergeAll, mergeMap, of, reduce, shareReplay, toArray } from "rxjs";
+import { assert, Guid, Id64 } from "@itwin/core-bentley";
 import { IModel } from "@itwin/core-common";
-import { pushToMap } from "../../common/Utils.js";
+import { CLASS_NAME_GeometricModel3d, CLASS_NAME_InformationPartitionElement, CLASS_NAME_Model, CLASS_NAME_SpatialCategory, CLASS_NAME_Subject, } from "../../common/internal/ClassNameDefinitions.js";
+import { ModelCategoryElementsCountCache } from "../../common/internal/ModelCategoryElementsCountCache.js";
+import { pushToMap } from "../../common/internal/Utils.js";
 /** @internal */
 export class ModelsTreeIdsCache {
-    constructor(_queryExecutor, _hierarchyConfig) {
-        this._queryExecutor = _queryExecutor;
-        this._hierarchyConfig = _hierarchyConfig;
-        this._categoryElementCounts = new ModelCategoryElementsCountCache(async (input) => this.queryCategoryElementCounts(input));
-        this._modelKeyPaths = new Map();
-        this._subjectKeyPaths = new Map();
-        this._categoryKeyPaths = new Map();
+    #categoryElementCounts;
+    #subjectInfos;
+    #parentSubjectIds; // the list should contain a subject id if its node should be shown as having children
+    #modelInfos;
+    #modelWithCategoryModeledElements;
+    #modelKeyPaths;
+    #subjectKeyPaths;
+    #categoryKeyPaths;
+    #queryExecutor;
+    #hierarchyConfig;
+    #componentId;
+    #componentName;
+    constructor(queryExecutor, hierarchyConfig, componentId) {
+        this.#queryExecutor = queryExecutor;
+        this.#hierarchyConfig = hierarchyConfig;
+        this.#componentId = componentId ?? Guid.createValue();
+        this.#componentName = "ModelsTreeIdsCache";
+        this.#categoryElementCounts = new ModelCategoryElementsCountCache(this.#queryExecutor, [this.#hierarchyConfig.elementClassSpecification], this.#componentId);
+        this.#modelKeyPaths = new Map();
+        this.#subjectKeyPaths = new Map();
+        this.#categoryKeyPaths = new Map();
     }
     [Symbol.dispose]() {
-        this._categoryElementCounts[Symbol.dispose]();
+        this.#categoryElementCounts[Symbol.dispose]();
     }
-    async *querySubjects() {
-        const subjectsQuery = `
-      SELECT
-        s.ECInstanceId id,
-        s.Parent.Id parentId,
-        (
-          SELECT m.ECInstanceId
-          FROM bis.GeometricModel3d m
-          WHERE m.ECInstanceId = HexToId(json_extract(s.JsonProperties, '$.Subject.Model.TargetPartition'))
-            AND NOT m.IsPrivate
-            AND EXISTS (SELECT 1 FROM ${this._hierarchyConfig.elementClassSpecification} WHERE Model.Id = m.ECInstanceId)
-        ) targetPartitionId,
-        CASE
-          WHEN (
-            json_extract(s.JsonProperties, '$.Subject.Job.Bridge') IS NOT NULL
-            OR json_extract(s.JsonProperties, '$.Subject.Model.Type') = 'Hierarchy'
-          ) THEN 1
-          ELSE 0
-        END hideInHierarchy
-      FROM bis.Subject s
-    `;
-        for await (const row of this._queryExecutor.createQueryReader({ ecsql: subjectsQuery }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded" })) {
-            yield { id: row.id, parentId: row.parentId, targetPartitionId: row.targetPartitionId, hideInHierarchy: !!row.hideInHierarchy };
-        }
+    querySubjects() {
+        return defer(() => {
+            const subjectsQuery = `
+        SELECT
+          s.ECInstanceId id,
+          s.Parent.Id parentId,
+          (
+            SELECT m.ECInstanceId
+            FROM ${CLASS_NAME_GeometricModel3d} m
+            WHERE m.ECInstanceId = HexToId(json_extract(s.JsonProperties, '$.Subject.Model.TargetPartition'))
+              AND NOT m.IsPrivate
+              AND EXISTS (SELECT 1 FROM ${this.#hierarchyConfig.elementClassSpecification} WHERE Model.Id = m.ECInstanceId)
+          ) targetPartitionId,
+          CASE
+            WHEN (
+              json_extract(s.JsonProperties, '$.Subject.Job.Bridge') IS NOT NULL
+              OR json_extract(s.JsonProperties, '$.Subject.Model.Type') = 'Hierarchy'
+            ) THEN 1
+            ELSE 0
+          END hideInHierarchy
+        FROM bis.Subject s
+      `;
+            return this.#queryExecutor.createQueryReader({ ecsql: subjectsQuery }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded", restartToken: `${this.#componentName}/${this.#componentId}/subjects` });
+        }).pipe(map((row) => {
+            return { id: row.id, parentId: row.parentId, targetPartitionId: row.targetPartitionId, hideInHierarchy: !!row.hideInHierarchy };
+        }));
     }
-    async *queryModels() {
-        const modelsQuery = `
-      SELECT p.ECInstanceId id, p.Parent.Id parentId
-      FROM bis.InformationPartitionElement p
-      INNER JOIN bis.GeometricModel3d m ON m.ModeledElement.Id = p.ECInstanceId
-      WHERE
-        NOT m.IsPrivate
-        ${this._hierarchyConfig.showEmptyModels ? "" : `AND EXISTS (SELECT 1 FROM ${this._hierarchyConfig.elementClassSpecification} WHERE Model.Id = m.ECInstanceId)`}
-    `;
-        for await (const row of this._queryExecutor.createQueryReader({ ecsql: modelsQuery }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded" })) {
-            yield { id: row.id, parentId: row.parentId };
-        }
+    queryModels() {
+        return defer(() => {
+            const modelsQuery = `
+        SELECT p.ECInstanceId id, p.Parent.Id parentId
+        FROM ${CLASS_NAME_InformationPartitionElement} p
+        INNER JOIN ${CLASS_NAME_GeometricModel3d} m ON m.ModeledElement.Id = p.ECInstanceId
+        WHERE
+          NOT m.IsPrivate
+          ${this.#hierarchyConfig.showEmptyModels ? "" : `AND EXISTS (SELECT 1 FROM ${this.#hierarchyConfig.elementClassSpecification} WHERE Model.Id = m.ECInstanceId)`}
+      `;
+            return this.#queryExecutor.createQueryReader({ ecsql: modelsQuery }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded", restartToken: `${this.#componentName}/${this.#componentId}/models` });
+        }).pipe(map((row) => {
+            return { id: row.id, parentId: row.parentId };
+        }));
     }
-    async getSubjectInfos() {
-        this._subjectInfos ??= (async () => {
-            const [subjectInfos, targetPartitionSubjects] = await Promise.all([
-                (async () => {
-                    const result = new Map();
-                    for await (const subject of this.querySubjects()) {
-                        const subjectInfo = {
-                            parentSubject: subject.parentId,
-                            hideInHierarchy: subject.hideInHierarchy,
-                            childSubjects: new Set(),
-                            childModels: new Set(),
-                        };
-                        if (subject.targetPartitionId) {
-                            subjectInfo.childModels.add(subject.targetPartitionId);
-                        }
-                        result.set(subject.id, subjectInfo);
-                    }
-                    return result;
-                })(),
-                (async () => {
-                    const result = new Map();
-                    for await (const model of this.queryModels()) {
-                        pushToMap(result, model.id, model.parentId);
+    getSubjectInfos() {
+        this.#subjectInfos ??= forkJoin({
+            subjectInfos: this.querySubjects().pipe(reduce((acc, subject) => {
+                const subjectInfo = {
+                    parentSubjectId: subject.parentId,
+                    hideInHierarchy: subject.hideInHierarchy,
+                    childSubjectIds: new Set(),
+                    childModelIds: new Set(),
+                };
+                if (subject.targetPartitionId) {
+                    subjectInfo.childModelIds.add(subject.targetPartitionId);
+                }
+                acc.set(subject.id, subjectInfo);
+                return acc;
+            }, new Map()), map((subjectInfos) => {
+                for (const [subjectId, { parentSubjectId: parentSubjectId }] of subjectInfos.entries()) {
+                    if (parentSubjectId) {
+                        const parentSubjectInfo = subjectInfos.get(parentSubjectId);
+                        assert(!!parentSubjectInfo);
+                        parentSubjectInfo.childSubjectIds.add(subjectId);
                     }
-                    return result;
-                })(),
-            ]);
-            for (const [subjectId, { parentSubject: parentSubjectId }] of subjectInfos.entries()) {
-                if (parentSubjectId) {
-                    const parentSubjectInfo = subjectInfos.get(parentSubjectId);
-                    assert(!!parentSubjectInfo);
-                    parentSubjectInfo.childSubjects.add(subjectId);
                 }
-            }
+                return subjectInfos;
+            })),
+            targetPartitionSubjects: this.queryModels().pipe(reduce((acc, model) => {
+                pushToMap(acc, model.id, model.parentId);
+                return acc;
+            }, new Map())),
+        }).pipe(map(({ subjectInfos, targetPartitionSubjects }) => {
             for (const [partitionId, subjectIds] of targetPartitionSubjects) {
                 subjectIds.forEach((subjectId) => {
                     const subjectInfo = subjectInfos.get(subjectId);
                     assert(!!subjectInfo);
-                    subjectInfo.childModels.add(partitionId);
+                    subjectInfo.childModelIds.add(partitionId);
                 });
             }
             return subjectInfos;
-        })();
-        return this._subjectInfos;
+        }), shareReplay());
+        return this.#subjectInfos;
     }
     /** Returns ECInstanceIDs of Subjects that either have direct Model or at least one child Subject with a Model. */
-    async getParentSubjectIds() {
-        this._parentSubjectIds ??= (async () => {
-            const subjectInfos = await this.getSubjectInfos();
+    getParentSubjectIds() {
+        this.#parentSubjectIds ??= this.getSubjectInfos().pipe(map((subjectInfos) => {
             const parentSubjectIds = new Set();
             subjectInfos.forEach((subjectInfo, subjectId) => {
-                if (subjectInfo.childModels.size > 0) {
+                if (subjectInfo.childModelIds.size > 0) {
                     parentSubjectIds.add(subjectId);
-                    let currParentId = subjectInfo.parentSubject;
+                    let currParentId = subjectInfo.parentSubjectId;
                     while (currParentId) {
                         parentSubjectIds.add(currParentId);
-                        currParentId = subjectInfos.get(currParentId)?.parentSubject;
+                        currParentId = subjectInfos.get(currParentId)?.parentSubjectId;
                     }
                 }
             });
             return [...parentSubjectIds];
-        })();
-        return this._parentSubjectIds;
+        }), shareReplay());
+        return this.#parentSubjectIds;
     }
     /**
      * Returns child subjects of the specified parent subjects as they're displayed in the hierarchy - taking into
      * account `hideInHierarchy` flag.
      */
-    async getChildSubjectIds(parentSubjectIds) {
-        const childSubjectIds = new Array();
-        const subjectInfos = await this.getSubjectInfos();
-        parentSubjectIds.forEach((subjectId) => {
-            forEachChildSubject(subjectInfos, subjectId, (childSubjectId, childSubjectInfo) => {
-                if (!childSubjectInfo.hideInHierarchy) {
-                    childSubjectIds.push(childSubjectId);
-                    return "break";
-                }
-                return "continue";
-            });
-        });
-        return childSubjectIds;
+    getChildSubjectIds(parentSubjectIds) {
+        return this.getSubjectInfos().pipe(map((subjectInfos) => {
+            const childSubjectIds = new Array();
+            for (const subjectId of Id64.iterable(parentSubjectIds)) {
+                forEachChildSubject(subjectInfos, subjectId, (childSubjectId, childSubjectInfo) => {
+                    if (!childSubjectInfo.hideInHierarchy) {
+                        childSubjectIds.push(childSubjectId);
+                        return "break";
+                    }
+                    return "continue";
+                });
+            }
+            return childSubjectIds;
+        }));
     }
     /** Returns ECInstanceIDs of all Models under specific parent Subjects, including their child Subjects, etc. */
-    async getSubjectModelIds(subjectIds) {
-        const subjectInfos = await this.getSubjectInfos();
-        const subjectStack = [...subjectIds];
-        const result = new Array();
-        while (true) {
-            const subjectId = subjectStack.pop();
-            if (subjectId === undefined) {
-                break;
+    getSubjectModelIds(subjectIds) {
+        return this.getSubjectInfos().pipe(mergeMap((subjectInfos) => {
+            const result = new Array();
+            const childSubjects = new Array();
+            for (const subjectId of Id64.iterable(subjectIds)) {
+                const subjectInfo = subjectInfos.get(subjectId);
+                if (!subjectInfo) {
+                    continue;
+                }
+                result.push(...subjectInfo.childModelIds);
+                childSubjects.push(...subjectInfo.childSubjectIds);
             }
-            const subjectInfo = subjectInfos.get(subjectId);
-            if (!subjectInfo) {
-                continue;
+            if (childSubjects.length === 0) {
+                return of(result);
             }
-            result.push(...subjectInfo.childModels);
-            subjectStack.push(...subjectInfo.childSubjects);
-        }
-        return result;
+            return this.getSubjectModelIds(childSubjects).pipe(map((modelsOfChildSubjects) => {
+                result.push(...modelsOfChildSubjects);
+                return result;
+            }));
+        }));
     }
     /** Returns ECInstanceIDs of Models under specific parent Subjects as they are displayed in the hierarchy. */
-    async getChildSubjectModelIds(parentSubjectIds) {
-        const subjectInfos = await this.getSubjectInfos();
-        const hiddenSubjectIds = new Array();
-        parentSubjectIds.forEach((subjectId) => {
-            forEachChildSubject(subjectInfos, subjectId, (childSubjectId, childSubjectInfo) => {
-                if (childSubjectInfo.hideInHierarchy) {
-                    hiddenSubjectIds.push(childSubjectId);
-                    return "continue";
-                }
-                return "break";
+    getChildSubjectModelIds(parentSubjectIds) {
+        return this.getSubjectInfos().pipe(map((subjectInfos) => {
+            const hiddenSubjectIds = new Array();
+            for (const subjectId of Id64.iterable(parentSubjectIds)) {
+                forEachChildSubject(subjectInfos, subjectId, (childSubjectId, childSubjectInfo) => {
+                    if (childSubjectInfo.hideInHierarchy) {
+                        hiddenSubjectIds.push(childSubjectId);
+                        return "continue";
+                    }
+                    return "break";
+                });
+            }
+            const modelIds = new Array();
+            [...parentSubjectIds, ...hiddenSubjectIds].forEach((subjectId) => {
+                const subjectInfo = subjectInfos.get(subjectId);
+                subjectInfo && modelIds.push(...subjectInfo.childModelIds);
             });
-        });
-        const modelIds = new Array();
-        [...parentSubjectIds, ...hiddenSubjectIds].forEach((subjectId) => {
-            const subjectInfo = subjectInfos.get(subjectId);
-            subjectInfo && modelIds.push(...subjectInfo.childModels);
-        });
-        return modelIds;
+            return modelIds;
+        }));
     }
-    async createSubjectInstanceKeysPath(targetSubjectId) {
-        let entry = this._subjectKeyPaths.get(targetSubjectId);
+    createSubjectInstanceKeysPath(targetSubjectId) {
+        let entry = this.#subjectKeyPaths.get(targetSubjectId);
         if (!entry) {
-            entry = (async () => {
-                const subjectInfos = await this.getSubjectInfos();
+            entry = this.getSubjectInfos().pipe(map((subjectInfos) => {
                 const result = new Array();
                 let currParentId = targetSubjectId;
-                while (currParentId && currParentId !== IModel.rootSubjectId) {
+                while (currParentId) {
+                    if (this.#hierarchyConfig.hideRootSubject && currParentId === IModel.rootSubjectId) {
+                        break;
+                    }
                     const parentInfo = subjectInfos.get(currParentId);
                     if (!parentInfo?.hideInHierarchy) {
-                        result.push({ className: "BisCore.Subject", id: currParentId });
+                        result.push({ className: CLASS_NAME_Subject, id: currParentId });
                     }
-                    currParentId = parentInfo?.parentSubject;
+                    currParentId = parentInfo?.parentSubjectId;
                 }
                 return result.reverse();
-            })();
-            this._subjectKeyPaths.set(targetSubjectId, entry);
+            }), shareReplay());
+            this.#subjectKeyPaths.set(targetSubjectId, entry);
         }
         return entry;
     }
-    async *queryModelElementCounts() {
-        const query = /* sql */ `
-      SELECT Model.Id modelId, COUNT(*) elementCount
-      FROM ${this._hierarchyConfig.elementClassSpecification}
-      GROUP BY Model.Id
-    `;
-        for await (const row of this._queryExecutor.createQueryReader({ ecsql: query }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded" })) {
-            yield { modelId: row.modelId, elementCount: row.elementCount };
-        }
+    queryModelCategories() {
+        return defer(() => {
+            const query = `
+        SELECT
+          this.Model.Id modelId,
+          this.Category.Id categoryId,
+          m.IsPrivate isModelPrivate,
+          MAX(IIF(this.Parent.Id IS NULL, 1, 0)) isRootElementCategory
+        FROM ${CLASS_NAME_Model} m
+        JOIN ${this.#hierarchyConfig.elementClassSpecification} this ON m.ECInstanceId = this.Model.Id
+        GROUP BY modelId, categoryId, isModelPrivate
+      `;
+            return this.#queryExecutor.createQueryReader({ ecsql: query }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded", restartToken: `${this.#componentName}/${this.#componentId}/model-categories` });
+        }).pipe(map((row) => {
+            return { modelId: row.modelId, categoryId: row.categoryId, isModelPrivate: !!row.isModelPrivate, isRootElementCategory: !!row.isRootElementCategory };
+        }));
     }
-    async *queryModelCategories() {
-        const query = /* sql */ `
-      SELECT Model.Id modelId, Category.Id categoryId
-      FROM ${this._hierarchyConfig.elementClassSpecification}
-      WHERE Parent.Id IS NULL
-      GROUP BY modelId, categoryId
+    queryModeledElements() {
+        return defer(() => {
+            const query = `
+      SELECT
+        pe.ECInstanceId modeledElementId,
+        pe.Category.Id categoryId,
+        pe.Model.Id modelId
+      FROM ${CLASS_NAME_Model} m
+      JOIN ${this.#hierarchyConfig.elementClassSpecification} pe ON pe.ECInstanceId = m.ModeledElement.Id
+      WHERE
+        m.IsPrivate = false
+        AND m.ECInstanceId IN (SELECT Model.Id FROM ${this.#hierarchyConfig.elementClassSpecification})
     `;
-        for await (const row of this._queryExecutor.createQueryReader({ ecsql: query }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded" })) {
-            yield { modelId: row.modelId, categoryId: row.categoryId };
-        }
+            return this.#queryExecutor.createQueryReader({ ecsql: query }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded", restartToken: `${this.#componentName}/${this.#componentId}/modeled-elements` });
+        }).pipe(map((row) => {
+            return { modelId: row.modelId, categoryId: row.categoryId, modeledElementId: row.modeledElementId };
+        }));
     }
-    async getModelInfos() {
-        this._modelInfos ??= (async () => {
-            const modelInfos = new Map();
-            await Promise.all([
-                (async () => {
-                    for await (const { modelId, categoryId } of this.queryModelCategories()) {
-                        const entry = modelInfos.get(modelId);
-                        if (entry) {
-                            entry.categories.add(categoryId);
-                        }
-                        else {
-                            modelInfos.set(modelId, { categories: new Set([categoryId]), elementCount: 0 });
-                        }
-                    }
-                })(),
-                (async () => {
-                    for await (const { modelId, elementCount } of this.queryModelElementCounts()) {
-                        const entry = modelInfos.get(modelId);
-                        if (entry) {
-                            entry.elementCount = elementCount;
-                        }
-                        else {
-                            modelInfos.set(modelId, { categories: new Set(), elementCount });
-                        }
-                    }
-                })(),
-            ]);
-            return modelInfos;
-        })();
-        return this._modelInfos;
+    getModelWithCategoryModeledElements() {
+        this.#modelWithCategoryModeledElements ??= this.queryModeledElements().pipe(reduce((acc, { modelId, categoryId, modeledElementId }) => {
+            const key = `${modelId}-${categoryId}`;
+            const entry = acc.get(key);
+            if (entry === undefined) {
+                acc.set(key, new Set([modeledElementId]));
+            }
+            else {
+                entry.add(modeledElementId);
+            }
+            return acc;
+        }, new Map()), shareReplay());
+        return this.#modelWithCategoryModeledElements;
+    }
+    getModelInfos() {
+        this.#modelInfos ??= this.queryModelCategories().pipe(reduce((acc, { modelId, categoryId, isModelPrivate, isRootElementCategory }) => {
+            const entry = acc.get(modelId);
+            if (entry) {
+                entry.categories.set(categoryId, { isRootElementCategory });
+                entry.isModelPrivate = isModelPrivate;
+            }
+            else {
+                acc.set(modelId, { categories: new Map([[categoryId, { isRootElementCategory }]]), isModelPrivate });
+            }
+            return acc;
+        }, new Map()), shareReplay());
+        return this.#modelInfos;
     }
-    async getModelCategories(modelId) {
-        const modelInfos = await this.getModelInfos();
-        const categories = modelInfos.get(modelId)?.categories;
-        return categories ? [...categories] : [];
+    getAllCategories() {
+        return this.getModelInfos().pipe(mergeMap((modelInfos) => modelInfos.values()), mergeMap(({ categories }) => categories.keys()), reduce((acc, categoryId) => {
+            acc.add(categoryId);
+            return acc;
+        }, new Set()));
     }
-    async getModelElementCount(modelId) {
-        const modelInfos = await this.getModelInfos();
-        return modelInfos.get(modelId)?.elementCount ?? 0;
+    getModelCategoryIds(modelId) {
+        return this.getModelInfos().pipe(mergeMap((modelInfos) => modelInfos.get(modelId)?.categories.keys() ?? []), toArray());
     }
-    async createModelInstanceKeyPaths(modelId) {
-        let entry = this._modelKeyPaths.get(modelId);
-        if (!entry) {
-            entry = (async () => {
-                const result = new Array();
-                const subjectInfos = (await this.getSubjectInfos()).entries();
-                for (const [modelSubjectId, subjectInfo] of subjectInfos) {
-                    if (subjectInfo.childModels.has(modelId)) {
-                        const subjectPath = await this.createSubjectInstanceKeysPath(modelSubjectId);
-                        result.push([...subjectPath, { className: "BisCore.GeometricModel3d", id: modelId }]);
-                    }
+    hasSubModel(elementId) {
+        return this.getModelInfos().pipe(map((modelInfos) => {
+            const modeledElementInfo = modelInfos.get(elementId);
+            if (!modeledElementInfo) {
+                return false;
+            }
+            return !modeledElementInfo.isModelPrivate;
+        }));
+    }
+    getCategoriesModeledElements(modelId, categoryIds) {
+        return this.getModelWithCategoryModeledElements().pipe(mergeMap((modelWithCategoryModeledElements) => from(Id64.iterable(categoryIds)).pipe(reduce((acc, categoryId) => {
+            const entry = modelWithCategoryModeledElements.get(`${modelId}-${categoryId}`);
+            if (entry !== undefined) {
+                acc.push(...entry);
+            }
+            return acc;
+        }, new Array()))));
+    }
+    getCategoriesElementModels(categoryIds) {
+        return this.getModelInfos().pipe(mergeMap((modelInfos) => from(Id64.iterable(categoryIds)).pipe(reduce((acc, categoryId) => {
+            const entry = new Array();
+            for (const [modelId, { categories }] of modelInfos.entries()) {
+                if (categories.has(categoryId)) {
+                    entry.push(modelId);
                 }
-                return result;
-            })();
-            this._modelKeyPaths.set(modelId, entry);
-        }
-        return entry;
+            }
+            acc.set(categoryId, entry);
+            return acc;
+        }, new Map()))));
     }
-    async queryCategoryElementCounts(input) {
-        const reader = this._queryExecutor.createQueryReader({
-            ctes: [
-                /* sql */ `
-            CategoryElements(id, modelId, categoryId) AS (
-              SELECT ECInstanceId, Model.Id, Category.Id
-              FROM ${this._hierarchyConfig.elementClassSpecification}
-              WHERE
-                Parent.Id IS NULL
-                AND (
-                  ${input.map(({ modelId, categoryId }) => `Model.Id = ${modelId} AND Category.Id = ${categoryId}`).join(" OR ")}
-                )
-
-              UNION ALL
-
-              SELECT c.ECInstanceId, p.modelId, p.categoryId
-              FROM ${this._hierarchyConfig.elementClassSpecification} c
-              JOIN CategoryElements p ON c.Parent.Id = p.id
-            )
-          `,
-            ],
-            ecsql: `
-          SELECT modelId, categoryId, COUNT(id) elementsCount
-          FROM CategoryElements
-          GROUP BY modelId, categoryId
-        `,
-        }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded" });
-        const result = new Array();
-        for await (const row of reader) {
-            result.push({ modelId: row.modelId, categoryId: row.categoryId, elementsCount: row.elementsCount });
+    createModelInstanceKeyPaths(modelId) {
+        let entry = this.#modelKeyPaths.get(modelId);
+        if (!entry) {
+            entry = this.getSubjectInfos().pipe(mergeMap((subjectInfos) => subjectInfos.entries()), filter(([_, subjectInfo]) => subjectInfo.childModelIds.has(modelId)), mergeMap(([modelSubjectId]) => this.createSubjectInstanceKeysPath(modelSubjectId).pipe(map((path) => [...path, { className: CLASS_NAME_GeometricModel3d, id: modelId }]))), toArray(), shareReplay());
+            this.#modelKeyPaths.set(modelId, entry);
         }
-        return result;
+        return entry;
     }
-    async getCategoryElementsCount(modelId, categoryId) {
-        return this._categoryElementCounts.getCategoryElementsCount(modelId, categoryId);
+    getCategoryElementsCount(modelId, categoryId) {
+        return this.#categoryElementCounts.getCategoryElementsCount(modelId, categoryId);
     }
-    async createCategoryInstanceKeyPaths(categoryId) {
-        let entry = this._categoryKeyPaths.get(categoryId);
+    createCategoryInstanceKeyPaths(categoryId) {
+        let entry = this.#categoryKeyPaths.get(categoryId);
         if (!entry) {
-            entry = (async () => {
-                const result = new Set();
-                const modelInfos = await this.getModelInfos();
-                modelInfos?.forEach((modelInfo, modelId) => {
-                    if (modelInfo.categories.has(categoryId)) {
-                        result.add(modelId);
-                    }
-                });
-                const categoryPaths = new Array();
-                for (const categoryModelId of [...result]) {
-                    const modelPaths = await this.createModelInstanceKeyPaths(categoryModelId);
-                    for (const modelPath of modelPaths) {
-                        categoryPaths.push([...modelPath, { className: "BisCore.SpatialCategory", id: categoryId }]);
-                    }
+            entry = this.getModelInfos().pipe(mergeMap((modelInfos) => modelInfos.entries()), reduce((acc, [modelId, modelInfo]) => {
+                const categoryEntry = modelInfo.categories.get(categoryId);
+                if (categoryEntry?.isRootElementCategory) {
+                    acc.add(modelId);
                 }
-                return categoryPaths;
-            })();
-            this._categoryKeyPaths.set(categoryId, entry);
+                return acc;
+            }, new Set()), mergeMap((categoryModels) => from(categoryModels)), mergeMap((categoryModelId) => this.createModelInstanceKeyPaths(categoryModelId)), mergeAll(), reduce((acc, modelPath) => {
+                acc.push([...modelPath, { className: CLASS_NAME_SpatialCategory, id: categoryId }]);
+                return acc;
+            }, new Array()), shareReplay());
+            this.#categoryKeyPaths.set(categoryId, entry);
         }
         return entry;
     }
@@ -341,7 +353,7 @@ export class ModelsTreeIdsCache {
 function forEachChildSubject(subjectInfos, parentSubject, cb) {
     const parentSubjectInfo = typeof parentSubject === "string" ? subjectInfos.get(parentSubject) : parentSubject;
     parentSubjectInfo &&
-        parentSubjectInfo.childSubjects.forEach((childSubjectId) => {
+        parentSubjectInfo.childSubjectIds.forEach((childSubjectId) => {
             const childSubjectInfo = subjectInfos.get(childSubjectId);
             if (cb(childSubjectId, childSubjectInfo) === "break") {
                 return;
@@ -349,34 +361,4 @@ function forEachChildSubject(subjectInfos, parentSubject, cb) {
             forEachChildSubject(subjectInfos, childSubjectInfo, cb);
         });
 }
-class ModelCategoryElementsCountCache {
-    constructor(_loader) {
-        this._loader = _loader;
-        this._cache = new Map();
-        this._requestsStream = new Subject();
-        this._subscription = this._requestsStream
-            .pipe(bufferTime(20), filter((requests) => requests.length > 0), mergeMap(async (requests) => this._loader(requests)), mergeAll())
-            .subscribe({
-            next: ({ modelId, categoryId, elementsCount }) => {
-                const subject = this._cache.get(`${modelId}${categoryId}`);
-                assert(!!subject);
-                subject.next(elementsCount);
-            },
-        });
-    }
-    [Symbol.dispose]() {
-        this._subscription.unsubscribe();
-    }
-    async getCategoryElementsCount(modelId, categoryId) {
-        const cacheKey = `${modelId}${categoryId}`;
-        let result = this._cache.get(cacheKey);
-        if (result !== undefined) {
-            return firstValueFrom(result);
-        }
-        result = new ReplaySubject(1);
-        this._cache.set(cacheKey, result);
-        this._requestsStream.next({ modelId, categoryId });
-        return firstValueFrom(result);
-    }
-}
 //# sourceMappingURL=ModelsTreeIdsCache.js.map