From db7065ea8731b6e070b15f2102d915207c1edadf Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Thu, 29 Aug 2024 16:38:01 -0700
Subject: [PATCH] Support expanding all layers through initial ui state.

Set deepestExpandedGroupNodeIds to ['___all___']

PiperOrigin-RevId: 669105037
---
 .../src/components/visualizer/common/utils.ts | 17 +++++++++++++----
 .../components/visualizer/webgl_renderer.ts   | 19 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/ui/src/components/visualizer/common/utils.ts b/src/ui/src/components/visualizer/common/utils.ts
index c9ef9217..6656edc2 100644
--- a/src/ui/src/components/visualizer/common/utils.ts
+++ b/src/ui/src/components/visualizer/common/utils.ts
@@ -171,6 +171,7 @@ export function getDeepestExpandedGroupNodeIds(
   root: GroupNode | undefined,
   modelGraph: ModelGraph,
   deepestExpandedGroupNodeIds: string[],
+  ignoreExpandedState = false,
 ) {
   let nsChildrenIds: string[] = [];
   if (root == null) {
@@ -183,10 +184,17 @@ export function getDeepestExpandedGroupNodeIds(
     if (!childNode) {
       continue;
     }
-    if (isGroupNode(childNode) && childNode.expanded) {
-      const isDeepest = (childNode.nsChildrenIds || [])
-        .filter((id) => isGroupNode(modelGraph.nodesById[id]))
-        .every((id) => !(modelGraph.nodesById[id] as GroupNode).expanded);
+    if (
+      isGroupNode(childNode) &&
+      (ignoreExpandedState || (!ignoreExpandedState && childNode.expanded))
+    ) {
+      const nsChildrenIds = childNode.nsChildrenIds || [];
+      const isDeepest = ignoreExpandedState
+        ? nsChildrenIds.filter((id) => isGroupNode(modelGraph.nodesById[id]))
+            .length === 0
+        : nsChildrenIds
+            .filter((id) => isGroupNode(modelGraph.nodesById[id]))
+            .every((id) => !(modelGraph.nodesById[id] as GroupNode).expanded);
       if (isDeepest) {
         deepestExpandedGroupNodeIds.push(childNode.id);
       }
@@ -194,6 +202,7 @@ export function getDeepestExpandedGroupNodeIds(
         childNode,
         modelGraph,
         deepestExpandedGroupNodeIds,
+        ignoreExpandedState,
       );
     }
   }
diff --git a/src/ui/src/components/visualizer/webgl_renderer.ts b/src/ui/src/components/visualizer/webgl_renderer.ts
index 583329d3..c9e17199 100644
--- a/src/ui/src/components/visualizer/webgl_renderer.ts
+++ b/src/ui/src/components/visualizer/webgl_renderer.ts
@@ -80,6 +80,7 @@ import {
 } from './common/types';
 import {
   genUid,
+  getDeepestExpandedGroupNodeIds,
   getHighQualityPixelRatio,
   getNodeStyleValue,
   hasNonEmptyQueries,
@@ -766,9 +767,25 @@ export class WebglRenderer implements OnInit, OnDestroy {
       if (!paneState) {
         initGraphFn();
       } else {
+        // Expand all layers if paneState.deepestExpandedGroupNodeIds has only
+        // one elemenet '___all___'.
+        let deepestExpandedGroupNodeIds = paneState.deepestExpandedGroupNodeIds;
+        if (
+          deepestExpandedGroupNodeIds.length === 1 &&
+          deepestExpandedGroupNodeIds[0] === '___all___'
+        ) {
+          const groupNodeIds: string[] = [];
+          getDeepestExpandedGroupNodeIds(
+            undefined,
+            this.curModelGraph,
+            groupNodeIds,
+            true,
+          );
+          deepestExpandedGroupNodeIds = groupNodeIds;
+        }
         this.sendRelayoutGraphRequest(
           paneState.selectedNodeId,
-          paneState.deepestExpandedGroupNodeIds,
+          deepestExpandedGroupNodeIds,
           true,
         );
         // This is needed for loading old perma-link.