Refactor ModelGraphVisualizer config to use signals and add highlight toggle.

PiperOrigin-RevId: 880804502

Author: Google AI Edge

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

@@ -21,6 +21,9 @@ import {IS_EXTERNAL} from '../../../common/flags';
 /** The height of the node label. */
 export const DEFAULT_NODE_LABEL_HEIGHT = 11;
 
+/** The padding of the label. */
+export const LABEL_PADDING = 24;
+
 /**
  * The padding between the label and the value in node's attrs table.
  */

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

@@ -1038,14 +1038,99 @@ export function splitLabel(label: string): string[] {
     .filter((line) => line !== '');
 }
 
+/** Wraps the label based on the given max width. */
+export function wrapLabel(
+  label: string,
+  maxWidth: number,
+  fontSize: number,
+  bold = false,
+): string[] {
+  const sections = label.split('\n');
+  const allLines: string[] = [];
+
+  for (const section of sections) {
+    if (section === '') {
+      allLines.push('');
+      continue;
+    }
+
+    const words = section.split(' ');
+    const lines: string[] = [];
+    let currentLine = words[0];
+
+    for (let i = 1; i < words.length; i++) {
+      const word = words[i];
+      const width = getLabelWidth(currentLine + ' ' + word, fontSize, bold);
+      if (width < maxWidth) {
+        currentLine += ' ' + word;
+      } else {
+        lines.push(currentLine);
+        currentLine = word;
+      }
+    }
+    lines.push(currentLine);
+
+    // Post-process to split any lines that are still too long (single words).
+    const finalLines: string[] = [];
+    for (const line of lines) {
+      if (getLabelWidth(line, fontSize, bold) <= maxWidth) {
+        finalLines.push(line);
+        continue;
+      }
+      // Split long word.
+      let curLine = line;
+      while (getLabelWidth(curLine, fontSize, bold) > maxWidth) {
+        // Find split point.
+        let low = 0;
+        let high = curLine.length;
+        let splitIndex = 0;
+        while (low <= high) {
+          const mid = Math.floor((low + high) / 2);
+          const sub = curLine.substring(0, mid);
+          if (getLabelWidth(sub, fontSize, bold) <= maxWidth) {
+            splitIndex = mid;
+            low = mid + 1;
+          } else {
+            high = mid - 1;
+          }
+        }
+        if (splitIndex === 0) {
+          // If even one char is too wide, just take one char (shouldn't happen with reasonable width).
+          splitIndex = 1;
+        }
+        finalLines.push(curLine.substring(0, splitIndex));
+        curLine = curLine.substring(splitIndex);
+      }
+      if (curLine.length > 0) {
+        finalLines.push(curLine);
+      }
+    }
+    allLines.push(...finalLines);
+  }
+
+  return allLines;
+}
+
 /** Get the extra height for multi-line label. */
 export function getMultiLineLabelExtraHeight(
   node: ModelNode,
   config?: VisualizerConfig,
 ): number {
-  return (
-    (splitLabel(node.label).length - 1) * getNodeLabelLineHeight(node, config)
-  );
+  let lineCount = 0;
+  if (config?.nodeLabelWidth) {
+    const lines = wrapLabel(
+      node.label,
+      config.nodeLabelWidth,
+      getNodeLabelHeight(node, config),
+      !isOpNode(node),
+    );
+    if (lines.length >= 1) {
+      return (lines.length - 1) * getNodeLabelLineHeight(node, config);
+    }
+  } else {
+    lineCount = splitLabel(node.label).length;
+  }
+  return (lineCount - 1) * getNodeLabelLineHeight(node, config);
 }
 
 /**
@@ -1329,9 +1414,14 @@ export function getLayoutMarginTop(
   config?: VisualizerConfig,
 ): number {
   const nodeLabelHeight = getNodeLabelHeight(node, config);
+  let extraHeight = 0;
+  if (config?.nodeLabelWidth) {
+    extraHeight = getMultiLineLabelExtraHeight(node, config);
+  }
+
   if (nodeLabelHeight === 11) {
-    return 36;
+    return 36 + extraHeight;
   }
   const nodeLabelYPadding = getNodeLabelYPadding(node, config);
-  return nodeLabelYPadding + nodeLabelHeight + 16;
+  return nodeLabelYPadding + nodeLabelHeight + extraHeight + 16;
 }

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

@@ -72,6 +72,14 @@ export declare interface VisualizerConfig {
   /** The maximum number of child nodes under a layer node. Default: 400. */
   artificialLayerNodeCountThreshold?: number;
 
+  /**
+   * The maximum width for node labels.
+   *
+   * If set, node labels will be wrapped to this width.
+   * If unset (default), the node labels will have a default maximum width.
+   */
+  nodeLabelWidth?: number;
+
   /** The font size of the edge label. Default: 7.5. */
   edgeLabelFontSize?: number;
 

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

@@ -103,6 +103,7 @@ import {
   processNodeStylerRules,
   splitLabel,
   splitNamespace,
+  wrapLabel,
 } from './common/utils';
 import {
   ExpandOrCollapseGroupNodeRequest,
@@ -2552,12 +2553,19 @@ export class WebglRenderer implements OnInit, OnChanges, OnDestroy {
 
         // Expand icon.
         const iconZ = y + this.getNodeLabelRelativeY(node) + 18.5;
-        const leftIconX = node.expanded
-          ? labelLeft - 13
-          : (x + labelLeft + 1) / 2 + 1;
-        const rightIconX = node.expanded
-          ? labelRight + 12
-          : (x + width + labelRight - 1) / 2 - 1;
+        let leftIconX = 0;
+        let rightIconX = 0;
+        if (this.appService.config()?.nodeLabelWidth) {
+          leftIconX = x + 12;
+          rightIconX = x + width - 12;
+        } else {
+          leftIconX = node.expanded
+            ? labelLeft - 13
+            : (x + labelLeft + 1) / 2 + 1;
+          rightIconX = node.expanded
+            ? labelRight + 12
+            : (x + width + labelRight - 1) / 2 - 1;
+        }
         groupNodeIcons.push({
           id: node.id,
           nodeId: node.id,
@@ -2784,7 +2792,19 @@ export class WebglRenderer implements OnInit, OnChanges, OnDestroy {
 
       // Font size.
       const labelHeight = getNodeLabelHeight(node, this.appService.config());
-      const lines = splitLabel(this.getNodeLabel(node));
+      let lines: string[] = [];
+      const config = this.appService.config();
+      if (config?.nodeLabelWidth) {
+        lines = wrapLabel(
+          this.getNodeLabel(node),
+          config.nodeLabelWidth,
+          labelHeight,
+          !isOpNode(node),
+        );
+      } else {
+        lines = splitLabel(this.getNodeLabel(node));
+      }
+
       for (let i = 0; i < lines.length; i++) {
         const curLineLabel = lines[i];
         labels.push({

src/ui/src/components/visualizer/webgl_renderer_search_results_service.ts

@@ -27,7 +27,9 @@ import {
   getNodeLabelHeight,
   getNodeLabelLineHeight,
   isGroupNode,
+  isOpNode,
   splitLabel,
+  wrapLabel,
 } from './common/utils';
 import {
   ColorVariable,
@@ -177,8 +179,19 @@ export class WebglRendererSearchResultsService {
           let y = 0;
           let height = 0;
           let width = 0;
-          const lines = splitLabel(node.label);
-          if (lines.length === 1) {
+          const config = this.webglRenderer.appService.config();
+          let lines: string[] = [];
+          if (config?.nodeLabelWidth) {
+            lines = wrapLabel(
+              node.label,
+              config.nodeLabelWidth,
+              nodeLabelHeight,
+              !isOpNode(node),
+            );
+          } else {
+            lines = splitLabel(node.label);
+          }
+          if (lines.length === 1 && !config?.nodeLabelWidth) {
             const labelSizes = this.webglRenderer.texts.getLabelSizes(
               node.label,
               isGroupNode(node) ? FontWeight.BOLD : FontWeight.MEDIUM,
@@ -191,8 +204,19 @@ export class WebglRendererSearchResultsService {
               this.webglRenderer.getNodeLabelRelativeY(node) -
               2 * scale;
           } else {
-            const {minX, maxX} = this.webglRenderer.getNodeLabelSizes(node);
-            width = (maxX - minX) * scale + 4;
+            let maxLineWidth = 0;
+            for (const line of lines) {
+              const labelSizes = this.webglRenderer.texts.getLabelSizes(
+                line,
+                isGroupNode(node) ? FontWeight.BOLD : FontWeight.MEDIUM,
+                nodeLabelHeight,
+              ).sizes;
+              maxLineWidth = Math.max(
+                maxLineWidth,
+                labelSizes.maxX - labelSizes.minX,
+              );
+            }
+            width = maxLineWidth * scale + 4;
             height =
               lines.length *
                 getNodeLabelLineHeight(

src/ui/src/components/visualizer/worker/graph_layout.ts

@@ -22,6 +22,7 @@ import {
   EXPANDED_NODE_DATA_PROVIDER_SUMMARY_ROW_HEIGHT,
   EXPANDED_NODE_DATA_PROVIDER_SUMMARY_TOP_PADDING,
   EXPANDED_NODE_DATA_PROVIDER_SYUMMARY_FONT_SIZE,
+  LABEL_PADDING,
   LAYOUT_MARGIN_X,
   MAX_IO_ROWS_IN_ATTRS_TABLE,
   NODE_ATTRS_TABLE_FONT_SIZE_TO_HEIGHT_RATIO,
@@ -69,6 +70,7 @@ import {
   isGroupNode,
   isOpNode,
   splitLabel,
+  wrapLabel,
 } from '../common/utils';
 import {VisualizerConfig} from '../common/visualizer_config';
 
@@ -83,8 +85,6 @@ export const LAYOUT_MARGIN_BOTTOM = 16;
 /** Node width for test cases. */
 export const NODE_WIDTH_FOR_TEST = 50;
 
-const LABEL_PADDING = 24;
-
 const MIN_NODE_WIDTH = 80;
 
 const ATTRS_TABLE_MARGIN_X = 8;
@@ -417,8 +417,8 @@ export function getNodeWidth(
   testMode = false,
   config?: VisualizerConfig,
 ) {
-  // Always return 32 in test mode.
-  if (testMode) {
+  // Always return 32 in test mode, unless nodeLabelWidth is set.
+  if (testMode && !config?.nodeLabelWidth) {
     return NODE_WIDTH_FOR_TEST;
   }
 
@@ -428,7 +428,17 @@ export function getNodeWidth(
     fontSize * NODE_ATTRS_TABLE_VALUE_MAX_CHAR_COUNT;
 
   const label = node.label;
-  const lines = splitLabel(label);
+  let lines: string[] = [];
+  if (config?.nodeLabelWidth) {
+    lines = wrapLabel(
+      label,
+      config.nodeLabelWidth,
+      getNodeLabelHeight(node, config),
+      isGroupNode(node),
+    );
+  } else {
+    lines = splitLabel(label);
+  }
   let labelWidth = 0;
   for (const line of lines) {
     labelWidth = Math.max(
@@ -601,7 +611,7 @@ export function getNodeHeight(
   forceRecalculate = false,
   config?: VisualizerConfig,
 ) {
-  if (testMode) {
+  if (testMode && !config?.nodeLabelWidth) {
     return NODE_HEIGHT_FOR_TEST;
   }