Optimize DexLimitTracker with strong string interning for faster Multidex sharding.

PiperOrigin-RevId: 974808907
Change-Id: Iedb642b5afd9ecab1abd34b8e641553d89a441b6

src/tools/java/com/google/devtools/build/android/dexer/DexLimitTracker.java

diff --git a/src/tools/java/com/google/devtools/build/android/dexer/DexLimitTracker.java b/src/tools/java/com/google/devtools/build/android/dexer/DexLimitTracker.java
index a836b07..6ff0251 100644
--- a/src/tools/java/com/google/devtools/build/android/dexer/DexLimitTracker.java
+++ b/src/tools/java/com/google/devtools/build/android/dexer/DexLimitTracker.java
@@ -20,6 +20,7 @@
 import com.android.dex.TypeList;
 import com.google.common.collect.Interner;
 import com.google.common.collect.Interners;
+import com.google.common.collect.Sets;
 import java.util.Collections;
 import java.util.HashSet;
 
@@ -30,15 +31,34 @@
 class DexLimitTracker {
 
   private static final ThreadLocal<Interner<String>> threadLocalInterner =
-      ThreadLocal.withInitial(Interners::newWeakInterner);
+      ThreadLocal.withInitial(Interners::newStrongInterner);
 
-  private final HashSet<String> fieldsSeen = new HashSet<>();
-  private final HashSet<String> methodsSeen = new HashSet<>();
-  private final HashSet<String> typesSeen = new HashSet<>();
+  /**
+   * Upper bound for initial HashSet capacity. Avoids repeated resizing and rehashing cycles (from
+   * the default capacity of 16) for typical multidex shards without prematurely allocating full
+   * 64K-entry backing tables for smaller archives.
+   */
+  private static final int MAX_INITIAL_EXPECTED_SIZE = 8192;
+
+  /**
+   * Removes the thread-local interner instance for the current thread to prevent accumulation
+   * across multiple builds in persistent worker processes.
+   */
+  public static void clearInterner() {
+    threadLocalInterner.remove();
+  }
+
+  private final HashSet<String> fieldsSeen;
+  private final HashSet<String> methodsSeen;
+  private final HashSet<String> typesSeen;
   private final int maxNumberOfIdxPerDex;
 
   public DexLimitTracker(int maxNumberOfIdxPerDex) {
     this.maxNumberOfIdxPerDex = maxNumberOfIdxPerDex;
+    int initialCapacity = Math.min(maxNumberOfIdxPerDex, MAX_INITIAL_EXPECTED_SIZE);
+    this.fieldsSeen = Sets.newHashSetWithExpectedSize(initialCapacity);
+    this.methodsSeen = Sets.newHashSetWithExpectedSize(initialCapacity);
+    this.typesSeen = Sets.newHashSetWithExpectedSize(initialCapacity);
   }
 
   /**
@@ -57,6 +77,7 @@
     fieldsSeen.clear();
     methodsSeen.clear();
     typesSeen.clear();
+    threadLocalInterner.remove();
   }
 
   public void track(Dex dexFile) {
@@ -81,44 +102,55 @@
     }
 
     static DexTrackerInfo create(Dex dexFile) {
+      int stringCount = dexFile.getTableOfContents().stringIds.size;
+      String[] stringCache = new String[stringCount];
+
       int typeCount = dexFile.typeIds().size();
       String[] types = new String[typeCount];
       for (int typeIndex = 0; typeIndex < typeCount; ++typeIndex) {
-        types[typeIndex] = typeName(dexFile, typeIndex);
+        int stringIndex = dexFile.typeIds().get(typeIndex);
+        types[typeIndex] = getString(dexFile, stringIndex, stringCache);
       }
 
       int fieldCount = dexFile.fieldIds().size();
       String[] fields = new String[fieldCount];
       for (int fieldIndex = 0; fieldIndex < fieldCount; ++fieldIndex) {
-        fields[fieldIndex] = fieldSignature(dexFile, fieldIndex, types);
+        fields[fieldIndex] = fieldSignature(dexFile, fieldIndex, types, stringCache);
       }
 
       int methodCount = dexFile.methodIds().size();
       String[] methods = new String[methodCount];
       for (int methodIndex = 0; methodIndex < methodCount; ++methodIndex) {
-        methods[methodIndex] = methodSignature(dexFile, methodIndex, types);
+        methods[methodIndex] = methodSignature(dexFile, methodIndex, types, stringCache);
       }
 
       return new DexTrackerInfo(fields, methods, types);
     }
   }
 
-  private static String typeName(Dex dex, int typeIndex) {
-    return threadLocalInterner.get().intern(dex.typeNames().get(typeIndex));
+  private static String getString(Dex dex, int stringIndex, String[] stringCache) {
+    String s = stringCache[stringIndex];
+    if (s == null) {
+      s = threadLocalInterner.get().intern(dex.strings().get(stringIndex));
+      stringCache[stringIndex] = s;
+    }
+    return s;
   }
 
-  private static String fieldSignature(Dex dex, int fieldIndex, String[] typeCache) {
+  private static String fieldSignature(
+      Dex dex, int fieldIndex, String[] typeCache, String[] stringCache) {
     FieldId field = dex.fieldIds().get(fieldIndex);
-    String name = dex.strings().get(field.getNameIndex());
+    String name = getString(dex, field.getNameIndex(), stringCache);
     String declaringClass = typeCache[field.getDeclaringClassIndex()];
     String type = typeCache[field.getTypeIndex()];
     return threadLocalInterner.get().intern(declaringClass + "." + name + ":" + type);
   }
 
-  private static String methodSignature(Dex dex, int methodIndex, String[] typeCache) {
+  private static String methodSignature(
+      Dex dex, int methodIndex, String[] typeCache, String[] stringCache) {
     MethodId method = dex.methodIds().get(methodIndex);
     ProtoId proto = dex.protoIds().get(method.getProtoIndex());
-    String name = dex.strings().get(method.getNameIndex());
+    String name = getString(dex, method.getNameIndex(), stringCache);
     String declaringClass = typeCache[method.getDeclaringClassIndex()];
     String returnType = typeCache[proto.getReturnTypeIndex()];
     TypeList parameterTypeIndices = dex.readTypeList(proto.getParametersOffset());
@@ -128,6 +160,6 @@
     }
     return threadLocalInterner
         .get()
-        .intern(declaringClass + "." + name + ":" + returnType + "(" + parameterTypes + ")");
+        .intern(declaringClass + "." + name + "(" + parameterTypes + ")" + returnType);
   }
 }

src/tools/javatests/com/google/devtools/build/android/dexer/DexLimitTrackerTest.java

diff --git a/src/tools/javatests/com/google/devtools/build/android/dexer/DexLimitTrackerTest.java b/src/tools/javatests/com/google/devtools/build/android/dexer/DexLimitTrackerTest.java
index 1788198..744c43f 100644
--- a/src/tools/javatests/com/google/devtools/build/android/dexer/DexLimitTrackerTest.java
+++ b/src/tools/javatests/com/google/devtools/build/android/dexer/DexLimitTrackerTest.java
@@ -122,10 +122,10 @@
     assertThat(info.methods)
         .asList()
         .containsExactly(
-            "Lbase/SimpleClass;.<init>:V()",
-            "Lbase/SimpleClass;.method:V()",
-            "Lbase/SimpleClass;.methodWithParams:V(ILjava/lang/String;)",
-            "Ljava/lang/Object;.<init>:V()");
+            "Lbase/SimpleClass;.<init>()V",
+            "Lbase/SimpleClass;.method()V",
+            "Lbase/SimpleClass;.methodWithParams(ILjava/lang/String;)V",
+            "Ljava/lang/Object;.<init>()V");
 
     // Verify types
     assertThat(info.types).asList().contains("Lbase/SimpleClass;");
@@ -206,4 +206,50 @@
       assertTrue("Field " + field.getName() + " must be final", Modifier.isFinal(modifiers));
     }
   }
+
+  @Test
+  public void testTrackMethodsAndTypesLimits() throws Exception {
+    Dex dexSimple = loadDexFromArchive(simpleDexArchive); // 4 methods, 1 field
+
+    // Set method limit to 3 (less than 4) -> should be outside limits
+    DexLimitTracker methodTracker = new DexLimitTracker(3);
+    methodTracker.track(dexSimple);
+    assertThat(methodTracker.outsideLimits()).isTrue();
+
+    // Set type limit to 1 -> should be outside limits
+    DexLimitTracker typeTracker = new DexLimitTracker(1);
+    typeTracker.track(dexSimple);
+    assertThat(typeTracker.outsideLimits()).isTrue();
+  }
+
+  @Test
+  public void testTrackerReuseAfterClear() throws Exception {
+    DexLimitTracker tracker = new DexLimitTracker(8);
+    Dex dexSimple = loadDexFromArchive(simpleDexArchive);
+    Dex dexFieldsTypes = loadDexFromArchive(fieldsTypesDexArchive);
+
+    tracker.track(dexFieldsTypes);
+    assertThat(tracker.outsideLimits()).isTrue();
+
+    // Clear and reuse tracker instance
+    tracker.clear();
+    assertThat(tracker.outsideLimits()).isFalse();
+
+    tracker.track(dexSimple);
+    assertThat(tracker.outsideLimits()).isFalse();
+  }
+
+  @Test
+  public void testClearInterner() throws Exception {
+    Dex dex1 = loadDexFromArchive(simpleDexArchive);
+    DexTrackerInfo info1 = DexTrackerInfo.create(dex1);
+    assertThat(info1.types).isNotEmpty();
+
+    // Calling clearInterner() clears the thread-local interner instance
+    DexLimitTracker.clearInterner();
+
+    Dex dex2 = loadDexFromArchive(simpleDexArchive);
+    DexTrackerInfo info2 = DexTrackerInfo.create(dex2);
+    assertThat(info2.types).isNotEmpty();
+  }
 }