Improvements on subgraphs:

- Fix a bug where duplicated entries would show up in graph selector when multiple nodes in a graph link to the same subgraph.
- Add support for multi-parent subgraph (i.e. a subgraph can be linked to multiple parent nodes)
- Add support for opening the subgraph in the split pane by alt-clicking the subgraph indicator.

PiperOrigin-RevId: 731096967

src/ui/src/components/visualizer/app_service.ts

@@ -24,7 +24,7 @@ import {
   LOCAL_STORAGE_KEY_SHOW_ON_EDGE_ITEM_TYPES,
   LOCAL_STORAGE_KEY_SHOW_ON_NODE_ITEM_TYPES,
 } from './common/consts';
-import {Graph, GraphCollection} from './common/input_graph';
+import {Graph, GraphCollection, GraphWithLevel} from './common/input_graph';
 import {ModelGraph, ModelNode} from './common/model_graph';
 import {
   AddSnapshotInfo,
@@ -174,7 +174,12 @@ export class AppService {
               for (const subgraphId of node.subgraphIds) {
                 const subgraph = graphById[subgraphId];
                 if (subgraph) {
-                  subgraph.parentGraphId = graph.id;
+                  if (subgraph.parentGraphIds == null) {
+                    subgraph.parentGraphIds = [];
+                  }
+                  if (subgraph.parentGraphIds.includes(graph.id)) {
+                    subgraph.parentGraphIds.push(graph.id);
+                  }
                 }
               }
             }
@@ -183,11 +188,11 @@ export class AppService {
 
         // Find the 'root' graphs.
         const rootGraphs: Graph[] = collection.graphs.filter(
-          (graph) => graph.parentGraphId == null,
+          (graph) => graph.parentGraphIds == null,
         );
 
         // DFS from root graphs.
-        const dfsOrderedGraphs: Graph[] = [];
+        const dfsOrderedGraphs: GraphWithLevel[] = [];
         const visitGraph = (root?: Graph, level = 0) => {
           let graphs: Graph[] = [];
           if (root == null) {
@@ -197,16 +202,27 @@ export class AppService {
               .map((id) => graphById[id])
               .filter((graphs) => graphs != null);
           }
+
+          // Dedup graphs by their ids.
+          const uniqueGraphs: Graph[] = [];
+          const seenIds: Record<string, boolean> = {};
+          for (const graph of graphs) {
+            if (!seenIds[graph.id]) {
+              uniqueGraphs.push(graph);
+              seenIds[graph.id] = true;
+            }
+          }
+          graphs = uniqueGraphs;
+
           // Sort by node count.
           graphs.sort((g1, g2) => g2.nodes.length - g1.nodes.length);
           for (const graph of graphs) {
-            graph.level = level;
-            dfsOrderedGraphs.push(graph);
+            dfsOrderedGraphs.push({graph, level});
             visitGraph(graph, level + 1);
           }
         };
         visitGraph();
-        collection.graphs = dfsOrderedGraphs;
+        collection.graphsWithLevel = dfsOrderedGraphs;
       }
       newCollections.push(...graphCollections);
       return newCollections;
@@ -275,20 +291,37 @@ export class AppService {
     graph: Graph,
     flattenLayers = false,
     initialLayout = true,
+    openToLeft = false,
   ) {
+    // Keep the current pane and remove the other pane when there are two panes.
+    if (this.panes().length === 2) {
+      this.panes.update((panes) => {
+        if (openToLeft) {
+          return [panes[1]];
+        } else {
+          return [panes[0]];
+        }
+      });
+    }
+
     // Add a new pane.
     const paneId = genUid();
     this.paneIdToGraph[paneId] = graph;
     this.panes.update((panes) => {
       const firstPane = panes[0];
       firstPane.widthFraction = 0.5;
-      panes.push({
+      const newPane: Pane = {
         id: paneId,
         widthFraction: 0.5,
         flattenLayers,
         showOnNodeItemTypes: {[paneId]: this.getSavedShowOnNodeItemTypes()},
         showOnEdgeItemTypes: {[paneId]: this.getSavedShowOnEdgeItemTypes()},
-      });
+      };
+      if (openToLeft) {
+        panes.unshift(newPane);
+      } else {
+        panes.push(newPane);
+      }
       return [...panes];
     });
 
@@ -329,6 +362,11 @@ export class AppService {
     this.workerService.worker.postMessage(processGraphReq);
   }
 
+  getIsGraphInRightPane(graphId: string): boolean {
+    const panes = this.panes();
+    return panes.length === 2 && panes[1].modelGraph?.id === graphId;
+  }
+
   processGraph(
     paneId: string,
     flattenLayers = false,

src/ui/src/components/visualizer/common/input_graph.ts

@@ -32,6 +32,12 @@ export declare interface GraphCollection {
 
   /** The graphs inside the collection. */
   graphs: Graph[];
+
+  //////////////////////////////////////////////////////////////////////////////
+  // The following fields are set by model explorer. Users don't need to set
+  // them.
+
+  graphsWithLevel?: GraphWithLevel[];
 }
 
 /** The collection sent from the built-in adapters. */
@@ -83,12 +89,14 @@ export declare interface Graph {
   // The ids of all its subgraphs.
   subGraphIds?: string[];
 
-  // The id of its parent graph. We assume each subgraph only belongs to one
-  // parent.
-  parentGraphId?: string;
+  // The ids of its parent graphs.
+  parentGraphIds?: string[];
+}
 
-  // The level in the graph tree.
-  level?: number;
+/** A graph with its level, used in the graph selector. */
+export declare interface GraphWithLevel {
+  graph: Graph;
+  level: number;
 }
 
 /** A single node in the graph. */

src/ui/src/components/visualizer/graph_selector.ts

@@ -52,6 +52,7 @@ export interface GraphCollectionItem {
 export interface GraphItem {
   id: string;
   graph: Graph;
+  level: number;
   nonHiddenNodeCount: number;
   width: number;
 }
@@ -109,7 +110,7 @@ export class GraphSelector {
         collection,
         graphs: [],
       };
-      for (const graph of collection.graphs) {
+      for (const {graph, level} of collection.graphsWithLevel ?? []) {
         if (filterText !== '' && !graph.id.toLowerCase().includes(filterText)) {
           continue;
         }
@@ -122,6 +123,7 @@ export class GraphSelector {
         collectionItem.graphs.push({
           id: graph.id,
           graph,
+          level,
           nonHiddenNodeCount,
           width,
         });

src/ui/src/components/visualizer/graph_selector_panel.ts

@@ -101,15 +101,15 @@ export class GraphSelectorPanel {
   }
 
   showIndentSymbol(graphItem: GraphItem): boolean {
-    return !this.hasFilteredOutGraphs && (graphItem.graph.level ?? 0) > 0;
+    return !this.hasFilteredOutGraphs && (graphItem.level ?? 0) > 0;
   }
 
   getGraphItemPaddingLeft(graphItem: GraphItem): number {
     // Don't show tree indentation in filter mode.
     if (this.hasFilteredOutGraphs) {
       return DEFAULT_PADDING_LEFT;
     }
-    return DEFAULT_PADDING_LEFT + (graphItem.graph.level ?? 0) * 12;
+    return DEFAULT_PADDING_LEFT + (graphItem.level ?? 0) * 12;
   }
 
   trackByCollection(

src/ui/src/components/visualizer/webgl_renderer.ng.html

@@ -99,10 +99,11 @@
       [style.left.px]="subgraphIndicatorLeft"
       [style.width.px]="subgraphIndicatorWidth"
       [style.height.px]="subgraphIndicatorHeight"
-      matTooltip="Jump to subgraph"
+      [matTooltip]="subgraphIndicatorTooltip"
+      matTooltipClass="multiline-tooltip-left"
       matTooltipPosition="above"
       (mousedown)="$event.stopPropagation()"
-      (click)="handleClickSubgraphIndicator()">
+      (click)="handleClickSubgraphIndicator($event)">
   </div>
 
   <!-- Hidden menu trigger for subgraph dropdown -->
@@ -117,11 +118,14 @@
   <mat-menu #menu="matMenu">
     @for (subgraphId of curSubgraphIdsForMenu; track subgraphId) {
       <div class="model-explorer-menu-item-with-icon"
-          (click)="handleClickSubgraphId(subgraphId)">
+          (click)="handleClickSubgraphId(subgraphId, $event)">
         <mat-icon>subdirectory_arrow_right</mat-icon>
         {{subgraphId}}
       </div>
     }
+    <div class="model-explorer-alt-click-info">
+      Alt-click to open in split pane
+    </div>
   </mat-menu>
 
   <!-- Range zoom -->

src/ui/src/components/visualizer/webgl_renderer.scss

@@ -210,3 +210,17 @@
     background-color: rgba(0, 0, 0, 0.04);
   }
 }
+
+::ng-deep .model-explorer-alt-click-info {
+  height: 32px;
+  min-height: 32px;
+  background-color: #f6f6f6;
+  border-top: 1px solid #ddd;
+  box-sizing: border-box;
+  padding: 0 10px;
+  font-size: 12px;
+  display: flex;
+  align-items: center;
+  color: #777;
+  margin-bottom: -8px;
+}

src/ui/src/components/visualizer/webgl_renderer.ts

@@ -1276,7 +1276,7 @@ export class WebglRenderer implements OnInit, OnDestroy {
     this.webglRendererIoHighlightService.handleClickIoPicker(isInput, nodeId);
   }
 
-  handleClickSubgraphIndicator() {
+  handleClickSubgraphIndicator(event: MouseEvent) {
     if (!this.hoveredSubgraphIndicatorId) {
       return;
     }
@@ -1295,7 +1295,7 @@ export class WebglRenderer implements OnInit, OnDestroy {
     // directly.
     const subgraphIds = node.subgraphIds!;
     if (subgraphIds.length === 1) {
-      this.openSubgraph(subgraphIds[0]);
+      this.clickSubgraph(subgraphIds[0], event);
     }
     // If there are multiple subgraphs linked to the node, open a menu to let
     // users select a subgraph to jump to.
@@ -1305,8 +1305,8 @@ export class WebglRenderer implements OnInit, OnDestroy {
     }
   }
 
-  handleClickSubgraphId(subgraphId: string) {
-    this.openSubgraph(subgraphId);
+  handleClickSubgraphId(subgraphId: string, event: MouseEvent) {
+    this.clickSubgraph(subgraphId, event);
   }
 
   handleDoubleClickOnGraph(altDown: boolean, shiftDown: boolean) {
@@ -1719,6 +1719,25 @@ export class WebglRenderer implements OnInit, OnDestroy {
     return this.webglRendererThreejsService.fps;
   }
 
+  get subgraphIndicatorTooltip(): string {
+    if (!this.hoveredSubgraphIndicatorId) {
+      return '';
+    }
+    const node = this.curModelGraph.nodesById[
+      this.hoveredSubgraphIndicatorId
+    ] as OpNode;
+    if (!isOpNode(node)) {
+      return '';
+    }
+
+    const subgraphIds = node.subgraphIds!;
+    if (subgraphIds.length === 1) {
+      return `Jump to subgraph "${subgraphIds[0]}"\n(alt-click to open in split pane)`;
+    } else {
+      return 'Jump to subgraph';
+    }
+  }
+
   private handleSelectNode(nodeId: string, triggerNavigationSync = true) {
     this.appService.selectNode(this.paneId, {
       nodeId,
@@ -3208,4 +3227,20 @@ export class WebglRenderer implements OnInit, OnDestroy {
     }
     return [...deepestExpandedGroupNodeIdsSet];
   }
+
+  private clickSubgraph(subgraphId: string, event: MouseEvent) {
+    if (!event.altKey) {
+      this.openSubgraph(subgraphId);
+    }
+    // Alt-clicking opens the subgraph in a split pane.
+    else {
+      const subgraph = this.appService.getGraphById(subgraphId);
+      if (subgraph) {
+        const openToLeft = this.appService.getIsGraphInRightPane(
+          this.curModelGraph.id,
+        );
+        this.appService.openGraphInSplitPane(subgraph, false, true, openToLeft);
+      }
+    }
+  }
 }