Improve resource DataKey serialization perf a bit

Address the TODO about using toString in the
TreeMap comparator by using DataKey#compareTo.
Also, use the entrySet to iterate K-V pairs
instead of calling get(key).

Synthetic benchmark w/ 10000 random keys:
Before: 260ms to serialize
After: 33ms to serialize

Less of a difference when there are few keys.

--
MOS_MIGRATED_REVID=128469407

4 files changed, 47 insertions, 25 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/AndroidDataSerializer.java b/src/tools/android/java/com/google/devtools/build/android/AndroidDataSerializer.java
index 89120c3f24..1b1c0fbb19 100644
--- a/src/tools/android/java/com/google/devtools/build/android/AndroidDataSerializer.java
+++ b/src/tools/android/java/com/google/devtools/build/android/AndroidDataSerializer.java
@@ -14,12 +14,10 @@
 package com.google.devtools.build.android;
 
 import com.google.common.base.Stopwatch;
-import com.google.common.collect.Ordering;
 import com.google.devtools.build.android.ParsedAndroidData.KeyValueConsumer;
 import com.google.devtools.build.android.proto.SerializeFormat;
 import com.google.devtools.build.android.proto.SerializeFormat.Header;
 import com.google.protobuf.InvalidProtocolBufferException;
-
 import java.io.BufferedOutputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -33,7 +31,7 @@ import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.NavigableMap;
-import java.util.NavigableSet;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
@@ -44,9 +42,7 @@ import java.util.logging.Logger;
 public class AndroidDataSerializer {
   private static final Logger logger = Logger.getLogger(AndroidDataSerializer.class.getName());
 
-  // TODO(corysmith): We might need a more performant comparison methodology than toString.
-  private final NavigableMap<DataKey, DataValue> entries =
-      new TreeMap<DataKey, DataValue>(Ordering.usingToString());
+  private final NavigableMap<DataKey, DataValue> entries = new TreeMap<>();
 
   public static AndroidDataSerializer create() {
     return new AndroidDataSerializer();
@@ -89,20 +85,21 @@ public class AndroidDataSerializer {
 
   private void writeKeyValuesTo(NavigableMap<DataKey, DataValue> map, OutputStream outStream)
       throws IOException {
-    NavigableSet<DataKey> keys = map.navigableKeySet();
-    int[] orderedValueSizes = new int[keys.size()];
+    Set<Entry<DataKey, DataValue>> entries = map.entrySet();
+    int[] orderedValueSizes = new int[entries.size()];
     int valueSizeIndex = 0;
     // Serialize all the values in sorted order to a intermediate buffer, so that the keys
     // can be associated with a value size.
     // TODO(corysmith): Tune the size of the byte array.
     ByteArrayOutputStream valuesOutputStream = new ByteArrayOutputStream(2048);
-    for (DataKey key : keys) {
-      orderedValueSizes[valueSizeIndex++] = map.get(key).serializeTo(key, valuesOutputStream);
+    for (Map.Entry<DataKey, DataValue> entry : entries) {
+      orderedValueSizes[valueSizeIndex++] = entry.getValue()
+          .serializeTo(entry.getKey(), valuesOutputStream);
     }
     // Serialize all the keys in sorted order
     valueSizeIndex = 0;
-    for (DataKey key : keys) {
-      key.serializeTo(outStream, orderedValueSizes[valueSizeIndex++]);
+    for (Map.Entry<DataKey, DataValue> entry : entries) {
+      entry.getKey().serializeTo(outStream, orderedValueSizes[valueSizeIndex++]);
     }
     // write the values to the output stream.
     outStream.write(valuesOutputStream.toByteArray());
diff --git a/src/tools/android/java/com/google/devtools/build/android/DataKey.java b/src/tools/android/java/com/google/devtools/build/android/DataKey.java
index e3a8b9d987..0fbbd1f9ed 100644
--- a/src/tools/android/java/com/google/devtools/build/android/DataKey.java
+++ b/src/tools/android/java/com/google/devtools/build/android/DataKey.java
@@ -26,7 +26,8 @@ import java.io.OutputStream;
  *
  * For Assets, it is the asset path from the assets directory.
  */
-public interface DataKey {
+public interface DataKey extends Comparable<DataKey> {
+
   /**
    * Writes the Key and the value size to a stream.
    *
@@ -40,4 +41,14 @@ public interface DataKey {
    * Returns a human readable string representation of the key.
    */
   String toPrettyString();
+
+  /**
+   * Defines a total ordering on the different key types to assist in compareTo operations.
+   */
+  enum KeyType {
+    ASSET_PATH,
+    FULL_QUALIFIED_NAME
+  }
+
+  KeyType getKeyType();
 }
diff --git a/src/tools/android/java/com/google/devtools/build/android/FullyQualifiedName.java b/src/tools/android/java/com/google/devtools/build/android/FullyQualifiedName.java
index 4f953c8bd4..ad3fe2427e 100644
--- a/src/tools/android/java/com/google/devtools/build/android/FullyQualifiedName.java
+++ b/src/tools/android/java/com/google/devtools/build/android/FullyQualifiedName.java
@@ -13,6 +13,9 @@
 // limitations under the License.
 package com.google.devtools.build.android;
 
+import com.android.ide.common.resources.configuration.FolderConfiguration;
+import com.android.ide.common.resources.configuration.ResourceQualifier;
+import com.android.resources.ResourceType;
 import com.google.common.base.Joiner;
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
@@ -20,11 +23,6 @@ import com.google.common.collect.ImmutableList.Builder;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.PeekingIterator;
 import com.google.devtools.build.android.proto.SerializeFormat;
-
-import com.android.ide.common.resources.configuration.FolderConfiguration;
-import com.android.ide.common.resources.configuration.ResourceQualifier;
-import com.android.resources.ResourceType;
-
 import java.io.IOException;
 import java.io.OutputStream;
 import java.lang.reflect.InvocationTargetException;
@@ -40,7 +38,6 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Logger;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
-
 import javax.annotation.CheckReturnValue;
 import javax.annotation.concurrent.Immutable;
 
@@ -50,7 +47,7 @@ import javax.annotation.concurrent.Immutable;
  * Each resource name consists of the resource package, name, type, and qualifiers.
  */
 @Immutable
-public class FullyQualifiedName implements DataKey, Comparable<FullyQualifiedName> {
+public class FullyQualifiedName implements DataKey {
   public static final String DEFAULT_PACKAGE = "res-auto";
   private static final Joiner DASH_JOINER = Joiner.on('-');
 
@@ -440,7 +437,11 @@ public class FullyQualifiedName implements DataKey, Comparable<FullyQualifiedNam
   }
 
   @Override
-  public int compareTo(FullyQualifiedName other) {
+  public int compareTo(DataKey otherKey) {
+    if (!(otherKey instanceof FullyQualifiedName)) {
+      return getKeyType().compareTo(otherKey.getKeyType());
+    }
+    FullyQualifiedName other = (FullyQualifiedName) otherKey;
     if (!pkg.equals(other.pkg)) {
       return pkg.compareTo(other.pkg);
     }
@@ -461,6 +462,11 @@ public class FullyQualifiedName implements DataKey, Comparable<FullyQualifiedNam
   }
 
   @Override
+  public KeyType getKeyType() {
+    return KeyType.FULL_QUALIFIED_NAME;
+  }
+
+  @Override
   public void serializeTo(OutputStream out, int valueSize) throws IOException {
     toSerializedBuilder().setValueSize(valueSize).build().writeDelimitedTo(out);
   }
diff --git a/src/tools/android/java/com/google/devtools/build/android/RelativeAssetPath.java b/src/tools/android/java/com/google/devtools/build/android/RelativeAssetPath.java
index b6f595db4e..9940d13d27 100644
--- a/src/tools/android/java/com/google/devtools/build/android/RelativeAssetPath.java
+++ b/src/tools/android/java/com/google/devtools/build/android/RelativeAssetPath.java
@@ -16,7 +16,6 @@ package com.google.devtools.build.android;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.devtools.build.android.proto.SerializeFormat;
-
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.file.FileSystem;
@@ -32,7 +31,7 @@ import java.util.Objects;
  *
  * Note: Assets have no qualifiers or packages.
  */
-public class RelativeAssetPath implements DataKey, Comparable<RelativeAssetPath> {
+public class RelativeAssetPath implements DataKey {
   /**
    * A Factory that creates RelativeAssetsPath objects whose paths are relative to a given path.
    */
@@ -103,8 +102,17 @@ public class RelativeAssetPath implements DataKey, Comparable<RelativeAssetPath>
   }
 
   @Override
-  public int compareTo(RelativeAssetPath relativeAssetPath) {
-    return this.relativeAssetPath.compareTo(relativeAssetPath.relativeAssetPath);
+  public int compareTo(DataKey otherKey) {
+    if (!(otherKey instanceof RelativeAssetPath)) {
+      return getKeyType().compareTo(otherKey.getKeyType());
+    }
+    RelativeAssetPath otherAssetPath = (RelativeAssetPath) otherKey;
+    return this.relativeAssetPath.compareTo(otherAssetPath.relativeAssetPath);
+  }
+
+  @Override
+  public KeyType getKeyType() {
+    return KeyType.ASSET_PATH;
   }
 
   @Override