From fb6b42b6373870aa047458335aa4ea5eda69e497 Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Tue, 5 Nov 2024 03:18:28 -0800
Subject: [PATCH] Add display limit to the attribute string

The custom op attributes may contain strings that are too long to display, and that may crash the browser. Truncate the string to a reasonable length.

PiperOrigin-RevId: 693280782
---
 src/builtin-adapter/translate_helpers.cc | 20 ++++++++++++++++----
 src/builtin-adapter/visualize_config.h   | 17 ++++++++++++++++-
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/src/builtin-adapter/translate_helpers.cc b/src/builtin-adapter/translate_helpers.cc
index 1f2bc568..b1ad5332 100644
--- a/src/builtin-adapter/translate_helpers.cc
+++ b/src/builtin-adapter/translate_helpers.cc
@@ -97,7 +97,8 @@ inline std::string GetTypeString(const mlir::Type& t) {
   return result;
 }
 
-void AppendNodeAttrs(const int const_element_count_limit, Operation& operation,
+void AppendNodeAttrs(const int const_element_count_limit,
+                     const int attribute_string_limit, Operation& operation,
                      GraphNodeBuilder& builder) {
   std::string value;
   llvm::raw_string_ostream sstream(value);
@@ -122,6 +123,14 @@ void AppendNodeAttrs(const int const_element_count_limit, Operation& operation,
       // Special handles `value` attribute to represent the tensor data.
       builder.AppendNodeAttribute(kValue, value);
     } else {
+      // The custom op attributes may contain strings that are too long to
+      // display, and that may crash the browser. Truncate the string to a
+      // reasonable length.
+      if (attribute_string_limit >= 0 &&
+          value.length() > attribute_string_limit) {
+        value.resize(attribute_string_limit);
+        absl::StrAppend(&value, "... [truncated]");
+      }
       builder.AppendNodeAttribute(name, value);
     }
     value.clear();
@@ -431,7 +440,8 @@ absl::StatusOr<Subgraph> TfFunctionToSubgraph(const VisualizeConfig& config,
     llvm::StringRef node_name = GetTfNodeName(operation);
     GraphNodeBuilder builder;
     builder.SetNodeInfo(node_id, node_label, node_name);
-    AppendNodeAttrs(config.const_element_count_limit, operation, builder);
+    AppendNodeAttrs(config.const_element_count_limit,
+                    config.attribute_string_limit, operation, builder);
     for (int input_index = 0; input_index < operation.getNumOperands();
          ++input_index) {
       mlir::Value val = operation.getOperand(input_index);
@@ -483,7 +493,8 @@ absl::StatusOr<Subgraph> TfliteFunctionToSubgraph(const VisualizeConfig& config,
     std::string node_name = GenerateTfliteNodeName(node_label, operation);
     GraphNodeBuilder builder;
     builder.SetNodeInfo(node_id, node_label, node_name);
-    AppendNodeAttrs(config.const_element_count_limit, operation, builder);
+    AppendNodeAttrs(config.const_element_count_limit,
+                    config.attribute_string_limit, operation, builder);
     absl::Status append_subgraph_status =
         TfliteMaybeAppendSubgraphs(operation, builder);
     if (!append_subgraph_status.ok()) {
@@ -565,7 +576,8 @@ absl::StatusOr<Subgraph> StablehloFunctionToSubgraph(
     builder.SetNodeId(node_id);
     builder.SetNodeLabel(node_label);
     AddJaxNodeNameAndAttribute(operation, builder);
-    AppendNodeAttrs(config.const_element_count_limit, operation, builder);
+    AppendNodeAttrs(config.const_element_count_limit,
+                    config.attribute_string_limit, operation, builder);
     absl::Status append_subgraph_status =
         StablehloMaybeAppendSubgraphs(operation);
     if (!append_subgraph_status.ok()) {
diff --git a/src/builtin-adapter/visualize_config.h b/src/builtin-adapter/visualize_config.h
index 57214ea6..a3becca5 100644
--- a/src/builtin-adapter/visualize_config.h
+++ b/src/builtin-adapter/visualize_config.h
@@ -22,12 +22,27 @@ namespace visualization_client {
 struct VisualizeConfig {
   VisualizeConfig() = default;
   explicit VisualizeConfig(const int const_element_count_limit)
-      : const_element_count_limit(const_element_count_limit) {}
+      : const_element_count_limit(const_element_count_limit) {
+    // If const_element_count_limit is set to -1, then we should also set
+    // attribute_string_limit to -1 to print all attributes.
+    if (const_element_count_limit == -1) {
+      attribute_string_limit = -1;
+    }
+  }
+  VisualizeConfig(const int const_element_count_limit,
+                  const int attribute_string_limit)
+      : const_element_count_limit(const_element_count_limit),
+        attribute_string_limit(attribute_string_limit) {}
 
   // The maximum number of constant elements to be displayed. If the number
   // exceeds this threshold, the rest of data will be elided. The default
   // threshold is set to 16 (use -1 to print all).
   int const_element_count_limit = 16;
+
+  // The maximum length of attribute string to be displayed. If the length
+  // exceeds this threshold, the rest of data will be elided. The default
+  // threshold is set to 1024 (use -1 to print all).
+  int attribute_string_limit = 1024;
 };
 
 }  // namespace visualization_client