From 1c730cdbcb7d1c2a64bbba2af5bc77368694b6c4 Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Wed, 16 Oct 2024 15:14:52 -0700
Subject: [PATCH] - Allow users to disable the fallback behavior when finding
 mapped node ids in sync navigation. - Use the reversed mapping to find mapped
 node when selecting node in right panel.

PiperOrigin-RevId: 686658582
---
 .../visualizer/common/sync_navigation.ts      | 19 +++++++++-
 .../visualizer/sync_navigation_service.ts     |  7 +++-
 .../components/visualizer/webgl_renderer.ts   | 38 +++++++++++--------
 3 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/src/ui/src/components/visualizer/common/sync_navigation.ts b/src/ui/src/components/visualizer/common/sync_navigation.ts
index b82a08e5..e2459afb 100644
--- a/src/ui/src/components/visualizer/common/sync_navigation.ts
+++ b/src/ui/src/components/visualizer/common/sync_navigation.ts
@@ -22,12 +22,27 @@ import {TaskData, TaskType} from './task';
 export declare interface SyncNavigationData extends TaskData {
   type: TaskType.SYNC_NAVIGATION;
 
+  /**
+   * Specifies the mapping for navigation syncing.
+   *
+   * When user selects a node on one side, Model Explorer will try to find the
+   * mapped node on the other side and select it automacitally. If the mapped
+   * node is not found, Model Explorer will try to find the node with the same
+   * node id on the other side. This fallback behavior can be disabled by
+   * setting `disableMappingFallback` below to true.
+   */
   mapping: SyncNavigationMapping;
+
+  /**
+   * Whether to disable the fallback behavior (find the node with the same id)
+   * when the mapped node is not found from the `mapping` field above.
+   */
+  disableMappingFallback?: boolean;
 }
 
 /**
- * The mapping for navigation syncing, from node id from one side to node id
- * from another side.
+ * The mapping for navigation syncing, from node id from left side to node id
+ * from right side.
  */
 export type SyncNavigationMapping = Record<string, string>;
 
diff --git a/src/ui/src/components/visualizer/sync_navigation_service.ts b/src/ui/src/components/visualizer/sync_navigation_service.ts
index af64156a..ad4d36ac 100644
--- a/src/ui/src/components/visualizer/sync_navigation_service.ts
+++ b/src/ui/src/components/visualizer/sync_navigation_service.ts
@@ -86,7 +86,12 @@ export class SyncNavigationService {
         // Fallback to the original node id if not found.
         const mapping = curSyncNavigationData?.mapping ?? {};
         const inversedMapping = curSyncNavigationData?.inversedMapping ?? {};
-        return mapping[nodeId] ?? inversedMapping[nodeId] ?? nodeId;
+        const targetMapping = paneIndex === 0 ? mapping : inversedMapping;
+        const mappedNodeId = targetMapping[nodeId];
+        if (mappedNodeId) {
+          return mappedNodeId;
+        }
+        return curSyncNavigationData?.disableMappingFallback ? '' : nodeId;
       }
       default:
         return nodeId;
diff --git a/src/ui/src/components/visualizer/webgl_renderer.ts b/src/ui/src/components/visualizer/webgl_renderer.ts
index db63e52d..acf381c5 100644
--- a/src/ui/src/components/visualizer/webgl_renderer.ts
+++ b/src/ui/src/components/visualizer/webgl_renderer.ts
@@ -740,27 +740,35 @@ export class WebglRenderer implements OnInit, OnDestroy {
         }
 
         if (data.paneIndex !== this.appService.getPaneIndexById(this.paneId)) {
-          const mappedNodeId = this.syncNavigationService.getMappedNodeId(
-            data.paneIndex,
-            data.nodeId,
-          );
-          const mappedNode = this.curModelGraph.nodesById[mappedNodeId];
-          const hideInLayout = isOpNode(mappedNode) && mappedNode.hideInLayout;
-          if (
-            mappedNode &&
-            mappedNode.id !== this.selectedNodeId &&
-            !hideInLayout
-          ) {
-            this.revealNode(mappedNodeId, false);
+          // Clicking on the empty space. Hide the no mapped node message.
+          if (data.nodeId === '') {
             this.syncNavigationService.showNoMappedNodeMessageTrigger$.next(
               undefined,
             );
-          } else {
-            if (mappedNodeId !== '' && (!mappedNode || hideInLayout)) {
+          }
+          // Clicking on a node.
+          else {
+            const mappedNodeId = this.syncNavigationService.getMappedNodeId(
+              data.paneIndex,
+              data.nodeId,
+            );
+            const mappedNode = this.curModelGraph.nodesById[mappedNodeId];
+            const hideInLayout =
+              isOpNode(mappedNode) && mappedNode.hideInLayout;
+            if (
+              mappedNode &&
+              mappedNode.id !== this.selectedNodeId &&
+              !hideInLayout
+            ) {
+              this.revealNode(mappedNodeId, false);
+              this.syncNavigationService.showNoMappedNodeMessageTrigger$.next(
+                undefined,
+              );
+            } else if (!mappedNode || hideInLayout) {
               this.syncNavigationService.showNoMappedNodeMessageTrigger$.next(
                 {},
               );
-            } else if (mappedNodeId === '') {
+            } else {
               this.syncNavigationService.showNoMappedNodeMessageTrigger$.next(
                 undefined,
               );