@itwin/tree-widget-react 3.16.1 → 3.17.1

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

@@ -2,11 +2,12 @@
  * Copyright (c) Bentley Systems, Incorporated. All rights reserved.
  * See LICENSE.md in the project root for license terms and full copyright notice.
  *--------------------------------------------------------------------------------------------*/
-import { bufferCount, bufferTime, filter, firstValueFrom, from, map, mergeAll, mergeMap, reduce, ReplaySubject, Subject } from "rxjs";
+import { bufferCount, bufferTime, defaultIfEmpty, defer, EMPTY, filter, firstValueFrom, forkJoin, from, map, mergeAll, mergeMap, reduce, ReplaySubject, shareReplay, Subject, take, tap, toArray, } from "rxjs";
 import { assert, Guid, Id64 } from "@itwin/core-bentley";
 import { IModel } from "@itwin/core-common";
-import { collect } from "../../common/Rxjs.js";
+import { releaseMainThreadOnItemsCount } from "../../common/internal/Utils.js";
 import { pushToMap } from "../../common/Utils.js";
+import { getOptimalBatchSize } from "../Utils.js";
 /** @internal */
 export class ModelsTreeIdsCache {
     #categoryElementCounts;
@@ -19,93 +20,221 @@ export class ModelsTreeIdsCache {
     #categoryKeyPaths;
     #queryExecutor;
     #hierarchyConfig;
+    #childrenMap;
+    /** Stores element ids which have children query scheduled to execute. */
+    #childrenLoadingMap;
     #componentId;
     #componentName;
     constructor(queryExecutor, hierarchyConfig, componentId) {
         this.#hierarchyConfig = hierarchyConfig;
         this.#queryExecutor = queryExecutor;
-        this.#categoryElementCounts = new ModelCategoryElementsCountCache(async (input) => this.queryCategoryElementCounts(input));
+        this.#categoryElementCounts = new ModelCategoryElementsCountCache((input) => this.queryCategoryElementCounts(input));
         this.#modelKeyPaths = new Map();
         this.#subjectKeyPaths = new Map();
         this.#categoryKeyPaths = new Map();
+        this.#childrenMap = new Map();
+        this.#childrenLoadingMap = new Map();
         this.#componentId = componentId ?? Guid.createValue();
         this.#componentName = "ModelsTreeIdsCache";
     }
     [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", restartToken: `${this.#componentName}/${this.#componentId}/subjects` })) {
-            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 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
+      `;
+            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 = `
+    queryModels() {
+        return defer(() => {
+            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", restartToken: `${this.#componentName}/${this.#componentId}/models` })) {
-            yield { id: row.id, parentId: row.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)`}
+      `;
+            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 };
+        }));
+    }
+    queryChildren({ elementIds }) {
+        if (elementIds.length === 0) {
+            return EMPTY;
         }
+        return defer(() => {
+            const ctes = [
+                `
+          ElementChildren(id, parentId) AS (
+            SELECT ECInstanceId id, Parent.Id parentId
+            FROM ${this.#hierarchyConfig.elementClassSpecification}
+            WHERE Parent.Id IN (${elementIds.join(", ")})
+
+            UNION ALL
+
+            SELECT c.ECInstanceId id, c.Parent.Id
+            FROM ${this.#hierarchyConfig.elementClassSpecification} c
+            JOIN ElementChildren p ON c.Parent.Id = p.id
+          )
+        `,
+            ];
+            const ecsql = `
+        SELECT id, parentId
+        FROM ElementChildren
+      `;
+            return this.#queryExecutor.createQueryReader({ ecsql, ctes }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded", restartToken: `${this.#componentName}/${this.#componentId}/children/${Guid.createValue()}` });
+        }).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);
+    getChildrenTreeFromMap({ elementIds }) {
+        const result = new Map();
+        if (Id64.sizeOf(elementIds) === 0 || this.#childrenMap.size === 0) {
+            return result;
+        }
+        for (const elementId of Id64.iterable(elementIds)) {
+            const entry = this.#childrenMap.get(elementId);
+            if (!entry?.children) {
+                continue;
+            }
+            const elementChildrenTree = new Map();
+            result.set(elementId, { children: elementChildrenTree });
+            entry.children.forEach((childId) => {
+                const childrenTreeOfChild = this.getChildrenTreeFromMap({ elementIds: childId });
+                // Need to add children tree created from childId. This tree includes childId as root element
+                // If child does not have children, children tree won't be created. Need to add childId with undefined children
+                elementChildrenTree.set(childId, { children: childrenTreeOfChild.size > 0 ? childrenTreeOfChild : undefined });
+            });
+        }
+        return result;
+    }
+    getChildrenCountMap({ elementIds }) {
+        const result = new Map();
+        for (const elementId of Id64.iterable(elementIds)) {
+            const entry = this.#childrenMap.get(elementId);
+            if (entry?.children) {
+                let totalChildrenCount = entry.children.length;
+                this.getChildrenCountMap({ elementIds: entry.children }).forEach((childrenOfChildCount) => (totalChildrenCount += childrenOfChildCount));
+                result.set(elementId, totalChildrenCount);
+            }
+        }
+        return result;
+    }
+    /**
+     * Populates #childrenLoadingMap with promises. When these promises resolve, they will populate #childrenMap with values and delete entries from #childrenLoadingMap.
+     */
+    createChildrenLoadingMapEntries({ elementsToQuery }) {
+        const getElementsToQueryPromise = async (batchedElementsToQuery) => firstValueFrom(this.queryChildren({ elementIds: batchedElementsToQuery }).pipe(
+        // Want to have void at the end instead of void[], so using reduce
+        reduce((acc, { parentId, id }) => {
+            let entry = this.#childrenMap.get(parentId);
+            if (!entry) {
+                entry = { children: [] };
+                this.#childrenMap.set(parentId, entry);
+            }
+            if (!entry.children) {
+                entry.children = [];
+            }
+            // Add child to parent's entry and add child to children map
+            entry.children.push(id);
+            if (!this.#childrenMap.has(id)) {
+                this.#childrenMap.set(id, { children: undefined });
+            }
+            return acc;
+        }, (() => { })()), tap({ complete: () => batchedElementsToQuery.forEach((elementId) => this.#childrenLoadingMap.delete(elementId)) }), defaultIfEmpty((() => { })())));
+        const maximumBatchSize = 1000;
+        const totalSize = elementsToQuery.length;
+        const optimalBatchSize = getOptimalBatchSize({ totalSize, maximumBatchSize });
+        const loadingMapEntries = new Array();
+        // Don't want to slice if its not necessary
+        if (totalSize <= maximumBatchSize) {
+            loadingMapEntries.push(getElementsToQueryPromise(elementsToQuery));
+        }
+        else {
+            for (let i = 0; i < elementsToQuery.length; i += optimalBatchSize) {
+                loadingMapEntries.push(getElementsToQueryPromise(elementsToQuery.slice(i, i + optimalBatchSize)));
+            }
+        }
+        elementsToQuery.forEach((elementId, index) => this.#childrenLoadingMap.set(elementId, loadingMapEntries[Math.floor(index / optimalBatchSize)]));
+        return { loadingMapEntries: Promise.all(loadingMapEntries).then(() => { }) };
+    }
+    createChildrenMapEntries({ elementIds }) {
+        return from(Id64.iterable(elementIds)).pipe(reduce((acc, elementId) => {
+            if (this.#childrenMap.has(elementId)) {
+                return acc;
+            }
+            const loadingPromise = this.#childrenLoadingMap.get(elementId);
+            if (loadingPromise) {
+                acc.existingPromises.push(loadingPromise);
+            }
+            else {
+                acc.elementsToQuery.push(elementId);
+            }
+            return acc;
+        }, { existingPromises: new Array(), elementsToQuery: new Array() }), mergeMap(async ({ elementsToQuery, existingPromises }) => {
+            existingPromises.push(this.createChildrenLoadingMapEntries({ elementsToQuery }).loadingMapEntries);
+            return Promise.all(existingPromises);
+        }));
+    }
+    getChildrenTree({ elementIds }) {
+        return this.createChildrenMapEntries({ elementIds }).pipe(map(() => this.getChildrenTreeFromMap({ elementIds })));
+    }
+    getAllChildrenCount({ elementIds }) {
+        return this.createChildrenMapEntries({ elementIds }).pipe(map(() => this.getChildrenCountMap({ elementIds })));
+    }
+    getSubjectInfos() {
+        this.#subjectInfos ??= forkJoin({
+            subjectInfos: this.querySubjects().pipe(reduce((acc, subject) => {
+                const subjectInfo = {
+                    parentSubject: subject.parentId,
+                    hideInHierarchy: subject.hideInHierarchy,
+                    childSubjects: new Set(),
+                    childModels: new Set(),
+                };
+                if (subject.targetPartitionId) {
+                    subjectInfo.childModels.add(subject.targetPartitionId);
+                }
+                acc.set(subject.id, subjectInfo);
+                return acc;
+            }, new Map()), map((subjectInfos) => {
+                for (const [subjectId, { parentSubject: parentSubjectId }] of subjectInfos.entries()) {
+                    if (parentSubjectId) {
+                        const parentSubjectInfo = subjectInfos.get(parentSubjectId);
+                        assert(!!parentSubjectInfo);
+                        parentSubjectInfo.childSubjects.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);
@@ -114,13 +243,12 @@ export class ModelsTreeIdsCache {
                 });
             }
             return 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) {
@@ -133,71 +261,73 @@ export class ModelsTreeIdsCache {
                 }
             });
             return [...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";
+    getChildSubjectIds(parentSubjectIds) {
+        return this.getSubjectInfos().pipe(map((subjectInfos) => {
+            const childSubjectIds = new Array();
+            parentSubjectIds.forEach((subjectId) => {
+                forEachChildSubject(subjectInfos, subjectId, (childSubjectId, childSubjectInfo) => {
+                    if (!childSubjectInfo.hideInHierarchy) {
+                        childSubjectIds.push(childSubjectId);
+                        return "break";
+                    }
+                    return "continue";
+                });
             });
-        });
-        return childSubjectIds;
+            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;
-            }
-            const subjectInfo = subjectInfos.get(subjectId);
-            if (!subjectInfo) {
-                continue;
+    getSubjectModelIds(subjectIds) {
+        return this.getSubjectInfos().pipe(map((subjectInfos) => {
+            const subjectStack = [...subjectIds];
+            const result = new Array();
+            while (true) {
+                const subjectId = subjectStack.pop();
+                if (subjectId === undefined) {
+                    break;
+                }
+                const subjectInfo = subjectInfos.get(subjectId);
+                if (!subjectInfo) {
+                    continue;
+                }
+                result.push(...subjectInfo.childModels);
+                subjectStack.push(...subjectInfo.childSubjects);
             }
-            result.push(...subjectInfo.childModels);
-            subjectStack.push(...subjectInfo.childSubjects);
-        }
-        return result;
+            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();
+            parentSubjectIds.forEach((subjectId) => {
+                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.childModels);
-        });
-        return modelIds;
+            const modelIds = new Array();
+            [...parentSubjectIds, ...hiddenSubjectIds].forEach((subjectId) => {
+                const subjectInfo = subjectInfos.get(subjectId);
+                subjectInfo && modelIds.push(...subjectInfo.childModels);
+            });
+            return modelIds;
+        }));
     }
-    async createSubjectInstanceKeysPath(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) {
@@ -211,128 +341,111 @@ export class ModelsTreeIdsCache {
                     currParentId = parentInfo?.parentSubject;
                 }
                 return result.reverse();
-            })();
+            }), shareReplay());
             this.#subjectKeyPaths.set(targetSubjectId, entry);
         }
         return entry;
     }
-    async *queryModelCategories() {
-        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 BisCore.Model m
-      JOIN ${this.#hierarchyConfig.elementClassSpecification} this ON m.ECInstanceId = this.Model.Id
-      GROUP BY modelId, categoryId, isModelPrivate
-    `;
-        for await (const row of this.#queryExecutor.createQueryReader({ ecsql: query }, { rowFormat: "ECSqlPropertyNames", limit: "unbounded", restartToken: `${this.#componentName}/${this.#componentId}/model-categories` })) {
-            yield { modelId: row.modelId, categoryId: row.categoryId, isModelPrivate: !!row.isModelPrivate, isRootElementCategory: !!row.isRootElementCategory };
-        }
+    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 BisCore.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 *queryModeledElements() {
-        const query = `
-      SELECT
-        pe.ECInstanceId modeledElementId,
-        pe.Category.Id categoryId,
-        pe.Model.Id modelId
-      FROM BisCore.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", restartToken: `${this.#componentName}/${this.#componentId}/modeled-elements` })) {
-            yield { modelId: row.modelId, categoryId: row.categoryId, modeledElementId: row.modeledElementId };
-        }
+    queryModeledElements() {
+        return defer(() => {
+            const query = `
+        SELECT
+          pe.ECInstanceId modeledElementId,
+          pe.Category.Id categoryId,
+          pe.Model.Id modelId
+        FROM BisCore.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})
+      `;
+            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 getModelWithCategoryModeledElements() {
-        this.#modelWithCategoryModeledElements ??= (async () => {
-            const modelWithCategoryModeledElements = new Map();
-            for await (const { modelId, categoryId, modeledElementId } of this.queryModeledElements()) {
-                const key = `${modelId}-${categoryId}`;
-                const entry = modelWithCategoryModeledElements.get(key);
-                if (entry === undefined) {
-                    modelWithCategoryModeledElements.set(key, new Set([modeledElementId]));
-                }
-                else {
-                    entry.add(modeledElementId);
-                }
+    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]));
             }
-            return modelWithCategoryModeledElements;
-        })();
+            else {
+                entry.add(modeledElementId);
+            }
+            return acc;
+        }, new Map()), shareReplay());
         return this.#modelWithCategoryModeledElements;
     }
-    async getModelInfos() {
-        this.#modelInfos ??= (async () => {
-            const modelInfos = new Map();
-            for await (const { modelId, categoryId, isModelPrivate, isRootElementCategory } of this.queryModelCategories()) {
-                const entry = modelInfos.get(modelId);
-                if (entry) {
-                    entry.categories.set(categoryId, { isRootElementCategory });
-                    entry.isModelPrivate = isModelPrivate;
-                }
-                else {
-                    modelInfos.set(modelId, { categories: new Map([[categoryId, { isRootElementCategory }]]), isModelPrivate });
-                }
+    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 modelInfos;
-        })();
+            return acc;
+        }, new Map()), shareReplay());
         return this.#modelInfos;
     }
-    async getAllCategories() {
-        const modelInfos = await this.getModelInfos();
-        const result = new Set();
-        modelInfos.forEach(({ categories }) => {
-            categories.forEach((_, categoryId) => result.add(categoryId));
-        });
-        return result;
+    getAllCategories() {
+        return this.getModelInfos().pipe(mergeMap((modelInfos) => modelInfos.values()), mergeMap(({ categories }) => categories.keys()), reduce((acc, categoryId) => {
+            acc.add(categoryId);
+            return acc;
+        }, new Set()));
     }
-    async getModelCategories(modelId) {
-        const modelInfos = await this.getModelInfos();
-        const categories = modelInfos.get(modelId)?.categories.keys();
-        return categories ? [...categories] : [];
+    getModelCategories(modelId) {
+        return this.getModelInfos().pipe(mergeMap((modelInfos) => modelInfos.get(modelId)?.categories.keys() ?? []), toArray());
     }
-    async hasSubModel(elementId) {
-        const modelInfos = await this.getModelInfos();
-        const modeledElementInfo = modelInfos.get(elementId);
-        if (!modeledElementInfo) {
-            return false;
-        }
-        return !modeledElementInfo.isModelPrivate;
+    hasSubModel(elementId) {
+        return this.getModelInfos().pipe(map((modelInfos) => {
+            const modeledElementInfo = modelInfos.get(elementId);
+            if (!modeledElementInfo) {
+                return false;
+            }
+            return !modeledElementInfo.isModelPrivate;
+        }));
     }
-    async getCategoriesModeledElements(modelId, categoryIds) {
-        const modelWithCategoryModeledElements = await this.getModelWithCategoryModeledElements();
-        const result = new Array();
-        for (const categoryId of Id64.iterable(categoryIds)) {
+    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) {
-                result.push(...entry);
+                acc.push(...entry);
             }
-        }
-        return result;
+            return acc;
+        }, new Array()))));
     }
-    async createModelInstanceKeyPaths(modelId) {
+    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 }]);
-                    }
-                }
-                return result;
-            })();
+            entry = this.getSubjectInfos().pipe(mergeMap((subjectInfos) => subjectInfos.entries()), filter(([_, subjectInfo]) => subjectInfo.childModels.has(modelId)), mergeMap(([modelSubjectId]) => this.createSubjectInstanceKeysPath(modelSubjectId).pipe(map((path) => [...path, { className: "BisCore.GeometricModel3d", id: modelId }]))), toArray(), shareReplay());
             this.#modelKeyPaths.set(modelId, entry);
         }
         return entry;
     }
-    async queryCategoryElementCounts(input) {
-        return collect(from(input).pipe(reduce((acc, { modelId, categoryId }) => {
+    queryCategoryElementCounts(input) {
+        return from(input).pipe(reduce((acc, { modelId, categoryId }) => {
             const entry = acc.get(modelId);
             if (!entry) {
                 acc.set(modelId, new Set([categoryId]));
@@ -344,73 +457,64 @@ export class ModelsTreeIdsCache {
         }, new Map()), mergeMap((modelCategoryMap) => modelCategoryMap.entries()), map(([modelId, categoryIds]) => `Model.Id = ${modelId} AND Category.Id IN (${[...categoryIds].join(", ")})`), 
         // we may have thousands of where clauses here, and sending a single query with all of them could take a
         // long time - instead, split it into smaller chunks
-        bufferCount(100), mergeMap(async (whereClauses) => {
-            const reader = this.#queryExecutor.createQueryReader({
-                ctes: [
-                    `
-                  CategoryElements(id, modelId, categoryId) AS (
-                    SELECT ECInstanceId, Model.Id, Category.Id
-                    FROM ${this.#hierarchyConfig.elementClassSpecification}
-                    WHERE
-                      Parent.Id IS NULL
-                      AND (
-                        ${whereClauses.join(" OR ")}
-                      )
+        bufferCount(100), mergeMap((whereClauses) => defer(() => this.#queryExecutor.createQueryReader({
+            ctes: [
+                `
+                    CategoryElements(id, modelId, categoryId) AS (
+                      SELECT ECInstanceId, Model.Id, Category.Id
+                      FROM ${this.#hierarchyConfig.elementClassSpecification}
+                      WHERE
+                        Parent.Id IS NULL
+                        AND (
+                          ${whereClauses.join(" OR ")}
+                        )
 
-                    UNION ALL
+                      UNION ALL
 
-                    SELECT c.ECInstanceId, p.modelId, p.categoryId
-                    FROM ${this.#hierarchyConfig.elementClassSpecification} c
-                    JOIN CategoryElements p ON c.Parent.Id = p.id
-                  )
+                      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
                 `,
-                ],
-                ecsql: `
-                SELECT modelId, categoryId, COUNT(id) elementsCount
-                FROM CategoryElements
-                GROUP BY modelId, categoryId
-              `,
-            }, {
-                rowFormat: "ECSqlPropertyNames",
-                limit: "unbounded",
-                restartToken: `${this.#componentName}/${this.#componentId}/category-element-counts/${Guid.createValue()}`,
+        }, {
+            rowFormat: "ECSqlPropertyNames",
+            limit: "unbounded",
+            restartToken: `${this.#componentName}/${this.#componentId}/category-element-counts/${Guid.createValue()}`,
+        }))), releaseMainThreadOnItemsCount(500), reduce(({ acc, createKey }, row) => {
+            acc.set(createKey({ modelId: row.modelId, categoryId: row.categoryId }), {
+                modelId: row.modelId,
+                categoryId: row.categoryId,
+                elementsCount: row.elementsCount,
             });
-            const result = new Array();
-            for await (const row of reader) {
-                result.push({ modelId: row.modelId, categoryId: row.categoryId, elementsCount: row.elementsCount });
-            }
+            return { acc, createKey };
+        }, {
+            acc: new Map(),
+            createKey: (keyProps) => `${keyProps.modelId}-${keyProps.categoryId}`,
+        }), mergeMap(({ acc: result, createKey }) => {
             input.forEach(({ modelId, categoryId }) => {
-                if (!result.some((queriedResult) => queriedResult.categoryId === categoryId && queriedResult.modelId === modelId)) {
-                    result.push({ categoryId, modelId, elementsCount: 0 });
+                if (!result.has(createKey({ modelId, categoryId }))) {
+                    result.set(createKey({ modelId, categoryId }), { categoryId, modelId, elementsCount: 0 });
                 }
             });
-            return result;
-        }), mergeAll()));
+            return from(result.values());
+        }), toArray());
     }
-    async getCategoryElementsCount(modelId, categoryId) {
+    getCategoryElementsCount(modelId, categoryId) {
         return this.#categoryElementCounts.getCategoryElementsCount(modelId, categoryId);
     }
-    async createCategoryInstanceKeyPaths(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) => {
-                    const categoryEntry = modelInfo.categories.get(categoryId);
-                    if (categoryEntry?.isRootElementCategory) {
-                        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 }]);
-                    }
-                }
-                return categoryPaths;
-            })();
+            entry = this.getModelInfos().pipe(mergeMap((modelInfos) => modelInfos.entries()), filter(([_, modelInfo]) => !!modelInfo.categories.get(categoryId)?.isRootElementCategory), mergeMap(([categoryModelId]) => this.createModelInstanceKeyPaths(categoryModelId)), mergeMap((modelPaths) => modelPaths), reduce((acc, modelPath) => {
+                acc.push([...modelPath, { className: "BisCore.SpatialCategory", id: categoryId }]);
+                return acc;
+            }, new Array()), shareReplay());
             this.#categoryKeyPaths.set(categoryId, entry);
         }
         return entry;
@@ -433,7 +537,7 @@ class ModelCategoryElementsCountCache {
     #subscription;
     constructor(loader) {
         this.#subscription = this.#requestsStream
-            .pipe(bufferTime(20), filter((requests) => requests.length > 0), mergeMap(async (requests) => loader(requests)), mergeAll())
+            .pipe(bufferTime(20), filter((requests) => requests.length > 0), mergeMap((requests) => loader(requests)), mergeAll())
             .subscribe({
             next: ({ modelId, categoryId, elementsCount }) => {
                 const subject = this.#cache.get(`${modelId}${categoryId}`);
@@ -445,16 +549,16 @@ class ModelCategoryElementsCountCache {
     [Symbol.dispose]() {
         this.#subscription.unsubscribe();
     }
-    async getCategoryElementsCount(modelId, categoryId) {
+    getCategoryElementsCount(modelId, categoryId) {
         const cacheKey = `${modelId}${categoryId}`;
         let result = this.#cache.get(cacheKey);
         if (result !== undefined) {
-            return firstValueFrom(result);
+            return from(result).pipe(take(1));
         }
         result = new ReplaySubject(1);
         this.#cache.set(cacheKey, result);
         this.#requestsStream.next({ modelId, categoryId });
-        return firstValueFrom(result);
+        return from(result).pipe(take(1));
     }
 }
 //# sourceMappingURL=ModelsTreeIdsCache.js.map