Support expanding all layers through initial ui state.

Set deepestExpandedGroupNodeIds to ['___all___'] PiperOrigin-RevId: 669105037
google-ai-edge · Aug 30, 2024 · db7065e · db7065e
1 parent 21874a7
commit db7065e
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git 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,17 +184,25 @@ 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);
       }
       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
@@ -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.