From 1968b8b3f9316a8327176e25211c8cc2b01459bc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 6 May 2017 20:46:07 -0800
Subject: [PATCH] Only make 1 collapsible for each run.

Previously, we could possibly make several collapsibles for each run because of how tag names were broken up.
Change: 155305237
---
 .../components/tf_dashboard_common/BUILD      |  1 +
 .../tf_dashboard_common/categorizer.ts        | 33 ++++++++++-------
 .../test/categorizerTest.ts                   | 12 +++++++
 .../tf-categorizer.ts                         | 35 +++++++++++--------
 4 files changed, 54 insertions(+), 27 deletions(-)

diff --git a/tensorflow/tensorboard/components/tf_dashboard_common/BUILD b/tensorflow/tensorboard/components/tf_dashboard_common/BUILD
index 880e6bd7125..efced92dcf3 100644
--- a/tensorflow/tensorboard/components/tf_dashboard_common/BUILD
+++ b/tensorflow/tensorboard/components/tf_dashboard_common/BUILD
@@ -46,6 +46,7 @@ tensorboard_typescript_genrule(
     ],
     typings = [
         "@org_definitelytyped//:d3.d.ts",
+        "@org_definitelytyped//:lodash.d.ts",
         "//tensorflow/tensorboard/components/vz_sorting:ts_typings",
     ],
 )
diff --git a/tensorflow/tensorboard/components/tf_dashboard_common/categorizer.ts b/tensorflow/tensorboard/components/tf_dashboard_common/categorizer.ts
index 42e7cbcff40..4c06462a981 100644
--- a/tensorflow/tensorboard/components/tf_dashboard_common/categorizer.ts
+++ b/tensorflow/tensorboard/components/tf_dashboard_common/categorizer.ts
@@ -72,24 +72,31 @@ module Categorizer {
       if (tags.length === 0) {
         return [];
       }
-      let sortedTags = tags.slice().sort(VZ.Sorting.compareTagNames);
-      let categories: Category[] = [];
-      let currentCategory = {
-        name: extractor(sortedTags[0]),
-        tags: [],
-      };
-      sortedTags.forEach((t: string) => {
-        let topLevel = extractor(t);
-        if (currentCategory.name !== topLevel) {
-          categories.push(currentCategory);
-          currentCategory = {
+
+      // Maps between top-level name and category. We use the mapping to avoid
+      // duplicating categories per run.
+      const categoryMapping: {[key: string]: Category} = {};
+
+      tags.forEach((t: string) => {
+        const topLevel = extractor(t);
+        if (!categoryMapping[topLevel]) {
+          const newCategory = {
             name: topLevel,
             tags: [],
           };
+          categoryMapping[topLevel] = newCategory;
         }
-        currentCategory.tags.push(t);
+
+        categoryMapping[topLevel].tags.push(t);
+      });
+
+      // Sort categories into alphabetical order.
+      const categories =
+          _.map(_.keys(categoryMapping).sort(), key => categoryMapping[key]);
+      _.forEach(categories, (category) => {
+        // Sort the tags within each category.
+        category.tags.sort(VZ.Sorting.compareTagNames);
       });
-      categories.push(currentCategory);
       return categories;
     };
   }
diff --git a/tensorflow/tensorboard/components/tf_dashboard_common/test/categorizerTest.ts b/tensorflow/tensorboard/components/tf_dashboard_common/test/categorizerTest.ts
index ea149fda47a..4e52b60f37f 100644
--- a/tensorflow/tensorboard/components/tf_dashboard_common/test/categorizerTest.ts
+++ b/tensorflow/tensorboard/components/tf_dashboard_common/test/categorizerTest.ts
@@ -62,6 +62,18 @@ module Categorizer {
         assert.deepEqual(
             topLevelNamespaceCategorizer(['a']), [{name: 'a', tags: ['a']}]);
       });
+
+      it('only create 1 category per run', () => {
+        // TensorBoard separates runs from tags using the / and _ characters
+        // *only* during sorting. The categorizer should group all tags under
+        // their correct categories - and create only 1 category per run.
+        const tags = ['foo/bar', 'foo_in_between_run/baz', 'foo/quux'];
+        const expected = [
+          {name: 'foo', tags: ['foo/bar', 'foo/quux']},
+          {name: 'foo_in_between_run', tags: ['foo_in_between_run/baz']},
+        ];
+        assert.deepEqual(topLevelNamespaceCategorizer(tags), expected);
+      });
     });
 
     describe('customCategorizer', () => {
diff --git a/tensorflow/tensorboard/components/tf_dashboard_common_d3v4/tf-categorizer.ts b/tensorflow/tensorboard/components/tf_dashboard_common_d3v4/tf-categorizer.ts
index 10271d2c085..c9da9a17eb0 100644
--- a/tensorflow/tensorboard/components/tf_dashboard_common_d3v4/tf-categorizer.ts
+++ b/tensorflow/tensorboard/components/tf_dashboard_common_d3v4/tf-categorizer.ts
@@ -73,24 +73,31 @@ function extractorToCategorizer(extractor: (s: string) => string): Categorizer {
     if (tags.length === 0) {
       return [];
     }
-    let sortedTags = tags.slice().sort(compareTagNames);
-    let categories: Category[] = [];
-    let currentCategory = {
-      name: extractor(sortedTags[0]),
-      tags: [],
-    };
-    sortedTags.forEach((t: string) => {
-      let topLevel = extractor(t);
-      if (currentCategory.name !== topLevel) {
-        categories.push(currentCategory);
-        currentCategory = {
+
+    // Maps between top-level name and category. We use the mapping to avoid
+    // duplicating categories per run.
+    const categoryMapping: {[key: string]: Category} = {};
+
+    tags.forEach((t: string) => {
+      const topLevel = extractor(t);
+      if (!categoryMapping[topLevel]) {
+        const newCategory = {
           name: topLevel,
           tags: [],
         };
+        categoryMapping[topLevel] = newCategory;
       }
-      currentCategory.tags.push(t);
+
+      categoryMapping[topLevel].tags.push(t);
+    });
+
+    // Sort categories into alphabetical order.
+    const categories =
+        _.map(_.keys(categoryMapping).sort(), key => categoryMapping[key]);
+    _.forEach(categories, (category) => {
+      // Sort the tags within each category.
+      category.tags.sort(compareTagNames);
     });
-    categories.push(currentCategory);
     return categories;
   };
 }
@@ -180,4 +187,4 @@ Polymer({
       this._setCategories(categories);
     })
   },
-});
\ No newline at end of file
+});