Collect the contexts of D8 compiler-synthesized classes.

This completes steps 2 and 3 in b/241351268#comment9
towards resolving #16368

PiperOrigin-RevId: 530238344
Change-Id: I569bf8e3bfa81f7005f3e9b79338d5a5b868e339

Author: Googler

src/test/java/com/google/devtools/build/android/r8/BUILD

@@ -51,13 +51,15 @@ java_test(
         ":arithmetic",
         ":barray",
         ":naming001",
+        ":testdata_lambda_desugared.jar",
         ":twosimpleclasses",
     ],
     jvm_flags = [
         "-DCompatDexBuilderTests.twosimpleclasses=$(location :twosimpleclasses)",
         "-DCompatDexBuilderTests.naming001=$(location :naming001)",
         "-DCompatDxTests.arithmetic=$(location :arithmetic)",
         "-DCompatDxTests.barray=$(location :barray)",
+        "-DCompatDexBuilderTests.lambda=$(location :testdata_lambda_desugared.jar)",
     ],
     runtime_deps = [
         ":tests",
@@ -103,3 +105,21 @@ java_library(
         "-target 8",
     ],
 )
+
+java_library(
+    name = "testdata_lambda",
+    srcs = glob(["testdata/lambda/*.java"]),
+)
+
+genrule(
+    name = "desugar_testdata_lambda",
+    srcs = [
+        ":testdata_lambda",
+        "@bazel_tools//tools/android:android_jar",
+    ],
+    outs = ["testdata_lambda_desugared.jar"],
+    cmd = "$(location //src/tools/android/java/com/google/devtools/build/android/r8:desugar) " +
+          "-i $(location :testdata_lambda) -o $@ " +
+          "--bootclasspath_entry $(location @bazel_tools//tools/android:android_jar)",
+    tools = ["//src/tools/android/java/com/google/devtools/build/android/r8:desugar"],
+)

src/test/java/com/google/devtools/build/android/r8/CompatDexBuilderTest.java

@@ -21,13 +21,15 @@
 import com.android.tools.r8.OutputMode;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.common.options.OptionsParsingException;
+import java.io.BufferedReader;
 import java.io.IOException;
+import java.io.InputStreamReader;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 import org.junit.Rule;
 import org.junit.Test;
@@ -45,7 +47,7 @@ public void compileManyClasses()
       throws IOException, InterruptedException, ExecutionException, OptionsParsingException {
     // Random set of classes from the R8 example test directory naming001.
     final String inputJar = System.getProperty("CompatDexBuilderTests.naming001");
-    final List<String> classNames =
+    final ImmutableList<String> classNames =
         ImmutableList.of(
             "A",
             "B",
@@ -86,6 +88,40 @@ public void compileManyClasses()
     assertThat(expectedNames).isEmpty();
   }
 
+  @Test
+  public void compileWithSyntheticLambdas() throws Exception {
+    final String contextName = "com/google/devtools/build/android/r8/testdata/lambda/Lambda";
+    final String inputJar = System.getProperty("CompatDexBuilderTests.lambda");
+    final Path outputZip = temp.getRoot().toPath().resolve("out.zip");
+    CompatDexBuilder.main(
+        new String[] {"--input_jar", inputJar, "--output_zip", outputZip.toString()});
+    assertThat(Files.exists(outputZip)).isTrue();
+
+    try (ZipFile zipFile = new ZipFile(outputZip.toFile(), UTF_8)) {
+      assertThat(zipFile.getEntry(contextName + ".class.dex")).isNotNull();
+      ZipEntry entry = zipFile.getEntry("META-INF/synthetic-contexts.map");
+      assertThat(entry).isNotNull();
+      try (BufferedReader reader =
+          new BufferedReader(new InputStreamReader(zipFile.getInputStream(entry), UTF_8))) {
+        String line = reader.readLine();
+        assertThat(line).isNotNull();
+        // Format of mapping is: <synthetic-binary-name>;<context-binary-name>\n
+        int sep = line.indexOf(';');
+        String syntheticNameInMap = line.substring(0, sep);
+        String contextNameInMap = line.substring(sep + 1);
+        // The synthetic will be prefixed by the context type. This checks the synthetic name
+        // is larger than the context to avoid hardcoding the synthetic names, which may change.
+        assertThat(syntheticNameInMap).startsWith(contextName);
+        assertThat(syntheticNameInMap).isNotEqualTo(contextName);
+        // Check expected context.
+        assertThat(contextNameInMap).isEqualTo(contextName);
+        // Only one synthetic and its context should be present.
+        line = reader.readLine();
+        assertThat(line).isNull();
+      }
+    }
+  }
+
   @Test
   public void compileTwoClassesAndRun() throws Exception {
     // Run CompatDexBuilder on dexMergeSample.jar

src/test/java/com/google/devtools/build/android/r8/testdata/lambda/Lambda.java

@@ -0,0 +1,30 @@
+// Copyright 2023 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.r8.testdata.lambda;
+
+import java.util.function.Supplier;
+
+/** Test class */
+public final class Lambda {
+
+  private Lambda() {}
+
+  private static <T> T foo(Supplier<T> fn) {
+    return fn.get();
+  }
+
+  public static void main(String[] args) {
+    String unused = foo(() -> "Hello, world!");
+  }
+}

src/tools/android/java/com/google/devtools/build/android/r8/BUILD

@@ -89,3 +89,25 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/worker:work_request_handlers",
     ],
 )
+
+java_binary(
+    name = "desugar",
+    jvm_flags = [
+        # b/71513487
+        "-XX:+TieredCompilation",
+        "-XX:TieredStopAtLevel=1",
+        "-Xms8g",
+        "-Xmx8g",
+        # b/172508621
+        "-Dcom.android.tools.r8.sortMethodsOnCfWriting",
+        "-Dcom.android.tools.r8.allowAllDesugaredInput",
+        "-Dcom.android.tools.r8.noCfMarkerForDesugaredCode",
+        "-Dcom.android.tools.r8.lambdaClassFieldsNotFinal",
+        "-Dcom.android.tools.r8.createSingletonsForStatelessLambdas",
+    ],
+    main_class = "com.google.devtools.build.android.r8.Desugar",
+    visibility = ["//src/test/java/com/google/devtools/build/android/r8:__pkg__"],
+    runtime_deps = [
+        ":r8",
+    ],
+)

src/tools/android/java/com/google/devtools/build/android/r8/CompatDexBuilder.java

@@ -24,8 +24,11 @@
 import com.android.tools.r8.D8Command;
 import com.android.tools.r8.DexIndexedConsumer;
 import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.SyntheticInfoConsumer;
+import com.android.tools.r8.SyntheticInfoConsumerData;
 import com.android.tools.r8.origin.ArchiveEntryOrigin;
 import com.android.tools.r8.origin.PathOrigin;
+import com.android.tools.r8.references.ClassReference;
 import com.google.auto.value.AutoValue;
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
@@ -67,8 +70,47 @@
 public class CompatDexBuilder {
   private static final long ONE_MEG = 1024 * 1024;
 
+  private static class ContextConsumer implements SyntheticInfoConsumer {
+
+    // After compilation this will be non-null iff the compiled class is a D8 synthesized class.
+    ClassReference sythesizedPrimaryClass = null;
+
+    // If the above is non-null then this will be the non-synthesized context class that caused
+    // D8 to synthesize the above class.
+    ClassReference contextOfSynthesizedClass = null;
+
+    @Nullable
+    String getContextMapping() {
+      if (sythesizedPrimaryClass != null) {
+        return sythesizedPrimaryClass.getBinaryName()
+            + ";"
+            + contextOfSynthesizedClass.getBinaryName();
+      }
+      return null;
+    }
+
+    @Override
+    public synchronized void acceptSyntheticInfo(SyntheticInfoConsumerData data) {
+      verify(
+          sythesizedPrimaryClass == null || sythesizedPrimaryClass.equals(data.getSyntheticClass()),
+          "The single input classfile should ensure this has one value.");
+      verify(
+          contextOfSynthesizedClass == null
+              || contextOfSynthesizedClass.equals(data.getSynthesizingContextClass()),
+          "The single input classfile should ensure this has one value.");
+      sythesizedPrimaryClass = data.getSyntheticClass();
+      contextOfSynthesizedClass = data.getSynthesizingContextClass();
+    }
+
+    @Override
+    public void finished() {
+      // Do nothing.
+    }
+  }
+
   private static class DexConsumer implements DexIndexedConsumer {
 
+    final ContextConsumer contextConsumer = new ContextConsumer();
     byte[] bytes;
 
     @Override
@@ -86,6 +128,10 @@ void setBytes(byte[] byteCode) {
       this.bytes = byteCode;
     }
 
+    ContextConsumer getContextConsumer() {
+      return contextConsumer;
+    }
+
     @Override
     public void finished(DiagnosticsHandler handler) {
       // Do nothing.
@@ -269,10 +315,23 @@ private void dexEntries(@Nullable Cache<DexingKeyR8, byte[]> dexCache, List<Stri
                           minSdkVersion,
                           executor)));
         }
+        StringBuilder contextMappingBuilder = new StringBuilder();
         for (int i = 0; i < futures.size(); i++) {
           ZipEntry entry = toDex.get(i);
           DexConsumer consumer = futures.get(i).get();
           ZipUtils.addEntry(entry.getName() + ".dex", consumer.getBytes(), ZipEntry.STORED, out);
+          String mapping = consumer.getContextConsumer().getContextMapping();
+          if (mapping != null) {
+            contextMappingBuilder.append(mapping).append('\n');
+          }
+        }
+        String contextMapping = contextMappingBuilder.toString();
+        if (!contextMapping.isEmpty()) {
+          ZipUtils.addEntry(
+              "META-INF/synthetic-contexts.map",
+              contextMapping.getBytes(UTF_8),
+              ZipEntry.STORED,
+              out);
         }
       }
     } finally {
@@ -292,6 +351,7 @@ private DexConsumer dexEntry(
     D8Command.Builder builder = D8Command.builder();
     builder
         .setProgramConsumer(consumer)
+        .setSyntheticInfoConsumer(consumer.getContextConsumer())
         .setMode(mode)
         .setMinApiLevel(minSdkVersion)
         .setDisableDesugaring(true)