fix(devtools): retain tree-viz snapped node on pre-render cleanup

Each tree-visualizer render cycle resets the snapped node in an effort to not end up with a non-existent snapped node. However, each render cycle doesn't constitute a completely different tree. This change retains the snapped node as long as it exists in the tree.

devtools/projects/ng-devtools/src/lib/devtools-tabs/injector-tree/injector-tree.component.html

@@ -37,6 +37,7 @@
                   #environmentTree
                   a11yTitle="Environment hierarchy visualization"
                   [root]="environmentInjectorTree()!"
+                  [nodeEqualityFn]="treeNodesEqualityFn"
                   [config]="environmentTreeConfig"
                   (nodeClick)="selectInjectorByNode($event)"
                   (render)="onTreeRender(environmentTree, $event)"
@@ -55,6 +56,7 @@
                   #elementTree
                   a11yTitle="Element hierarchy visualization"
                   [root]="elementInjectorTree()!"
+                  [nodeEqualityFn]="treeNodesEqualityFn"
                   [config]="elementTreeConfig"
                   (nodeClick)="selectInjectorByNode($event)"
                   (render)="onTreeRender(elementTree, $event)"

devtools/projects/ng-devtools/src/lib/devtools-tabs/injector-tree/injector-tree.component.ts

@@ -30,6 +30,7 @@ import {TreeD3Node, TreeVisualizerConfig} from '../../shared/tree-visualizer/tre
 import {TreeVisualizerComponent} from '../../shared/tree-visualizer/tree-visualizer.component';
 import {InjectorProvidersComponent} from './injector-providers/injector-providers.component';
 import {
+  areInjectorTreeNodesEqual,
   areInjectorTreesEqual,
   d3InjectorTreeLinkModifier,
   d3InjectorTreeNodeModifier,
@@ -93,6 +94,8 @@ export class InjectorTreeComponent {
   protected readonly componentExplorerView = input.required<ComponentExplorerView | null>();
   protected readonly hidden = input(false);
 
+  protected readonly treeNodesEqualityFn = areInjectorTreeNodesEqual;
+
   protected readonly diDebugAPIsAvailable = computed<boolean>(() => {
     const view = this.componentExplorerView();
     return !!(view && view.forest.length && view.forest[0].resolutionPath);

devtools/projects/ng-devtools/src/lib/shared/tree-visualizer/tree-visualizer.component.ts

@@ -19,7 +19,13 @@ import {
   signal,
   viewChild,
 } from '@angular/core';
-import {TreeD3Node, TreeNode, TreeVisualizer, TreeVisualizerConfig} from './tree-visualizer';
+import {
+  TreeD3Node,
+  TreeNode,
+  TreeNodeEqualityFn,
+  TreeVisualizer,
+  TreeVisualizerConfig,
+} from './tree-visualizer';
 
 let instanceIdx = 0;
 
@@ -45,6 +51,7 @@ export class TreeVisualizerComponent<T extends TreeNode = TreeNode> {
   protected readonly group = viewChild.required<ElementRef>('group');
 
   readonly root = input.required<T>();
+  protected readonly nodeEqualityFn = input<TreeNodeEqualityFn<T> | null>(null);
   protected readonly config = input<Partial<TreeVisualizerConfig<T>>>();
   protected readonly a11yTitle = input.required<string>();
   protected readonly a11yTitleId = `tree-vis-host-${++instanceIdx}`;
@@ -68,6 +75,7 @@ export class TreeVisualizerComponent<T extends TreeNode = TreeNode> {
           new TreeVisualizer<T>(
             this.container().nativeElement,
             this.group().nativeElement,
+            this.nodeEqualityFn(),
             this.config(),
           ),
         );

devtools/projects/ng-devtools/src/lib/shared/tree-visualizer/tree-visualizer.ts

@@ -38,6 +38,8 @@ export type SvgD3Link<T extends TreeNode> = d3.Selection<
   TreeD3Node<T>
 >;
 
+export type TreeNodeEqualityFn<T extends TreeNode> = (a: T, b: T) => boolean;
+
 export interface TreeVisualizerConfig<T extends TreeNode> {
   /** WARNING: For vertically-oriented trees, use separation greater than `1` */
   orientation: 'horizontal' | 'vertical';
@@ -77,6 +79,8 @@ export class TreeVisualizer<T extends TreeNode = TreeNode> extends GraphRenderer
   constructor(
     private readonly containerElement: HTMLElement,
     private readonly graphElement: HTMLElement,
+    // Used for improved focus/snapped node retention on tree update.
+    private readonly nodeEqualityFn: TreeNodeEqualityFn<T> | null,
     config: Partial<TreeVisualizerConfig<T>> = {},
   ) {
     super();
@@ -98,12 +102,14 @@ export class TreeVisualizer<T extends TreeNode = TreeNode> extends GraphRenderer
     }
   }
 
+  /** Snaps to root node. NOTE: Relies on container size. */
   override snapToRoot(scale = 1): void {
     if (this.root) {
       this.snapToD3Node(this.root, scale);
     }
   }
 
+  /** Snaps to a provided node. NOTE: Relies on container size. */
   override snapToNode(node: T, scale = 1): void {
     const d3Node = this.findD3NodeByDataNode(node);
     if (d3Node) {
@@ -124,7 +130,7 @@ export class TreeVisualizer<T extends TreeNode = TreeNode> extends GraphRenderer
   override cleanup(): void {
     super.cleanup();
     d3.select(this.graphElement).selectAll('*').remove();
-    this.snappedNode = null;
+    this.smartSnappedNodeReset();
   }
 
   override dispose(): void {
@@ -405,4 +411,39 @@ export class TreeVisualizer<T extends TreeNode = TreeNode> extends GraphRenderer
 
     return null;
   }
+
+  /**
+   * Resets `snappedNode` only if the current tree doesn't contain it.
+   * Will perform deep equality check, if `nodeEqualityFn` is provided;
+   * or fallback to a shallow comparison, if not.
+   */
+  private smartSnappedNodeReset() {
+    const snappedNode = this.snappedNode?.node.data;
+    if (!this.root?.data || !snappedNode) {
+      // Null input data; reset.
+      this.snappedNode = null;
+      return;
+    }
+
+    const stack = [this.root];
+
+    while (stack.length) {
+      const node = stack.pop()!;
+
+      if (
+        this.nodeEqualityFn
+          ? this.nodeEqualityFn(node.data, snappedNode)
+          : node.data === snappedNode
+      ) {
+        // Node found; do not reset.
+        return;
+      }
+      for (const child of node?.children || []) {
+        stack.push(child);
+      }
+    }
+
+    // Node not found; reset.
+    this.snappedNode = null;
+  }
 }