Don't layout the graph twice when restoring pane state.

PiperOrigin-RevId: 675702384
google-ai-edge · Sep 17, 2024 · 873e757 · 873e757
1 parent ae97c48
commit 873e757
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 9 deletions.
diff --git a/src/ui/src/components/visualizer/app_service.ts b/src/ui/src/components/visualizer/app_service.ts
@@ -213,6 +213,7 @@ export class AppService {
     paneIndex: number,
     flattenLayers = false,
     snapshot?: SnapshotData,
+    initialLayout = true,
   ) {
     if (paneIndex === 1 && this.panes().length === 1) {
       this.openGraphInSplitPane(graph);
@@ -247,23 +248,29 @@ export class AppService {
     }
 
     // Process the graph.
-    this.processGraph(paneId, flattenLayers, snapshot);
+    this.processGraph(paneId, flattenLayers, snapshot, initialLayout);
   }
 
   selectGraphInCurrentPane(
     graph: Graph,
     flattenLayers = false,
     snapshot?: SnapshotData,
+    initialLayout = true,
   ) {
     this.selectGraphInPane(
       graph,
       this.getPaneIndexById(this.selectedPaneId()),
       flattenLayers,
       snapshot,
+      initialLayout,
     );
   }
 
-  openGraphInSplitPane(graph: Graph, flattenLayers = false) {
+  openGraphInSplitPane(
+    graph: Graph,
+    flattenLayers = false,
+    initialLayout = true,
+  ) {
     // Add a new pane.
     const paneId = genUid();
     this.paneIdToGraph[paneId] = graph;
@@ -312,6 +319,7 @@ export class AppService {
         this.getGroupNodeChildrenCountThreshold(),
       flattenLayers,
       keepLayersWithASingleChild: this.config()?.keepLayersWithASingleChild,
+      initialLayout,
     };
     this.workerService.worker.postMessage(processGraphReq);
   }
@@ -320,6 +328,7 @@ export class AppService {
     paneId: string,
     flattenLayers = false,
     snapshotToRestore?: SnapshotData,
+    initialLayout = true,
   ) {
     // Store snapshotToResotre into pane if set.
     if (snapshotToRestore != null) {
@@ -344,6 +353,7 @@ export class AppService {
         this.getGroupNodeChildrenCountThreshold(),
       flattenLayers,
       keepLayersWithASingleChild: this.config()?.keepLayersWithASingleChild,
+      initialLayout,
     };
     this.workerService.worker.postMessage(processGraphReq);
   }

diff --git a/src/ui/src/components/visualizer/common/worker_events.ts b/src/ui/src/components/visualizer/common/worker_events.ts
@@ -59,6 +59,7 @@ export declare interface ProcessGraphRequest extends WorkerEventBase {
   groupNodeChildrenCountThreshold?: number;
   flattenLayers?: boolean;
   keepLayersWithASingleChild?: boolean;
+  initialLayout?: boolean;
 }
 
 /** The response for processing an input graph. */

diff --git a/src/ui/src/components/visualizer/model_graph_visualizer.ts b/src/ui/src/components/visualizer/model_graph_visualizer.ts
@@ -234,6 +234,9 @@ export class ModelGraphVisualizer implements OnInit, OnDestroy, OnChanges {
       // One pane.
       if (this.initialUiState.paneStates.length === 1) {
         const paneState = this.initialUiState.paneStates[0];
+        const initialLayout =
+          paneState.selectedNodeId === '' &&
+          paneState.deepestExpandedGroupNodeIds.length === 0;
         const selectedGraph = this.findGraphFromCollections(
           paneState.selectedCollectionLabel,
           paneState.selectedGraphId,
@@ -243,11 +246,18 @@ export class ModelGraphVisualizer implements OnInit, OnDestroy, OnChanges {
           this.appService.selectGraphInCurrentPane(
             selectedGraph,
             flattenLayers,
+            undefined,
+            initialLayout,
           );
         } else {
           // Fall back to first graph.
           const firstGraph = this.graphCollections[0].graphs[0];
-          this.appService.selectGraphInCurrentPane(firstGraph, flattenLayers);
+          this.appService.selectGraphInCurrentPane(
+            firstGraph,
+            flattenLayers,
+            undefined,
+            initialLayout,
+          );
         }
         this.appService.setFlattenLayersInCurrentPane(flattenLayers);
       }

diff --git a/src/ui/src/components/visualizer/webgl_renderer.ts b/src/ui/src/components/visualizer/webgl_renderer.ts
@@ -783,11 +783,18 @@ export class WebglRenderer implements OnInit, OnDestroy {
           );
           deepestExpandedGroupNodeIds = groupNodeIds;
         }
-        this.sendRelayoutGraphRequest(
-          paneState.selectedNodeId,
-          deepestExpandedGroupNodeIds,
-          true,
-        );
+        if (
+          paneState.selectedNodeId != '' ||
+          deepestExpandedGroupNodeIds.length > 0
+        ) {
+          this.sendRelayoutGraphRequest(
+            paneState.selectedNodeId,
+            deepestExpandedGroupNodeIds,
+            true,
+          );
+        } else {
+          initGraphFn();
+        }
         // This is needed for loading old perma-link.
         this.uiStateService.setDeepestExpandedGroupNodeIds(
           paneState.deepestExpandedGroupNodeIds,

diff --git a/src/ui/src/components/visualizer/worker/worker.ts b/src/ui/src/components/visualizer/worker/worker.ts
@@ -68,6 +68,7 @@ self.addEventListener('message', (event: Event) => {
         workerEvent.groupNodeChildrenCountThreshold,
         workerEvent.flattenLayers,
         workerEvent.keepLayersWithASingleChild,
+        workerEvent.initialLayout,
       );
       cacheModelGraph(modelGraph, workerEvent.paneId);
       const resp: ProcessGraphResponse = {
@@ -206,6 +207,7 @@ function handleProcessGraph(
   groupNodeChildrenCountThreshold?: number,
   flattenLayers?: boolean,
   keepLayersWithASingleChild?: boolean,
+  initialLayout?: boolean,
 ): ModelGraph {
   let error: string | undefined = undefined;
 
@@ -231,7 +233,7 @@ function handleProcessGraph(
   }
 
   // Do the initial layout.
-  if (!error) {
+  if (!error && initialLayout) {
     const layout = new GraphLayout(
       modelGraph,
       dagre,