From 68cb4b2d25d3bc149e4f8e80c030a16db75959c2 Mon Sep 17 00:00:00 2001
From: Oleksandr Saulyak <oleksandr.saulyak@codefresh.io>
Date: Tue, 12 Sep 2023 05:25:47 +0300
Subject: [PATCH] fix(module:tree-view): re-rendering fix (#8035)

[ node.ts ]: 'hasContextShapeChanged' func previously didn't allow customization of how to compare data objects, now it will use treeControl.trackBy func to do it in proper and optimized way
---
 components/tree-view/node.ts                     | 14 +++++++++++++-
 components/tree-view/tree-virtual-scroll-view.ts | 15 +++++++++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/components/tree-view/node.ts b/components/tree-view/node.ts
index 1842594a06b..d860595ca19 100644
--- a/components/tree-view/node.ts
+++ b/components/tree-view/node.ts
@@ -131,6 +131,7 @@ export class NzTreeNodeDefDirective<T> extends CdkTreeNodeDef<T> {
 export class NzTreeVirtualScrollNodeOutletDirective<T> implements OnChanges {
   private _viewRef: EmbeddedViewRef<NzSafeAny> | null = null;
   @Input() data!: NzTreeVirtualNodeData<T>;
+  @Input() compareBy?: ((value: T) => T | string | number) | null;
 
   constructor(private _viewContainerRef: ViewContainerRef) {}
 
@@ -170,11 +171,22 @@ export class NzTreeVirtualScrollNodeOutletDirective<T> implements OnChanges {
           return true;
         }
       }
-      return ctxChange.previousValue?.data !== ctxChange.currentValue?.data;
+      return (
+        this.innerCompareBy(ctxChange.previousValue?.data ?? null) !==
+        this.innerCompareBy(ctxChange.currentValue?.data ?? null)
+      );
     }
     return true;
   }
 
+  get innerCompareBy(): (value: T | null) => T | string | number | null {
+    return value => {
+      if (value === null) return value;
+      if (this.compareBy) return this.compareBy(value as T);
+      return value;
+    };
+  }
+
   private updateExistingContext(ctx: NzSafeAny): void {
     for (const propName of Object.keys(ctx)) {
       this._viewRef!.context[propName] = (this.data.context as NzSafeAny)[propName];
diff --git a/components/tree-view/tree-virtual-scroll-view.ts b/components/tree-view/tree-virtual-scroll-view.ts
index 7454d1f2c95..f6f0969b910 100644
--- a/components/tree-view/tree-virtual-scroll-view.ts
+++ b/components/tree-view/tree-virtual-scroll-view.ts
@@ -4,7 +4,7 @@
  */
 
 import { CdkVirtualScrollViewport } from '@angular/cdk/scrolling';
-import { CdkTree, CdkTreeNodeOutletContext } from '@angular/cdk/tree';
+import { CdkTree, CdkTreeNodeOutletContext, BaseTreeControl } from '@angular/cdk/tree';
 import {
   ChangeDetectionStrategy,
   Component,
@@ -16,6 +16,8 @@ import {
   ViewEncapsulation
 } from '@angular/core';
 
+import { NzSafeAny } from 'ng-zorro-antd/core/types';
+
 import { NzTreeVirtualNodeData } from './node';
 import { NzTreeNodeOutletDirective } from './outlet';
 import { NzTreeView } from './tree';
@@ -34,7 +36,7 @@ const DEFAULT_SIZE = 28;
         [maxBufferPx]="nzMaxBufferPx"
       >
         <ng-container *cdkVirtualFor="let item of nodes; let i = index; trackBy: innerTrackBy">
-          <ng-template nzTreeVirtualScrollNodeOutlet [data]="item"></ng-template>
+          <ng-template nzTreeVirtualScrollNodeOutlet [data]="item" [compareBy]="compareBy"></ng-template>
         </ng-container>
       </cdk-virtual-scroll-viewport>
     </div>
@@ -74,6 +76,15 @@ export class NzTreeVirtualScrollViewComponent<T> extends NzTreeView<T> implement
     }
   }
 
+  get compareBy(): ((value: T) => NzSafeAny) | null {
+    const baseTreeControl = this.treeControl as BaseTreeControl<T, NzSafeAny>;
+    if (baseTreeControl.trackBy) {
+      return baseTreeControl.trackBy;
+    }
+
+    return null;
+  }
+
   override renderNodeChanges(data: T[] | readonly T[]): void {
     this.nodes = new Array(...data).map((n, i) => this.createNode(n, i));
     this._dataSourceChanged.next();