Reload tree when imodel changes (#101)

* Reload tree when imodel changes

* Fix comments

* Extract duplicate code into function

* Change files

change/@itwin-presentation-components-2ee7274e-62d6-4c1a-90af-6a1a9bf03484.json

@@ -0,0 +1,7 @@
+{
+  "type": "minor",
+  "comment": "Unified hierarchy update handling in all situations (after iModel change, ruleset change or ruleset variables change)",
+  "packageName": "@itwin/presentation-components",
+  "email": "24278440+saskliutas@users.noreply.github.com",
+  "dependentChangeType": "patch"
+}

packages/components/src/presentation-components/tree/controlled/TreeReloader.ts

@@ -5,6 +5,7 @@
 
 import { Observable } from "rxjs/internal/Observable";
 import { concat } from "rxjs/internal/observable/concat";
+import { defer } from "rxjs/internal/observable/defer";
 import { EMPTY } from "rxjs/internal/observable/empty";
 import { from } from "rxjs/internal/observable/from";
 import { concatMap } from "rxjs/internal/operators/concatMap";
@@ -14,10 +15,12 @@ import { filter } from "rxjs/internal/operators/filter";
 import { finalize } from "rxjs/internal/operators/finalize";
 import { ignoreElements } from "rxjs/internal/operators/ignoreElements";
 import { map } from "rxjs/internal/operators/map";
+import { mergeMap } from "rxjs/internal/operators/mergeMap";
 import { take } from "rxjs/internal/operators/take";
 import { tap } from "rxjs/internal/operators/tap";
 import {
-  isTreeModelNode, PagedTreeNodeLoader, TreeModel, TreeModelNode, TreeModelRootNode, TreeModelSource, TreeNodeLoadResult,
+  computeVisibleNodes, isTreeModelNode, isTreeModelNodePlaceholder, PagedTreeNodeLoader, RenderedItemsRange, TreeModel, TreeModelNode,
+  TreeModelNodePlaceholder, TreeModelRootNode, TreeModelSource, TreeNodeLoadResult, VisibleTreeNodes,
 } from "@itwin/components-react";
 import { assert, isIDisposable } from "@itwin/core-bentley";
 import { IPresentationTreeDataProvider } from "../IPresentationTreeDataProvider";
@@ -28,16 +31,18 @@ import { toRxjsObservable } from "../Utils";
  * @param treeModel Previous tree model.
  * @param dataProvider Tree node provider.
  * @param pageSize Data provider's page size.
+ * @param itemsRange Range describing rendered items that are visible.
  * @returns An observable which will emit a new [TreeModelSource]($components-react) and complete.
  * @internal
  */
 export function reloadTree(
   treeModel: TreeModel,
   dataProvider: IPresentationTreeDataProvider,
   pageSize: number,
+  itemsRange?: RenderedItemsRange,
 ): Observable<TreeModelSource> {
   const modelSource = new TreeModelSource();
-  const nodeLoader = new TreeReloader(dataProvider, modelSource, pageSize, treeModel);
+  const nodeLoader = new TreeReloader(dataProvider, modelSource, pageSize, treeModel, itemsRange);
   return nodeLoader.reloadTree().pipe(
     endWith(modelSource),
     finalize(() => isIDisposable(nodeLoader) && /* istanbul ignore next */ nodeLoader.dispose()),
@@ -50,6 +55,7 @@ class TreeReloader extends PagedTreeNodeLoader<IPresentationTreeDataProvider> {
     modelSource: TreeModelSource,
     pageSize: number,
     private previousTreeModel: TreeModel,
+    private itemsRange?: RenderedItemsRange,
   ) {
     super(dataProvider, modelSource, pageSize);
   }
@@ -59,54 +65,90 @@ class TreeReloader extends PagedTreeNodeLoader<IPresentationTreeDataProvider> {
     return concat(
       // We need to know root node count before continuing
       this.loadNode(this.modelSource.getModel().getRootNode(), 0),
-      from(previouslyExpandedNodes)
+      this.reloadPreviouslyExpandedNodes(previouslyExpandedNodes),
+      this.reloadVisibleNodes(),
+    ).pipe(
+      ignoreElements()
+    );
+  }
+
+  private reloadPreviouslyExpandedNodes(previouslyExpandedNodes: ExpandedNode[]) {
+    return from(previouslyExpandedNodes)
+      .pipe(
+        // Process expanded nodes recursively, breadth first
+        expand((expandedNode) => {
+          const node = this.modelSource.getModel().getNode(expandedNode.id);
+          if (node !== undefined) {
+            // The expanded node is already loaded in the new tree model, now load and expand its children recursively
+            return concat(this.loadChildren(node), expandedNode.expandedChildren);
+          }
+
+          // The expanded node is either not loaded yet, or does not exist in the new tree hierarchy
+          const parentNode = getTreeNode(this.modelSource.getModel(), expandedNode.parentId);
+          if (parentNode === undefined || parentNode.numChildren === undefined) {
+            // Cannot determine sibling count. Assume parent is missing from the new tree or something went wrong.
+            return EMPTY;
+          }
+
+          if (parentNode.numChildren === 0) {
+            // Parent node no longer has any children, thus we will not find the expanded node
+            return EMPTY;
+          }
+
+          // Try to make the expanded node appear in the new tree hierarchy. Test three locations: at, a page before,
+          // and a page after previous known location.
+
+          // TODO: We should keep a list of nodes that we failed to find. There is a chance that we will load them
+          // accidentally while searching for other expanded nodes under the same parent.
+          return from([
+            Math.min(expandedNode.index, parentNode.numChildren - 1),
+            Math.min(Math.max(0, expandedNode.index - this.pageSize), parentNode.numChildren - 1),
+            Math.min(expandedNode.index + this.pageSize, parentNode.numChildren - 1),
+          ])
+            .pipe(
+              // For each guess, load the corresponding page
+              concatMap((index) => this.loadNode(parentNode, index)),
+              // Stop making guesses when the node is found
+              map(() => this.modelSource.getModel().getNode(expandedNode.id)),
+              filter((loadedNode) => loadedNode !== undefined),
+              take(1),
+              // If the node is found, load and expand its children recursively
+              concatMap((loadedNode) => {
+                assert(loadedNode !== undefined);
+                return concat(this.loadChildren(loadedNode), expandedNode.expandedChildren);
+              }),
+            );
+        }),
+      );
+  }
+
+  private reloadVisibleNodes() {
+    return defer(() => {
+      // if visible range is not provided do not load any more nodes
+      if (!this.itemsRange)
+        return EMPTY;
+
+      // collect not loaded (placeholder) nodes that are in visible range
+      const visibleNodes = computeVisibleNodes(this.modelSource.getModel());
+      const visibleRange = getVisibleRange(this.itemsRange, visibleNodes);
+      const notLoadedNode: TreeModelNodePlaceholder[] = [];
+      for (let i = visibleRange.start; i <= visibleRange.end; i++) {
+        const node = visibleNodes.getAtIndex(i);
+        if (!node || !isTreeModelNodePlaceholder(node))
+          continue;
+        notLoadedNode.push(node);
+      }
+
+      // load all placeholder nodes in visible range
+      return from(notLoadedNode)
         .pipe(
-          // Process expanded nodes recursively, breadth first
-          expand((expandedNode) => {
-            const node = this.modelSource.getModel().getNode(expandedNode.id);
-            if (node !== undefined) {
-              // The expanded node is already loaded in the new tree model, now load and expand its children recursively
-              return concat(this.loadChildren(node), expandedNode.expandedChildren);
-            }
-
-            // The expanded node is either not loaded yet, or does not exist in the new tree hierarchy
-            const parentNode = getTreeNode(this.modelSource.getModel(), expandedNode.parentId);
-            if (parentNode === undefined || parentNode.numChildren === undefined) {
-              // Cannot determine sibling count. Assume parent is missing from the new tree or something went wrong.
-              return EMPTY;
-            }
-
-            if (parentNode.numChildren === 0) {
-              // Parent node no longer has any children, thus we will not find the expanded node
-              return EMPTY;
-            }
-
-            // Try to make the expanded node appear in the new tree hierarchy. Test three locations: at, a page before,
-            // and a page after previous known location.
-
-            // TODO: We should keep a list of nodes that we failed to find. There is a chance that we will load them
-            // accidentally while searching for other expanded nodes under the same parent.
-            return from([
-              Math.min(expandedNode.index, parentNode.numChildren - 1),
-              Math.min(Math.max(0, expandedNode.index - this.pageSize), parentNode.numChildren - 1),
-              Math.min(expandedNode.index + this.pageSize, parentNode.numChildren - 1),
-            ])
-              .pipe(
-                // For each guess, load the corresponding page
-                concatMap((index) => this.loadNode(parentNode, index)),
-                // Stop making guesses when the node is found
-                map(() => this.modelSource.getModel().getNode(expandedNode.id)),
-                filter((loadedNode) => loadedNode !== undefined),
-                take(1),
-                // If the node is found, load and expand its children recursively
-                concatMap((loadedNode) => {
-                  assert(loadedNode !== undefined);
-                  return concat(this.loadChildren(loadedNode), expandedNode.expandedChildren);
-                }),
-              );
+          mergeMap((placeholder) => {
+            const parentNode = placeholder.parentId ? this.modelSource.getModel().getNode(placeholder.parentId) : this.modelSource.getModel().getRootNode();
+            assert(parentNode !== undefined);
+            return toRxjsObservable(super.loadNode(parentNode, placeholder.childIndex));
           }),
-        ),
-    ).pipe(ignoreElements());
+        );
+    });
   }
 
   private loadChildren(parentNode: TreeModelNode): Observable<never> {
@@ -172,3 +214,16 @@ function collectExpandedNodes(rootNodeId: string | undefined, treeModel: TreeMod
 function getTreeNode(treeModel: TreeModel, nodeId: string | undefined): TreeModelNode | TreeModelRootNode | undefined {
   return nodeId === undefined ? treeModel.getRootNode() : treeModel.getNode(nodeId);
 }
+
+function getVisibleRange(itemsRange: RenderedItemsRange, visibleNodes: VisibleTreeNodes) {
+  if (itemsRange.visibleStopIndex < visibleNodes.getNumNodes())
+    return { start: itemsRange.visibleStartIndex, end: itemsRange.visibleStopIndex };
+
+  const visibleNodesCount = itemsRange.visibleStopIndex - itemsRange.visibleStartIndex;
+  const endPosition = visibleNodes.getNumNodes() - 1;
+  const startPosition = endPosition - visibleNodesCount;
+  return {
+    start: startPosition < 0 ? 0 : startPosition,
+    end: endPosition < 0 ? 0 : endPosition,
+  };
+}

packages/components/src/test/tree/Utils.test.ts

@@ -45,6 +45,12 @@ describe("Utils", () => {
       expect(treeNode).to.matchSnapshot();
     });
 
+    it("creates tree node with children", () => {
+      const node = createTestECInstancesNode({ hasChildren: true });
+      const treeNode = createTreeNodeItem(node);
+      expect(treeNode.hasChildren).to.be.true;
+    });
+
     it("appends grouped nodes count if requested", () => {
       const node: Node = {
         key: createTestECClassGroupingNodeKey({ groupedInstancesCount: 999 }),