[Partition] Add test for R8 emitting a zip file container as mapping

This will also introduce a ProguardMapPartitionConsumer to have R8 automatically create a partioned output.

Bug: b/274735214
Change-Id: Iffddb5af7afb1db7ee4fac67a6b4e4a1904cd227
diff --git a/src/main/java/com/android/tools/r8/naming/ProguardMapPartitionConsumer.java b/src/main/java/com/android/tools/r8/naming/ProguardMapPartitionConsumer.java
new file mode 100644
index 0000000..451fd6c
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/naming/ProguardMapPartitionConsumer.java
@@ -0,0 +1,96 @@
+// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.naming;
+
+import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.ProguardMapConsumer;
+import com.android.tools.r8.errors.Unreachable;
+import com.android.tools.r8.retrace.MappingPartition;
+import com.android.tools.r8.retrace.MappingPartitionMetadata;
+import com.android.tools.r8.retrace.ProguardMapPartitioner;
+import com.android.tools.r8.retrace.internal.ProguardMapProducerInternal;
+import java.io.IOException;
+import java.util.function.Consumer;
+
+public class ProguardMapPartitionConsumer extends ProguardMapConsumer {
+
+  private final Consumer<MappingPartition> mappingPartitionConsumer;
+  private final Consumer<MappingPartitionMetadata> metadataConsumer;
+  private final Runnable finishedConsumer;
+  private final DiagnosticsHandler diagnosticsHandler;
+
+  private ProguardMapPartitionConsumer(
+      Consumer<MappingPartition> mappingPartitionConsumer,
+      Consumer<MappingPartitionMetadata> metadataConsumer,
+      Runnable finishedConsumer,
+      DiagnosticsHandler diagnosticsHandler) {
+    this.mappingPartitionConsumer = mappingPartitionConsumer;
+    this.metadataConsumer = metadataConsumer;
+    this.finishedConsumer = finishedConsumer;
+    this.diagnosticsHandler = diagnosticsHandler;
+  }
+
+  @Override
+  public void accept(ProguardMapMarkerInfo makerInfo, ClassNameMapper classNameMapper) {
+    try {
+      // TODO(b/274735214): Ensure we get markerInfo consumed as well.
+      MappingPartitionMetadata run =
+          ProguardMapPartitioner.builder(diagnosticsHandler)
+              .setProguardMapProducer(new ProguardMapProducerInternal(classNameMapper))
+              .setPartitionConsumer(mappingPartitionConsumer)
+              // Setting these do not actually do anything currently since there is no parsing.
+              .setAllowEmptyMappedRanges(false)
+              .setAllowExperimentalMapping(false)
+              .build()
+              .run();
+      metadataConsumer.accept(run);
+    } catch (IOException exception) {
+      throw new Unreachable("IOExceptions should only occur when parsing");
+    }
+  }
+
+  @Override
+  public void finished(DiagnosticsHandler handler) {
+    finishedConsumer.run();
+  }
+
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  public static class Builder {
+
+    private Consumer<MappingPartition> mappingPartitionConsumer;
+    private Consumer<MappingPartitionMetadata> metadataConsumer;
+    private Runnable finishedConsumer;
+    private DiagnosticsHandler diagnosticsHandler;
+
+    public Builder setMappingPartitionConsumer(
+        Consumer<MappingPartition> mappingPartitionConsumer) {
+      this.mappingPartitionConsumer = mappingPartitionConsumer;
+      return this;
+    }
+
+    public Builder setMetadataConsumer(Consumer<MappingPartitionMetadata> metadataConsumer) {
+      this.metadataConsumer = metadataConsumer;
+      return this;
+    }
+
+    public Builder setFinishedConsumer(Runnable finishedConsumer) {
+      this.finishedConsumer = finishedConsumer;
+      return this;
+    }
+
+    public Builder setDiagnosticsHandler(DiagnosticsHandler diagnosticsHandler) {
+      this.diagnosticsHandler = diagnosticsHandler;
+      return this;
+    }
+
+    public ProguardMapPartitionConsumer build() {
+      return new ProguardMapPartitionConsumer(
+          mappingPartitionConsumer, metadataConsumer, finishedConsumer, diagnosticsHandler);
+    }
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierImpl.java b/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierImpl.java
index 915fc81..7d663fb 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierImpl.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/PartitionMappingSupplierImpl.java
@@ -77,6 +77,7 @@
   }
 
   private void registerKeyUse(String key) {
+    // TODO(b/274735214): only call the register partition if we have a partition for it.
     if (!builtKeys.contains(key) && pendingKeys.add(key)) {
       registerPartitionCallback.register(key);
     }
@@ -103,9 +104,14 @@
     }
     for (String pendingKey : pendingKeys) {
       try {
+        byte[] suppliedPartition = partitionSupplier.get(pendingKey);
+        // TODO(b/274735214): only expect a partition if have generated one for the key.
+        if (suppliedPartition == null) {
+          continue;
+        }
         LineReader reader =
             new ProguardMapReaderWithFilteringInputBuffer(
-                new ByteArrayInputStream(partitionSupplier.get(pendingKey)), alwaysTrue(), true);
+                new ByteArrayInputStream(suppliedPartition), alwaysTrue(), true);
         classNameMapper =
             ClassNameMapper.mapperFromLineReaderWithFiltering(
                     reader, metadata.getMapVersion(), diagnosticsHandler, true, allowExperimental)
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapPartitionerOnClassNameToText.java b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapPartitionerOnClassNameToText.java
index 0298816..68adc36 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapPartitionerOnClassNameToText.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapPartitionerOnClassNameToText.java
@@ -25,6 +25,7 @@
 import com.android.tools.r8.utils.ExceptionDiagnostic;
 import com.android.tools.r8.utils.StringDiagnostic;
 import com.android.tools.r8.utils.StringUtils;
+import com.android.tools.r8.utils.TriConsumer;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
@@ -33,7 +34,6 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Map.Entry;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -62,8 +62,29 @@
     this.mappingPartitionKeyStrategy = mappingPartitionKeyStrategy;
   }
 
-  @Override
-  public MappingPartitionMetadata run() throws IOException {
+  private ClassNameMapper getPartitionsFromProguardMapProducer(
+      TriConsumer<ClassNameMapper, ClassNamingForNameMapper, String> consumer) throws IOException {
+    if (proguardMapProducer instanceof ProguardMapProducerInternal) {
+      return getPartitionsFromInternalProguardMapProducer(consumer);
+    } else {
+      return getPartitionsFromStringBackedProguardMapProducer(consumer);
+    }
+  }
+
+  private ClassNameMapper getPartitionsFromInternalProguardMapProducer(
+      TriConsumer<ClassNameMapper, ClassNamingForNameMapper, String> consumer) {
+    ClassNameMapper classNameMapper =
+        ((ProguardMapProducerInternal) proguardMapProducer).getClassNameMapper();
+    classNameMapper
+        .getClassNameMappings()
+        .forEach(
+            (key, mappingForClass) ->
+                consumer.accept(classNameMapper, mappingForClass, mappingForClass.toString()));
+    return classNameMapper;
+  }
+
+  private ClassNameMapper getPartitionsFromStringBackedProguardMapProducer(
+      TriConsumer<ClassNameMapper, ClassNamingForNameMapper, String> consumer) throws IOException {
     PartitionLineReader reader =
         new PartitionLineReader(
             proguardMapProducer.isFileBacked()
@@ -75,9 +96,11 @@
     // mappings.
     ClassNameMapper classMapper =
         ClassNameMapper.mapperFromLineReaderWithFiltering(
-            reader, MapVersion.MAP_VERSION_UNKNOWN, diagnosticsHandler, true, true);
-    HashSet<String> keys = new LinkedHashSet<>();
-    // We can then iterate over all sections.
+            reader,
+            MapVersion.MAP_VERSION_UNKNOWN,
+            diagnosticsHandler,
+            allowEmptyMappedRanges,
+            allowExperimentalMapping);
     reader.forEachClassMapping(
         (classMapping, entries) -> {
           try {
@@ -85,33 +108,44 @@
             ClassNameMapper classNameMapper =
                 ClassNameMapper.mapperFromString(
                     payload, null, allowEmptyMappedRanges, allowExperimentalMapping, false);
-            if (classNameMapper.getClassNameMappings().size() != 1) {
+            Map<String, ClassNamingForNameMapper> classNameMappings =
+                classNameMapper.getClassNameMappings();
+            if (classNameMappings.size() != 1) {
               diagnosticsHandler.error(
                   new StringDiagnostic("Multiple class names in payload\n: " + payload));
               return;
             }
-            Entry<String, ClassNamingForNameMapper> currentClassMapping =
-                classNameMapper.getClassNameMappings().entrySet().iterator().next();
-            ClassNamingForNameMapper value = currentClassMapping.getValue();
-            Set<String> seenMappings = new HashSet<>();
-            StringBuilder payloadWithClassReferences = new StringBuilder();
-            value.visitAllFullyQualifiedReferences(
-                holder -> {
-                  if (seenMappings.add(holder)) {
-                    payloadWithClassReferences.append(
-                        getSourceFileMapping(holder, classMapper.getSourceFile(holder)));
-                  }
-                });
-            payloadWithClassReferences.append(payload);
-            mappingPartitionConsumer.accept(
-                new MappingPartitionImpl(
-                    currentClassMapping.getKey(),
-                    payloadWithClassReferences.toString().getBytes(StandardCharsets.UTF_8)));
-            keys.add(currentClassMapping.getKey());
+            consumer.accept(classMapper, classNameMappings.values().iterator().next(), payload);
           } catch (IOException e) {
             diagnosticsHandler.error(new ExceptionDiagnostic(e));
           }
         });
+    return classMapper;
+  }
+
+  @Override
+  public MappingPartitionMetadata run() throws IOException {
+    HashSet<String> keys = new LinkedHashSet<>();
+    // We can then iterate over all sections.
+    ClassNameMapper classMapper =
+        getPartitionsFromProguardMapProducer(
+            (classNameMapper, classNamingForNameMapper, payload) -> {
+              Set<String> seenMappings = new HashSet<>();
+              StringBuilder payloadWithClassReferences = new StringBuilder();
+              classNamingForNameMapper.visitAllFullyQualifiedReferences(
+                  holder -> {
+                    if (seenMappings.add(holder)) {
+                      payloadWithClassReferences.append(
+                          getSourceFileMapping(holder, classNameMapper.getSourceFile(holder)));
+                    }
+                  });
+              payloadWithClassReferences.append(payload);
+              mappingPartitionConsumer.accept(
+                  new MappingPartitionImpl(
+                      classNamingForNameMapper.renamedName,
+                      payloadWithClassReferences.toString().getBytes(StandardCharsets.UTF_8)));
+              keys.add(classNamingForNameMapper.renamedName);
+            });
     MapVersion mapVersion = MapVersion.MAP_VERSION_UNKNOWN;
     MapVersionMappingInformation mapVersionInfo = classMapper.getFirstMapVersionInformation();
     if (mapVersionInfo != null) {
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapProducerInternal.java b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapProducerInternal.java
new file mode 100644
index 0000000..6650458
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMapProducerInternal.java
@@ -0,0 +1,29 @@
+// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.retrace.internal;
+
+import com.android.tools.r8.errors.Unreachable;
+import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.retrace.ProguardMapProducer;
+import java.io.IOException;
+import java.io.InputStream;
+
+public class ProguardMapProducerInternal implements ProguardMapProducer {
+
+  private final ClassNameMapper classNameMapper;
+
+  public ProguardMapProducerInternal(ClassNameMapper classNameMapper) {
+    this.classNameMapper = classNameMapper;
+  }
+
+  public ClassNameMapper getClassNameMapper() {
+    return classNameMapper;
+  }
+
+  @Override
+  public InputStream get() throws IOException {
+    throw new Unreachable("Should never get on ProguardMapProducerInternal");
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/naming/retrace/StackTrace.java b/src/test/java/com/android/tools/r8/naming/retrace/StackTrace.java
index 0446cd5..82a318d 100644
--- a/src/test/java/com/android/tools/r8/naming/retrace/StackTrace.java
+++ b/src/test/java/com/android/tools/r8/naming/retrace/StackTrace.java
@@ -11,6 +11,7 @@
 import com.android.tools.r8.SingleTestRunResult;
 import com.android.tools.r8.ToolHelper.DexVm;
 import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.retrace.MappingSupplier;
 import com.android.tools.r8.retrace.ProguardMapProducer;
 import com.android.tools.r8.retrace.ProguardMappingSupplier;
 import com.android.tools.r8.retrace.Retrace;
@@ -385,17 +386,21 @@
   }
 
   public StackTrace retrace(String map, boolean allowExperimentalMapping) {
+    return retrace(
+        ProguardMappingSupplier.builder()
+            .setProguardMapProducer(ProguardMapProducer.fromString(map))
+            .setAllowExperimental(allowExperimentalMapping)
+            .build());
+  }
+
+  public StackTrace retrace(MappingSupplier<?> mappingSupplier) {
     Box<List<String>> box = new Box<>();
     List<String> stackTrace =
         stackTraceLines.stream().map(line -> line.originalLine).collect(Collectors.toList());
     stackTrace.add(0, exceptionLine);
     Retrace.run(
         RetraceCommand.builder()
-            .setMappingSupplier(
-                ProguardMappingSupplier.builder()
-                    .setProguardMapProducer(ProguardMapProducer.fromString(map))
-                    .setAllowExperimental(allowExperimentalMapping)
-                    .build())
+            .setMappingSupplier(mappingSupplier)
             .setStackTrace(stackTrace)
             .setRetracedStackTraceConsumer(box::set)
             .build());
diff --git a/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java b/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
new file mode 100644
index 0000000..211121c
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
@@ -0,0 +1,154 @@
+// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.retrace.partition;
+
+import static com.android.tools.r8.naming.retrace.StackTrace.isSame;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertTrue;
+
+import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.ProguardMapConsumer;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestDiagnosticMessagesImpl;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import com.android.tools.r8.naming.ProguardMapPartitionConsumer;
+import com.android.tools.r8.naming.retrace.StackTrace;
+import com.android.tools.r8.naming.retrace.StackTrace.StackTraceLine;
+import com.android.tools.r8.retrace.PartitionMappingSupplier;
+import com.android.tools.r8.retrace.partition.testclasses.R8ZipContainerMappingFileTestClasses;
+import com.android.tools.r8.retrace.partition.testclasses.R8ZipContainerMappingFileTestClasses.Main;
+import com.android.tools.r8.utils.ZipUtils.ZipBuilder;
+import com.google.common.io.ByteStreams;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipFile;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class R8ZipContainerMappingFileTest extends TestBase {
+
+  private static final String SOURCE_FILE = "R8ZipContainerMappingFileTestClasses.java";
+
+  private final StackTrace EXPECTED =
+      StackTrace.builder()
+          .add(
+              StackTraceLine.builder()
+                  .setClassName(typeName(R8ZipContainerMappingFileTestClasses.Thrower.class))
+                  .setMethodName("throwError")
+                  .setFileName(SOURCE_FILE)
+                  .setLineNumber(13)
+                  .build())
+          .add(
+              StackTraceLine.builder()
+                  .setClassName(typeName(R8ZipContainerMappingFileTestClasses.Main.class))
+                  .setMethodName("main")
+                  .setFileName(SOURCE_FILE)
+                  .setLineNumber(21)
+                  .build())
+          .build();
+
+  @Parameter() public TestParameters parameters;
+
+  @Parameters(name = "{0}")
+  public static TestParametersCollection data() {
+    return getTestParameters().withAllRuntimesAndApiLevels().build();
+  }
+
+  @Test
+  public void testRuntime() throws Exception {
+    testForRuntime(parameters)
+        .addInnerClasses(R8ZipContainerMappingFileTestClasses.class)
+        .run(parameters.getRuntime(), Main.class)
+        .assertFailureWithErrorThatThrows(RuntimeException.class)
+        .inspectStackTrace(
+            stackTrace -> {
+              assertThat(stackTrace, isSame(EXPECTED));
+            });
+  }
+
+  @Test
+  public void testR8() throws Exception {
+    Path pgMapFile = temp.newFile("mapping.zip").toPath();
+    DiagnosticsHandler diagnosticsHandler = new TestDiagnosticMessagesImpl();
+    ProguardMapConsumer partitionZipConsumer =
+        createPartitionZipConsumer(pgMapFile, diagnosticsHandler);
+    StackTrace originalStackTrace =
+        testForR8(parameters.getBackend())
+            .addInnerClasses(R8ZipContainerMappingFileTestClasses.class)
+            .setMinApi(parameters)
+            .addKeepMainRule(Main.class)
+            .addKeepAttributeSourceFile()
+            .addKeepAttributeLineNumberTable()
+            .addOptionsModification(options -> options.proguardMapConsumer = partitionZipConsumer)
+            .run(parameters.getRuntime(), Main.class)
+            .assertFailureWithErrorThatThrows(RuntimeException.class)
+            .getOriginalStackTrace();
+
+    assertTrue(Files.exists(pgMapFile));
+    assertThat(
+        originalStackTrace.retrace(createMappingSupplierFromPartitionZip(pgMapFile)),
+        isSame(EXPECTED));
+  }
+
+  private ProguardMapConsumer createPartitionZipConsumer(
+      Path pgMapFile, DiagnosticsHandler diagnosticsHandler) throws IOException {
+    ZipBuilder zipBuilder = ZipBuilder.builder(pgMapFile);
+    return ProguardMapPartitionConsumer.builder()
+        .setMappingPartitionConsumer(
+            mappingPartition -> {
+              try {
+                zipBuilder.addBytes(mappingPartition.getKey(), mappingPartition.getPayload());
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            })
+        .setMetadataConsumer(
+            metadata -> {
+              try {
+                zipBuilder.addBytes("METADATA", metadata.getBytes());
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            })
+        .setFinishedConsumer(
+            () -> {
+              try {
+                zipBuilder.build();
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            })
+        .setDiagnosticsHandler(diagnosticsHandler)
+        .build();
+  }
+
+  private PartitionMappingSupplier createMappingSupplierFromPartitionZip(Path pgMapFile)
+      throws IOException {
+    ZipFile zipFile = new ZipFile(pgMapFile.toFile());
+    byte[] metadata = ByteStreams.toByteArray(zipFile.getInputStream(zipFile.getEntry("METADATA")));
+    return PartitionMappingSupplier.builder()
+        .setMetadata(metadata)
+        .setMappingPartitionFromKeySupplier(
+            key -> {
+              try {
+                // TODO(b/274735214): The key should exist.
+                ZipEntry entry = zipFile.getEntry(key);
+                return entry == null
+                    ? null
+                    : ByteStreams.toByteArray(zipFile.getInputStream(entry));
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            })
+        .build();
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/retrace/partition/testclasses/R8ZipContainerMappingFileTestClasses.java b/src/test/java/com/android/tools/r8/retrace/partition/testclasses/R8ZipContainerMappingFileTestClasses.java
new file mode 100644
index 0000000..e81d031
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/partition/testclasses/R8ZipContainerMappingFileTestClasses.java
@@ -0,0 +1,24 @@
+// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.retrace.partition.testclasses;
+
+public class R8ZipContainerMappingFileTestClasses {
+
+  public static class Thrower {
+
+    public static void throwError() {
+      if (System.currentTimeMillis() > 0) {
+        throw new RuntimeException("Hello World");
+      }
+    }
+  }
+
+  public static class Main {
+
+    public static void main(String[] args) {
+      Thrower.throwError();
+    }
+  }
+}