Optimize DexFileSplitter by improving memory locality. PiperOrigin-RevId: 935770875 Change-Id: Ie519b7ecd7a9006483f1006d3ff94e7be78d2f8e
diff --git a/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java b/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java index 4a37deb..91bdd7b 100644 --- a/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java +++ b/src/tools/java/com/google/devtools/build/android/dexer/DexFileSplitter.java
@@ -44,12 +44,12 @@ import java.nio.file.Path; import java.nio.file.StandardOpenOption; import java.util.ArrayList; -import java.util.Collection; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Scanner; import java.util.Set; +import java.util.SortedSet; import java.util.TreeMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -229,10 +229,6 @@ StandardOpenOption.WRITE)); } - private int shardsWritten() { - return curShard; - } - @Override public void close() throws IOException { if (curOut != null) { @@ -249,7 +245,7 @@ Map<String, ZipFile> dexFilesAndContainingZip, TreeMultimap<String, String> contextClassesToSyntheticClasses, Predicate<String> filter) - throws ExecutionException, InterruptedException, IOException { + throws ExecutionException, InterruptedException { Set<String> syntheticClasses = new HashSet<>(contextClassesToSyntheticClasses.values()); int nThreads = Runtime.getRuntime().availableProcessors(); @@ -268,7 +264,7 @@ if (filter.test(filename)) { if (!syntheticClasses.contains(filename)) { ZipFile zipFile = entry.getValue(); - Collection<String> synths = contextClassesToSyntheticClasses.get(filename); + SortedSet<String> synths = contextClassesToSyntheticClasses.get(filename); ListenableFuture<ZipEntryAndContent> contextFuture = listeningExecutor.submit(() -> readAndParseDex(zipFile, filename)); @@ -391,15 +387,13 @@ entry.getSize()); Dex dexFile = new Dex(content); - DexLimitTracker.DexTrackerInfo trackerInfo = new DexLimitTracker.DexTrackerInfo(dexFile); + DexLimitTracker.DexTrackerInfo trackerInfo = DexLimitTracker.DexTrackerInfo.create(dexFile); return new ZipEntryAndContent(entry, content, trackerInfo); } catch (IOException e) { throw new UncheckedIOException(e); } } - - private static final class ZipEntryAndContent { final ZipEntry zipEntry; final byte[] content;
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 14d74c5..a836b07 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
@@ -18,9 +18,9 @@ import com.android.dex.MethodId; import com.android.dex.ProtoId; import com.android.dex.TypeList; -import com.google.common.collect.ImmutableList; import com.google.common.collect.Interner; import com.google.common.collect.Interners; +import java.util.Collections; import java.util.HashSet; /** @@ -60,43 +60,46 @@ } public void track(Dex dexFile) { - track(new DexTrackerInfo(dexFile)); + track(DexTrackerInfo.create(dexFile)); } public void track(DexTrackerInfo info) { - fieldsSeen.addAll(info.fields); - methodsSeen.addAll(info.methods); - typesSeen.addAll(info.types); + Collections.addAll(fieldsSeen, info.fields); + Collections.addAll(methodsSeen, info.methods); + Collections.addAll(typesSeen, info.types); } - static class DexTrackerInfo { - final ImmutableList<String> fields; - final ImmutableList<String> methods; - final ImmutableList<String> types; + static final class DexTrackerInfo { + final String[] fields; + final String[] methods; + final String[] types; - DexTrackerInfo(Dex dexFile) { - int fieldCount = dexFile.fieldIds().size(); - ImmutableList.Builder<String> fieldsBuilder = - ImmutableList.builderWithExpectedSize(fieldCount); - for (int fieldIndex = 0; fieldIndex < fieldCount; ++fieldIndex) { - fieldsBuilder.add(fieldSignature(dexFile, fieldIndex)); + DexTrackerInfo(String[] fields, String[] methods, String[] types) { + this.fields = fields; + this.methods = methods; + this.types = types; + } + + static DexTrackerInfo create(Dex dexFile) { + int typeCount = dexFile.typeIds().size(); + String[] types = new String[typeCount]; + for (int typeIndex = 0; typeIndex < typeCount; ++typeIndex) { + types[typeIndex] = typeName(dexFile, typeIndex); } - fields = fieldsBuilder.build(); + + int fieldCount = dexFile.fieldIds().size(); + String[] fields = new String[fieldCount]; + for (int fieldIndex = 0; fieldIndex < fieldCount; ++fieldIndex) { + fields[fieldIndex] = fieldSignature(dexFile, fieldIndex, types); + } int methodCount = dexFile.methodIds().size(); - ImmutableList.Builder<String> methodsBuilder = - ImmutableList.builderWithExpectedSize(methodCount); + String[] methods = new String[methodCount]; for (int methodIndex = 0; methodIndex < methodCount; ++methodIndex) { - methodsBuilder.add(methodSignature(dexFile, methodIndex)); + methods[methodIndex] = methodSignature(dexFile, methodIndex, types); } - methods = methodsBuilder.build(); - int typeCount = dexFile.typeIds().size(); - ImmutableList.Builder<String> typesBuilder = ImmutableList.builderWithExpectedSize(typeCount); - for (int typeIndex = 0; typeIndex < typeCount; ++typeIndex) { - typesBuilder.add(typeName(dexFile, typeIndex)); - } - types = typesBuilder.build(); + return new DexTrackerInfo(fields, methods, types); } } @@ -104,24 +107,24 @@ return threadLocalInterner.get().intern(dex.typeNames().get(typeIndex)); } - private static String fieldSignature(Dex dex, int fieldIndex) { + private static String fieldSignature(Dex dex, int fieldIndex, String[] typeCache) { FieldId field = dex.fieldIds().get(fieldIndex); String name = dex.strings().get(field.getNameIndex()); - String declaringClass = typeName(dex, field.getDeclaringClassIndex()); - String type = typeName(dex, field.getTypeIndex()); + 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) { + private static String methodSignature(Dex dex, int methodIndex, String[] typeCache) { MethodId method = dex.methodIds().get(methodIndex); ProtoId proto = dex.protoIds().get(method.getProtoIndex()); String name = dex.strings().get(method.getNameIndex()); - String declaringClass = typeName(dex, method.getDeclaringClassIndex()); - String returnType = typeName(dex, proto.getReturnTypeIndex()); + String declaringClass = typeCache[method.getDeclaringClassIndex()]; + String returnType = typeCache[proto.getReturnTypeIndex()]; TypeList parameterTypeIndices = dex.readTypeList(proto.getParametersOffset()); StringBuilder parameterTypes = new StringBuilder(); for (short parameterTypeIndex : parameterTypeIndices.getTypes()) { - parameterTypes.append(typeName(dex, parameterTypeIndex & 0xFFFF)); + parameterTypes.append(typeCache[parameterTypeIndex & 0xFFFF]); } return threadLocalInterner .get()
diff --git a/src/tools/javatests/com/google/devtools/build/android/dexer/AllTests.java b/src/tools/javatests/com/google/devtools/build/android/dexer/AllTests.java index 4f936b6..15e2237 100644 --- a/src/tools/javatests/com/google/devtools/build/android/dexer/AllTests.java +++ b/src/tools/javatests/com/google/devtools/build/android/dexer/AllTests.java
@@ -20,5 +20,6 @@ @RunWith(Suite.class) @Suite.SuiteClasses({ DexFileSplitterTest.class, + DexLimitTrackerTest.class, }) public class AllTests {}
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 new file mode 100644 index 0000000..7e2faf6 --- /dev/null +++ b/src/tools/javatests/com/google/devtools/build/android/dexer/DexLimitTrackerTest.java
@@ -0,0 +1,209 @@ +// Copyright 2026 The Bazel Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +package com.google.devtools.build.android.dexer; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +import com.android.dex.Dex; +import com.google.devtools.build.android.dexer.DexLimitTracker.DexTrackerInfo; +import com.google.devtools.build.android.r8.CompatDexBuilder; +import com.google.devtools.build.runfiles.Runfiles; +import java.io.IOException; +import java.io.InputStream; +import java.lang.reflect.Field; +import java.lang.reflect.Modifier; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.zip.ZipEntry; +import java.util.zip.ZipFile; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link DexLimitTracker}. */ +@RunWith(JUnit4.class) +public class DexLimitTrackerTest { + + private static final Path SIMPLE_JAR; + private static final Path FIELDS_TYPES_JAR; + + @Rule public TemporaryFolder tmp = new TemporaryFolder(); + + static { + try { + Runfiles runfiles = Runfiles.preload().unmapped(); + SIMPLE_JAR = getRunfile(runfiles, "simplejar"); + FIELDS_TYPES_JAR = getRunfile(runfiles, "fields_types_jar"); + } catch (Exception e) { + throw new ExceptionInInitializerError(e); + } + } + + private static Path getRunfile(Runfiles runfiles, String property) { + String path = System.getProperty(property); + String google3Path = runfiles.rlocation("google3/" + path); + if (google3Path != null && Files.exists(Path.of(google3Path))) { + return Path.of(google3Path); + } + String rulesAndroidPath = runfiles.rlocation("rules_android/" + path); + if (rulesAndroidPath != null && Files.exists(Path.of(rulesAndroidPath))) { + return Path.of(rulesAndroidPath); + } + throw new RuntimeException("Could not find runfile for property " + property + ": " + path); + } + + private Path simpleDexArchive; + private Path fieldsTypesDexArchive; + + @Before + public void setUp() throws Exception { + simpleDexArchive = buildDexArchive(SIMPLE_JAR, "simple.dex.zip"); + fieldsTypesDexArchive = buildDexArchive(FIELDS_TYPES_JAR, "fields_types.dex.zip"); + } + + private Path buildDexArchive(Path inputJar, String outputZip) throws Exception { + Path outputZipPath = tmp.getRoot().toPath().resolve(outputZip); + CompatDexBuilder.main( + new String[] { + "--input_jar", + inputJar.toString(), + "--output_zip", + outputZipPath.toString(), + "--num-threads=1", + "--positions=lines" + }); + return outputZipPath; + } + + private Dex loadDexFromArchive(Path dexArchive) throws IOException { + try (ZipFile zip = new ZipFile(dexArchive.toFile())) { + ZipEntry entry = + zip.stream() + .filter(e -> e.getName().endsWith(".class.dex")) + .findFirst() + .orElseThrow(() -> new IOException("No .dex entry found")); + try (InputStream is = zip.getInputStream(entry)) { + return new Dex(is.readAllBytes()); + } + } + } + + @Test + public void testTrackerInfoCreationAndTracking() throws Exception { + Dex dex = loadDexFromArchive(simpleDexArchive); + DexTrackerInfo info = DexTrackerInfo.create(dex); + + // Verify fields + assertThat(info.fields).hasLength(1); + assertThat(info.fields[0]).isEqualTo("Lbase/SimpleClass;.field:I"); + + // Verify methods + assertThat(info.methods) + .hasLength( + 4); // constructor, method(), methodWithParams(int, String), and superclass constructor + // Object.<init> + 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()"); + + // Verify types + assertThat(info.types).asList().contains("Lbase/SimpleClass;"); + } + + @Test + public void testMemoizationCache() throws Exception { + Dex dex1 = loadDexFromArchive(simpleDexArchive); + DexTrackerInfo info1 = DexTrackerInfo.create(dex1); + + Dex dex2 = loadDexFromArchive(simpleDexArchive); + DexTrackerInfo info2 = DexTrackerInfo.create(dex2); + + // Verify that equal type names across tracker runs refer to the same String instance (interned) + String type1 = null; + for (String t : info1.types) { + if (t.equals("Lbase/SimpleClass;")) { + type1 = t; + } + } + String type2 = null; + for (String t : info2.types) { + if (t.equals("Lbase/SimpleClass;")) { + type2 = t; + } + } + + assertThat(type1).isNotNull(); + assertThat(type2).isNotNull(); + assertSame("Type names should be referentially equal (interned)", type1, type2); + + // Verify that equal field signatures across tracker runs refer to the same String instance + assertSame( + "Field signatures should be referentially equal (interned)", + info1.fields[0], + info2.fields[0]); + + // Verify that equal method signatures across tracker runs refer to the same String instance + String[] methods1 = info1.methods.clone(); + String[] methods2 = info2.methods.clone(); + Arrays.sort(methods1); + Arrays.sort(methods2); + for (int i = 0; i < methods1.length; i++) { + assertSame( + "Method signatures should be referentially equal (interned)", methods1[i], methods2[i]); + } + } + + @Test + public void testTrackCountsAndLimits() throws Exception { + DexLimitTracker tracker = new DexLimitTracker(8); // low limit of 8 + Dex dexSimple = loadDexFromArchive(simpleDexArchive); // 1 field, 4 methods, some types + Dex dexFieldsTypes = + loadDexFromArchive(fieldsTypesDexArchive); // 10 fields, 2 methods, some types + // Track simple dex first + tracker.track(dexSimple); + assertThat(tracker.outsideLimits()).isFalse(); + + // Track fieldsTypes (10 fields) which exceeds limit of 8 fields/types + tracker.track(dexFieldsTypes); + assertThat(tracker.outsideLimits()).isTrue(); + + // Verify clear resets tracking + tracker.clear(); + assertThat(tracker.outsideLimits()).isFalse(); + } + + @Test + public void testTrackerInfoImmutability() throws Exception { + // Check that all fields of DexTrackerInfo are final + Class<DexTrackerInfo> clazz = DexTrackerInfo.class; + for (Field field : clazz.getDeclaredFields()) { + // Ignore synthetic fields (like $jacocoData if code coverage is enabled) + if (field.isSynthetic()) { + continue; + } + int modifiers = field.getModifiers(); + assertTrue("Field " + field.getName() + " must be final", Modifier.isFinal(modifiers)); + } + } +}