Handle overlapping edges better in edge overlay.

PiperOrigin-RevId: 686520719

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

@@ -983,12 +983,19 @@ export function getMultiLineLabelExtraHeight(label: string): number {
  * Calculates the closest intersection points of a line (L) connecting
  * the centers of two rectangles (rect1 and rect2) with the sides of these
  * rectangles.
+ *
+ * xOffsetFactor is used to shift the center of the rectangle to the left or
+ * right by a certain factor of the width of the rectangle.
  */
-export function getIntersectionPoints(rect1: Rect, rect2: Rect) {
+export function getIntersectionPoints(
+ rect1: Rect,
+ rect2: Rect,
+ xOffsetFactor = 0,
+) {
  // Function to calculate the center of a rectangle
  function getCenter(rect: Rect) {
  return {
- x: rect.x + rect.width / 2,
+ x: rect.x + rect.width / 2 + xOffsetFactor * rect.width,
  y: rect.y + rect.height / 2,
  };
  }

src/ui/src/components/visualizer/webgl_renderer_edge_overlays_service.ts

@@ -82,6 +82,23 @@ export class WebglRendererEdgeOverlaysService {
  return;
  }
 
+ // Keep track of number of edges for a given pair of nodes. If there are
+ // more than 1 edges, we will shift the edges to avoid overlapping.
+ //
+ // From sorted edge key (nodeId1->nodeId2) to the number of edges for that
+ // pair.
+ const seenEdgePairs: Record<string, number> = {};
+ const totalEdgePairs: Record<string, number> = {};
+
+ // Populate totalEdgePairs.
+ for (let i = 0; i < this.curOverlays.length; i++) {
+ const subgraph = this.curOverlays[i];
+ for (const {sourceNodeId, targetNodeId, label} of subgraph.edges) {
+ this.addToEdgePairs(sourceNodeId, targetNodeId, totalEdgePairs);
+ }
+ }
+ console.log(totalEdgePairs);
+
  for (let i = 0; i < this.curOverlays.length; i++) {
  const subgraph = this.curOverlays[i];
  const edgeWidth = subgraph.edgeWidth ?? DEFAULT_EDGE_WIDTH;
@@ -98,9 +115,18 @@ export class WebglRendererEdgeOverlaysService {
  const targetNode = this.webglRenderer.curModelGraph.nodesById[
  targetNodeId
  ] as OpNode;
+ const curEdgesCount = this.addToEdgePairs(
+ sourceNodeId,
+ targetNodeId,
+ seenEdgePairs,
+ );
+ const totalEdgesCount =
+ totalEdgePairs[this.getEdgeKey(sourceNodeId, targetNodeId)];
+ const xOffsetFactor = (1 / (totalEdgesCount + 1)) * curEdgesCount - 0.5;
  const {intersection1, intersection2} = getIntersectionPoints(
  this.webglRenderer.getNodeRect(sourceNode),
  this.webglRenderer.getNodeRect(targetNode),
+ xOffsetFactor,
  );
  // Edge.
  edges.push({
@@ -190,4 +216,23 @@ export class WebglRendererEdgeOverlaysService {
  }
  return [...ids];
  }
+
+ private addToEdgePairs(
+ nodeId1: string,
+ nodeId2: string,
+ pairs: Record<string, number>,
+ ): number {
+ const key = this.getEdgeKey(nodeId1, nodeId2);
+ if (pairs[key] === undefined) {
+ pairs[key] = 0;
+ }
+ pairs[key]++;
+ return pairs[key];
+ }
+
+ private getEdgeKey(nodeId1: string, nodeId2: string): string {
+ return nodeId1.localeCompare(nodeId2) < 0
+ ? `${nodeId1}___${nodeId2}`
+ : `${nodeId2}___${nodeId1}`;
+ }
 }