Teach DexMapper to not separate synthetic classes from their context …

…classes when sharding dexes, like DexFileSplitter in 9092303. Fixes #16368

RELNOTES: None
PiperOrigin-RevId: 546041575
Change-Id: I12fa8cff15b13534ca9a5682b7e9b43e6c860ef8

src/test/java/com/google/devtools/build/android/ziputils/SplitterTest.java

@@ -51,7 +51,7 @@ public void testAssign() {
       filter.add("dir" + i + ARCHIVE_FILE_SEPARATOR + "file" + i + CLASS_SUFFIX);
     }
     Splitter splitter = new Splitter(size + 1, input.size());
-    splitter.assign(filter);
+    splitter.assignAllToCurrentShard(filter);
     splitter.nextShard();
     output = new LinkedHashMap<>();
     for (String path : input) {

src/test/shell/bazel/android/BUILD

@@ -189,11 +189,14 @@ android_sh_test(
     name = "DexFileSplitter_synthetic_classes_test",
     size = "large",
     srcs = ["DexFileSplitter_synthetic_classes_test.sh"],
+    # Fixes to DexMapper are not released yet.
+    create_test_with_released_tools = False,
     data = [
         ":android_helper",
         "//external:android_sdk_for_testing",
         "//src/test/shell/bazel:test-deps",
     ],
+    shard_count = 2,
     tags = [
         "no_windows",
     ],

src/test/shell/bazel/android/DexFileSplitter_synthetic_classes_test.sh

@@ -26,18 +26,20 @@
 # Load the test setup defined in the parent directory
 CURRENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
+source "${CURRENT_DIR}/../../integration_test_setup.sh" \
+  || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
+
 source "${CURRENT_DIR}/android_helper.sh" \
   || { echo "android_helper.sh not found!" >&2; exit 1; }
 fail_if_no_android_sdk
 
-source "${CURRENT_DIR}/../../integration_test_setup.sh" \
-  || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
-
 resolve_android_toolchains "$1"
 
-function test_DexFileSplitter_synthetic_classes_crossing_dexfiles() {
-  create_new_workspace
-  setup_android_sdk_support
+function create_test_app() {
+
+  inner_class_count="$1"
+  lambda_count="$2"
+  dex_shard_count="$3"
 
   mkdir -p java/com/testapp
 
@@ -81,7 +83,7 @@ public class MainActivity extends Activity {
 }
 EOF
 
-  generate_java_file_with_many_synthetic_classes > java/com/testapp/BigLib.java
+  generate_java_file_with_many_synthetic_classes "$1" "$2" > java/com/testapp/BigLib.java
 
   cat > java/com/testapp/BUILD <<EOF
 android_binary(
@@ -90,26 +92,22 @@ android_binary(
         "MainActivity.java",
         ":BigLib.java",
     ],
+    dex_shards = $dex_shard_count,
     manifest = "AndroidManifest.xml",
 )
 EOF
-
-  bazel build java/com/testapp || fail "Test app should have built succesfully"
-
-  dex_file_count=$(unzip -l bazel-bin/java/com/testapp/testapp.apk | grep "classes[0-9]*.dex" | wc -l)
-  if [[ ! "$dex_file_count" -ge "2" ]]; then
-    echo "Expected at least 2 dexes in app, found: $dex_file_count"
-    exit 1
-  fi
 }
 
-
 function generate_java_file_with_many_synthetic_classes() {
+
+  inner_class_count="$1"
+  lambda_count="$2"
+
   echo "package com.testapp;"
   echo "public class BigLib {"
 
   # First generate enough inner classes to fill up most of the dex
-  for (( i = 0; i < 21400; i++ )) do
+  for (( i = 0; i < $inner_class_count; i++ )) do
 
     echo "  public static class Bar$i {"
     echo "    public int bar() {"
@@ -129,7 +127,7 @@ function generate_java_file_with_many_synthetic_classes() {
   echo "    public IntSupplier[] foo() {"
   echo "      return new IntSupplier[] {"
 
-  for ((i = 0; i < 6000; i++ )) do
+  for ((i = 0; i < $lambda_count; i++ )) do
     echo "        () -> $i,"
   done
 
@@ -139,4 +137,42 @@ function generate_java_file_with_many_synthetic_classes() {
   echo "}"
 }
 
+function test_DexFileSplitter_synthetic_classes_crossing_dexfiles() {
+  create_new_workspace
+  setup_android_sdk_support
+
+  # dex_shards default is 1
+  create_test_app 21400 6000 1
+
+  bazel build java/com/testapp || fail "Test app should have built succesfully"
+
+  dex_file_count=$(unzip -l bazel-bin/java/com/testapp/testapp.apk | grep "classes[0-9]*.dex" | wc -l)
+  if [[ ! "$dex_file_count" -ge "2" ]]; then
+    echo "Expected at least 2 dexes in app, found: $dex_file_count"
+    exit 1
+  fi
+}
+
+function test_DexMapper_synthetic_classes_crossing_dexfiles() {
+  create_new_workspace
+  setup_android_sdk_support
+
+  # 3 inner classes, 6 lambdas (and thus 6 synthetics from D8) and 5 dex_shards
+  # is one magic combination to repro synthetics being separated from their
+  # context / enclosing classes.
+  create_test_app 3 6 5
+
+  echo here2
+  echo $TEST_TMPDIR/bazelrc
+  cat $TEST_TMPDIR/bazelrc
+
+  bazel build java/com/testapp || fail "Test app should have built succesfully"
+
+  dex_file_count=$(unzip -l bazel-bin/java/com/testapp/testapp.apk | grep "classes[0-9]*.dex" | wc -l)
+  if [[ ! "$dex_file_count" -eq "5" ]]; then
+    echo "Expected 5 dexes in app, found: $dex_file_count"
+    exit 1
+  fi
+}
+
 run_suite "Tests for DexFileSplitter with synthetic classes crossing dexfiles"

src/test/shell/bazel/android/android_sh_test.bzl

@@ -29,20 +29,36 @@ CHECK_FOR_ANDROID_SDK = select(
     no_match_error = "This test requires an android SDK, and one isn't present. Make sure to uncomment the android rules in the WORKSPACE.",
 )
 
-def android_sh_test(**kwargs):
+def android_sh_test(create_test_with_released_tools = True, **kwargs):
+    """Creates versions of the test with and without platforms and head android tools.
+
+    Args:
+        create_test_with_released_tools: Whether to create a version of the test with the released
+        android tools, for when the code under test relies on not-yet-released code.
+        **kwargs: Args to sh_test
+    """
     name = kwargs.pop("name")
     data = kwargs.pop("data")
     if not data:
         data = []
     data = data + CHECK_FOR_ANDROID_SDK
 
-    # Test with released android_tools version.
-    native.sh_test(
-        name = name,
-        args = ["--without_platforms"],
-        data = data,
-        **kwargs
-    )
+    if create_test_with_released_tools:
+        # Test with released android_tools version.
+        native.sh_test(
+            name = name,
+            args = ["--without_platforms"],
+            data = data,
+            **kwargs
+        )
+
+        # Test with platform-based toolchain resolution.
+        native.sh_test(
+            name = name + "_with_platforms",
+            data = data,
+            args = ["--with_platforms"],
+            **kwargs
+        )
 
     # Test with android_tools version that's built at the same revision
     # as the test itself.
@@ -54,11 +70,3 @@ def android_sh_test(**kwargs):
         ],
         **kwargs
     )
-
-    # Test with platform-based toolchain resolution.
-    native.sh_test(
-        name = name + "_with_platforms",
-        data = data,
-        args = ["--with_platforms"],
-        **kwargs
-    )

src/tools/android/java/com/google/devtools/build/android/ziputils/SplitZip.java

@@ -27,9 +27,13 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
@@ -39,10 +43,12 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Scanner;
 import java.util.Set;
 import java.util.TreeSet;
 
@@ -416,12 +422,22 @@ private void checkConfig() throws IOException {
     filter = filterFile == null && filterInputStream == null ? null : readPaths(filterFile);
   }
 
-  /**
-   * Parses the entries and assign each entry to an output file.
-   */
-  private void split() {
+  /** Parses the entries and assign each entry to an output file. */
+  private void split() throws IOException {
+
+    // A map of class (a "context") to its inner synthetic classes from D8.
+    SetMultimap<String, String> syntheticClassContexts = HashMultimap.create();
+
     for (ZipIn in : inputs) {
       CentralDirectory cdir = centralDirectories.get(in.getFilename());
+
+      for (DirectoryEntry entry : cdir.list()) {
+        if (entry.getFilename().equals("META-INF/synthetic-contexts.map")) {
+          parseSyntheticContextsMap(in.entryFor(entry).getContent(), syntheticClassContexts);
+          break;
+        }
+      }
+
       for (DirectoryEntry entry : cdir.list()) {
         String filename = normalizedFilename(entry.getFilename());
         if (!inputFilter.apply(filename)) {
@@ -438,17 +454,31 @@ private void split() {
         }
       }
     }
+
     Splitter splitter = new Splitter(outputs.size(), classes.size());
+
     if (filter != null) {
       // Assign files in the filter to the first output file.
-      splitter.assign(Sets.filter(filter, inputFilter));
+      splitter.assignAllToCurrentShard(Sets.filter(filter, inputFilter));
       splitter.nextShard(); // minimal initial shard
     }
+
+    Set<String> allSyntheticClasses = new HashSet<>(syntheticClassContexts.values());
+
     for (String path : classes) {
-      // Use normalized filename so the filter file doesn't have to change
-      int assignment = splitter.assign(path);
-      Preconditions.checkState(assignment >= 0 && assignment < zipOuts.length);
-      assignments.put(path, zipOuts[assignment]);
+      if (!allSyntheticClasses.contains(path)) {
+
+        // Use normalized filename so the filter file doesn't have to change
+        int assignment = splitter.assign(path);
+        Set<String> syntheticClasses = syntheticClassContexts.get(path);
+        splitter.assignAllToCurrentShard(syntheticClasses);
+        Preconditions.checkState(assignment >= 0 && assignment < zipOuts.length);
+
+        assignments.put(path, zipOuts[assignment]);
+        for (String syntheticClass : syntheticClasses) {
+          assignments.put(syntheticClass, zipOuts[assignment]);
+        }
+      }
     }
   }
 
@@ -494,6 +524,23 @@ private String fixPath(String path) {
     return path;
   }
 
+  private static void parseSyntheticContextsMap(
+      ByteBuffer byteBuffer, Multimap<String, String> syntheticClassContexts) {
+    // The ByteBuffer returned from the Splitter's zip library is not backed by an accessible array,
+    // so ByteBuffer.array() is not supported, so we must go the long way.
+    byte[] bytes = new byte[byteBuffer.remaining()];
+    byteBuffer.get(bytes);
+    Scanner scanner = new Scanner(new ByteArrayInputStream(bytes), UTF_8);
+    scanner.useDelimiter("[;\n]");
+    while (scanner.hasNext()) {
+      String syntheticClass = scanner.next();
+      String context = scanner.next();
+      // The context map uses class names, whereas SplitZip uses class file names, so add ".class"
+      // here to make it easier to work with the map in the rest of the code.
+      syntheticClassContexts.put(context + ".class", syntheticClass + ".class");
+    }
+  }
+
   private void verbose(String msg) {
     if (verbose) {
       System.out.println(msg);

src/tools/android/java/com/google/devtools/build/android/ziputils/Splitter.java

@@ -24,10 +24,10 @@ class Splitter {
 
   private final int numberOfShards;
   private final Map<String, Integer> assigned;
-  private int size = 0;
-  private int shard = 0;
+  private int size = 0; // Number of classes in the current shard
+  private int shard = 0; // Current shard number
   private String prevPath = null;
-  private int remaining;
+  private int remaining; // Number of classes remaining to be assigned shards
   private int idealSize;
   private int almostFull;
 
@@ -51,15 +51,15 @@ public Splitter(int numberOfShards, int expectedSize) {
    * remaining entries to process is adjusted, by subtracting the number of as-of-yet unassigned
    * entries from the filter.
    */
-  public void assign(Collection<String> filter) {
-    if (filter != null) {
-      for (String s : filter) {
-        if (!assigned.containsKey(s)) {
+  public void assignAllToCurrentShard(Collection<String> entries) {
+    if (entries != null) {
+      for (String e : entries) {
+        if (!assigned.containsKey(e)) {
           remaining--;
         }
-        assigned.put(s, shard);
+        assigned.put(e, shard);
       }
-      size = filter.size();
+      size += entries.size();
     }
   }