Extend R8CommandBuilder with setPartitionMapConsumer

Bug: b/274735214
Change-Id: I385fd76bba2711ffa902f37fed78844bce8bce59
diff --git a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
index 285ef1b..cc62da7 100644
--- a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
+++ b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
@@ -22,6 +22,7 @@
 import com.android.tools.r8.utils.FileUtils;
 import com.android.tools.r8.utils.InternalOptions.DesugarState;
 import com.android.tools.r8.utils.ListUtils;
+import com.android.tools.r8.utils.MapConsumerUtils;
 import com.android.tools.r8.utils.ProgramConsumerUtils;
 import com.android.tools.r8.utils.Reporter;
 import com.android.tools.r8.utils.ThreadUtils;
@@ -281,28 +282,25 @@
     private int minApiLevel = 0;
     private int threadCount = ThreadUtils.NOT_SPECIFIED;
     protected DesugarState desugarState = DesugarState.ON;
-    private List<StringResource> desugaredLibrarySpecificationResources = new ArrayList<>();
+    private final List<StringResource> desugaredLibrarySpecificationResources = new ArrayList<>();
     private boolean includeClassesChecksum = false;
     private boolean optimizeMultidexForLinearAlloc = false;
     private BiPredicate<String, Long> dexClassChecksumFilter = (name, checksum) -> true;
-    private List<AssertionsConfiguration> assertionsConfiguration = new ArrayList<>();
-    private List<Consumer<Inspector>> outputInspections = new ArrayList<>();
+    private final List<AssertionsConfiguration> assertionsConfiguration = new ArrayList<>();
+    private final List<Consumer<Inspector>> outputInspections = new ArrayList<>();
     protected StringConsumer proguardMapConsumer = null;
+    protected PartitionMapConsumer partitionMapConsumer = null;
     private DumpInputFlags dumpInputFlags = DumpInputFlags.getDefault();
     private MapIdProvider mapIdProvider = null;
     private SourceFileProvider sourceFileProvider = null;
     private boolean isAndroidPlatformBuild = false;
-    private List<ArtProfileForRewriting> artProfilesForRewriting = new ArrayList<>();
-    private List<StartupProfileProvider> startupProfileProviders = new ArrayList<>();
+    private final List<ArtProfileForRewriting> artProfilesForRewriting = new ArrayList<>();
+    private final List<StartupProfileProvider> startupProfileProviders = new ArrayList<>();
     private ClassConflictResolver classConflictResolver = null;
     private CancelCompilationChecker cancelCompilationChecker = null;
 
     abstract CompilationMode defaultCompilationMode();
 
-    Builder() {
-      mode = defaultCompilationMode();
-    }
-
     Builder(DiagnosticsHandler diagnosticsHandler) {
       super(diagnosticsHandler);
       mode = defaultCompilationMode();
@@ -395,6 +393,33 @@
     }
 
     /**
+     * Set an output destination to which partition-map content should be written.
+     *
+     * <p>This is a short-hand for setting a {@link PartitionMapConsumer} using {@link
+     * #setPartitionMapConsumer}. Note that any subsequent call to this method or {@link
+     * #setPartitionMapConsumer} will override the previous setting.
+     *
+     * @param partitionMapOutput File-system path to write output at.
+     */
+    B setPartitionMapOutputPath(Path partitionMapOutput) {
+      assert partitionMapOutput != null;
+      return setPartitionMapConsumer(MapConsumerUtils.createZipConsumer(partitionMapOutput));
+    }
+
+    /**
+     * Set a consumer for receiving the partition map content.
+     *
+     * <p>Note that any subsequent call to this method or {@link #setPartitionMapOutputPath} will
+     * override the previous setting.
+     *
+     * @param partitionMapConsumer Consumer to receive the content once produced.
+     */
+    B setPartitionMapConsumer(PartitionMapConsumer partitionMapConsumer) {
+      this.partitionMapConsumer = partitionMapConsumer;
+      return self();
+    }
+
+    /**
      * Get the main dex list consumer that will receive the final complete main dex list.
      */
     public StringConsumer getMainDexListConsumer() {
diff --git a/src/main/java/com/android/tools/r8/D8Command.java b/src/main/java/com/android/tools/r8/D8Command.java
index 137cb8b..a97ec22 100644
--- a/src/main/java/com/android/tools/r8/D8Command.java
+++ b/src/main/java/com/android/tools/r8/D8Command.java
@@ -665,13 +665,10 @@
     internal.setSyntheticInfoConsumer(syntheticInfoConsumer);
     internal.desugarGraphConsumer = desugarGraphConsumer;
     internal.mainDexKeepRules = mainDexKeepRules;
-    internal.proguardMapConsumer =
+    internal.mapConsumer =
         proguardMapConsumer == null
             ? null
-            : ProguardMapStringConsumer.builder()
-                .setStringConsumer(proguardMapConsumer)
-                .setDiagnosticsHandler(getReporter())
-                .build();
+            : ProguardMapStringConsumer.builder().setStringConsumer(proguardMapConsumer).build();
     internal.lineNumberOptimization =
         !internal.debug && proguardMapConsumer != null
             ? LineNumberOptimization.ON
diff --git a/src/main/java/com/android/tools/r8/MapConsumerToPartitionMapConsumer.java b/src/main/java/com/android/tools/r8/MapConsumerToPartitionMapConsumer.java
new file mode 100644
index 0000000..fcd3113
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/MapConsumerToPartitionMapConsumer.java
@@ -0,0 +1,49 @@
+// 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;
+
+import com.android.tools.r8.errors.Unreachable;
+import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.naming.MapConsumer;
+import com.android.tools.r8.naming.ProguardMapMarkerInfo;
+import com.android.tools.r8.retrace.ProguardMapPartitioner;
+import com.android.tools.r8.retrace.internal.ProguardMapProducerInternal;
+import java.io.IOException;
+
+public class MapConsumerToPartitionMapConsumer implements MapConsumer {
+
+  protected final PartitionMapConsumer partitionMapConsumer;
+
+  protected MapConsumerToPartitionMapConsumer(PartitionMapConsumer partitionMapConsumer) {
+    assert partitionMapConsumer != null;
+    this.partitionMapConsumer = partitionMapConsumer;
+  }
+
+  @Override
+  public void accept(
+      DiagnosticsHandler diagnosticsHandler,
+      ProguardMapMarkerInfo makerInfo,
+      ClassNameMapper classNameMapper) {
+    try {
+      classNameMapper.setPreamble(makerInfo.toPreamble());
+      partitionMapConsumer.acceptMappingPartitionMetadata(
+          ProguardMapPartitioner.builder(diagnosticsHandler)
+              .setProguardMapProducer(new ProguardMapProducerInternal(classNameMapper))
+              .setPartitionConsumer(partitionMapConsumer::acceptMappingPartition)
+              // Modifying these do not actually do anything currently since there is no parsing.
+              .setAllowEmptyMappedRanges(false)
+              .setAllowExperimentalMapping(false)
+              .build()
+              .run());
+    } catch (IOException exception) {
+      throw new Unreachable("IOExceptions should only occur when parsing");
+    }
+  }
+
+  @Override
+  public void finished(DiagnosticsHandler handler) {
+    partitionMapConsumer.finished(handler);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/PartitionMapConsumer.java b/src/main/java/com/android/tools/r8/PartitionMapConsumer.java
new file mode 100644
index 0000000..ea232d4
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/PartitionMapConsumer.java
@@ -0,0 +1,16 @@
+// 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;
+
+import com.android.tools.r8.retrace.MappingPartition;
+import com.android.tools.r8.retrace.MappingPartitionMetadata;
+
+@Keep
+public interface PartitionMapConsumer extends Finishable {
+
+  void acceptMappingPartition(MappingPartition mappingPartition);
+
+  void acceptMappingPartitionMetadata(MappingPartitionMetadata mappingPartitionMetadata);
+}
diff --git a/src/main/java/com/android/tools/r8/ProguardMapConsumer.java b/src/main/java/com/android/tools/r8/ProguardMapConsumer.java
deleted file mode 100644
index f00b428..0000000
--- a/src/main/java/com/android/tools/r8/ProguardMapConsumer.java
+++ /dev/null
@@ -1,13 +0,0 @@
-// 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;
-
-import com.android.tools.r8.naming.ClassNameMapper;
-import com.android.tools.r8.naming.ProguardMapMarkerInfo;
-
-public abstract class ProguardMapConsumer implements Finishable {
-
-  public abstract void accept(ProguardMapMarkerInfo makerInfo, ClassNameMapper classNameMapper);
-}
diff --git a/src/main/java/com/android/tools/r8/R8Command.java b/src/main/java/com/android/tools/r8/R8Command.java
index edec018..39e7996 100644
--- a/src/main/java/com/android/tools/r8/R8Command.java
+++ b/src/main/java/com/android/tools/r8/R8Command.java
@@ -4,6 +4,7 @@
 package com.android.tools.r8;
 
 import static com.android.tools.r8.utils.InternalOptions.DETERMINISTIC_DEBUGGING;
+import static com.android.tools.r8.utils.MapConsumerUtils.wrapExistingMapConsumerIfNotNull;
 
 import com.android.tools.r8.ProgramResource.Kind;
 import com.android.tools.r8.dex.Marker.Tool;
@@ -18,6 +19,7 @@
 import com.android.tools.r8.keepanno.asm.KeepEdgeReader;
 import com.android.tools.r8.keepanno.ast.KeepEdge;
 import com.android.tools.r8.keepanno.keeprules.KeepRuleExtractor;
+import com.android.tools.r8.naming.MapConsumer;
 import com.android.tools.r8.naming.ProguardMapStringConsumer;
 import com.android.tools.r8.naming.SourceFileRewriter;
 import com.android.tools.r8.origin.Origin;
@@ -299,6 +301,34 @@
     }
 
     /**
+     * Set an output destination to which r8-map content should be written.
+     *
+     * <p>This is a short-hand for setting a {@link MapConsumerToPartitionMapConsumer} using {@link
+     * #setPartitionMapConsumer}. Note that any subsequent call to this method or {@link
+     * #setPartitionMapConsumer} will override the previous setting.
+     *
+     * @param partitionMapOutput File-system path to write output at.
+     */
+    @Override
+    public Builder setPartitionMapOutputPath(Path partitionMapOutput) {
+      assert partitionMapOutput != null;
+      return super.setPartitionMapOutputPath(partitionMapOutput);
+    }
+
+    /**
+     * Set a consumer for receiving the r8-map content.
+     *
+     * <p>Note that any subsequent call to this method or {@link #setPartitionMapOutputPath} will
+     * override the previous setting.
+     *
+     * @param partitionMapConsumer Consumer to receive the content once produced.
+     */
+    @Override
+    public Builder setPartitionMapConsumer(PartitionMapConsumer partitionMapConsumer) {
+      return super.setPartitionMapConsumer(partitionMapConsumer);
+    }
+
+    /**
      * Set a consumer for receiving the keep rules to use when compiling the desugared library for
      * the program being compiled in this compilation.
      *
@@ -632,6 +662,7 @@
               forceProguardCompatibility,
               includeDataResources,
               proguardMapConsumer,
+              partitionMapConsumer,
               proguardUsageConsumer,
               proguardSeedsConsumer,
               proguardConfigurationConsumer,
@@ -832,6 +863,7 @@
   private final boolean forceProguardCompatibility;
   private final Optional<Boolean> includeDataResources;
   private final StringConsumer proguardMapConsumer;
+  private final PartitionMapConsumer partitionMapConsumer;
   private final StringConsumer proguardUsageConsumer;
   private final StringConsumer proguardSeedsConsumer;
   private final StringConsumer proguardConfigurationConsumer;
@@ -912,6 +944,7 @@
       boolean forceProguardCompatibility,
       Optional<Boolean> includeDataResources,
       StringConsumer proguardMapConsumer,
+      PartitionMapConsumer partitionMapConsumer,
       StringConsumer proguardUsageConsumer,
       StringConsumer proguardSeedsConsumer,
       StringConsumer proguardConfigurationConsumer,
@@ -969,6 +1002,7 @@
     this.forceProguardCompatibility = forceProguardCompatibility;
     this.includeDataResources = includeDataResources;
     this.proguardMapConsumer = proguardMapConsumer;
+    this.partitionMapConsumer = partitionMapConsumer;
     this.proguardUsageConsumer = proguardUsageConsumer;
     this.proguardSeedsConsumer = proguardSeedsConsumer;
     this.proguardConfigurationConsumer = proguardConfigurationConsumer;
@@ -992,6 +1026,7 @@
     forceProguardCompatibility = false;
     includeDataResources = null;
     proguardMapConsumer = null;
+    partitionMapConsumer = null;
     proguardUsageConsumer = null;
     proguardSeedsConsumer = null;
     proguardConfigurationConsumer = null;
@@ -1070,13 +1105,15 @@
             proguardMapConsumer,
             proguardConfiguration.isPrintMapping(),
             proguardConfiguration.getPrintMappingFile());
-    internal.proguardMapConsumer =
-        stringConsumer == null
-            ? null
-            : ProguardMapStringConsumer.builder()
-                .setStringConsumer(stringConsumer)
-                .setDiagnosticsHandler(getReporter())
-                .build();
+    MapConsumer mapConsumer =
+        wrapExistingMapConsumerIfNotNull(
+            internal.mapConsumer, partitionMapConsumer, MapConsumerToPartitionMapConsumer::new);
+    internal.mapConsumer =
+        wrapExistingMapConsumerIfNotNull(
+            mapConsumer,
+            stringConsumer,
+            nonNullStringConsumer ->
+                ProguardMapStringConsumer.builder().setStringConsumer(stringConsumer).build());
 
     // Amend the usage information consumer with options from the proguard configuration.
     internal.usageInformationConsumer =
diff --git a/src/main/java/com/android/tools/r8/bisect/Bisect.java b/src/main/java/com/android/tools/r8/bisect/Bisect.java
index caa1c27..c4c8170 100644
--- a/src/main/java/com/android/tools/r8/bisect/Bisect.java
+++ b/src/main/java/com/android/tools/r8/bisect/Bisect.java
@@ -5,7 +5,6 @@
 
 import com.android.tools.r8.OutputMode;
 import com.android.tools.r8.ProgramConsumer;
-import com.android.tools.r8.ProguardMapConsumer;
 import com.android.tools.r8.bisect.BisectOptions.Result;
 import com.android.tools.r8.dex.ApplicationReader;
 import com.android.tools.r8.dex.ApplicationWriter;
@@ -14,6 +13,7 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexApplication;
 import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.naming.MapConsumer;
 import com.android.tools.r8.synthesis.SyntheticItems.GlobalSyntheticsStrategy;
 import com.android.tools.r8.utils.AndroidApp;
 import com.android.tools.r8.utils.AndroidAppConsumers;
@@ -184,7 +184,7 @@
     InternalOptions options = app.options;
     // Save the original consumers, so they can be unwrapped after write.
     ProgramConsumer programConsumer = options.programConsumer;
-    ProguardMapConsumer proguardMapConsumer = options.proguardMapConsumer;
+    MapConsumer mapConsumer = options.mapConsumer;
     AndroidAppConsumers compatSink = new AndroidAppConsumers(options);
     ApplicationWriter writer =
         ApplicationWriter.create(
@@ -196,7 +196,7 @@
     compatSink.build().writeToDirectory(output, OutputMode.DexIndexed);
     // Restore original consumers.
     options.programConsumer = programConsumer;
-    options.proguardMapConsumer = proguardMapConsumer;
+    options.mapConsumer = mapConsumer;
   }
 
   public static DexProgramClass run(BisectOptions options) throws Exception {
diff --git a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
index a664d9d..37fad8c 100644
--- a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
@@ -289,7 +289,7 @@
   }
 
   private boolean willComputeProguardMap() {
-    return options.proguardMapConsumer != null;
+    return options.mapConsumer != null;
   }
 
   /** Writer that never needs the input app to deal with mapping info for kotlin. */
diff --git a/src/main/java/com/android/tools/r8/jar/CfApplicationWriter.java b/src/main/java/com/android/tools/r8/jar/CfApplicationWriter.java
index 5c92ac9..19febfd 100644
--- a/src/main/java/com/android/tools/r8/jar/CfApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/jar/CfApplicationWriter.java
@@ -112,7 +112,7 @@
   }
 
   public void write(ClassFileConsumer consumer) {
-    assert options.proguardMapConsumer == null;
+    assert options.mapConsumer == null;
     write(consumer, null);
   }
 
@@ -138,7 +138,7 @@
 
   private void writeApplication(AndroidApp inputApp, ClassFileConsumer consumer) {
     ProguardMapId proguardMapId = null;
-    if (options.proguardMapConsumer != null) {
+    if (options.mapConsumer != null) {
       assert marker.isPresent();
       proguardMapId =
           runAndWriteMap(
diff --git a/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java b/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java
index 7b17f6d..d5fd889 100644
--- a/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java
+++ b/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java
@@ -222,7 +222,7 @@
   private final Map<Signature, Signature> signatureMap = new ConcurrentHashMap<>();
   private final LinkedHashSet<MapVersionMappingInformation> mapVersions;
   private final Map<String, String> originalSourceFiles;
-  private final List<String> preamble;
+  private List<String> preamble;
 
   private ClassNameMapper(
       ImmutableMap<String, ClassNamingForNameMapper> classNameMappings,
@@ -243,6 +243,10 @@
     return preamble;
   }
 
+  public void setPreamble(List<String> preamble) {
+    this.preamble = preamble;
+  }
+
   private Signature canonicalizeSignature(Signature signature) {
     Signature result = signatureMap.get(signature);
     if (result != null) {
diff --git a/src/main/java/com/android/tools/r8/naming/MapConsumer.java b/src/main/java/com/android/tools/r8/naming/MapConsumer.java
new file mode 100644
index 0000000..236c2b7
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/naming/MapConsumer.java
@@ -0,0 +1,20 @@
+// 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.Finishable;
+
+/**
+ * This is an internal consumer that can accept our internal representation of a mapping format.
+ * This should not be exposed.
+ */
+public interface MapConsumer extends Finishable {
+
+  void accept(
+      DiagnosticsHandler diagnosticsHandler,
+      ProguardMapMarkerInfo makerInfo,
+      ClassNameMapper classNameMapper);
+}
diff --git a/src/main/java/com/android/tools/r8/naming/MultiProguardMapConsumer.java b/src/main/java/com/android/tools/r8/naming/MultiProguardMapConsumer.java
deleted file mode 100644
index 5901fd3..0000000
--- a/src/main/java/com/android/tools/r8/naming/MultiProguardMapConsumer.java
+++ /dev/null
@@ -1,47 +0,0 @@
-// 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 java.util.ArrayList;
-import java.util.List;
-
-public class MultiProguardMapConsumer extends ProguardMapConsumer {
-
-  private final List<ProguardMapConsumer> proguardMapConsumers;
-
-  public MultiProguardMapConsumer(List<ProguardMapConsumer> proguardMapConsumers) {
-    this.proguardMapConsumers = proguardMapConsumers;
-  }
-
-  @Override
-  public void finished(DiagnosticsHandler handler) {
-    proguardMapConsumers.forEach(consumer -> consumer.finished(handler));
-  }
-
-  @Override
-  public void accept(ProguardMapMarkerInfo markerInfo, ClassNameMapper classNameMapper) {
-    proguardMapConsumers.forEach(consumer -> consumer.accept(markerInfo, classNameMapper));
-  }
-
-  public static Builder builder() {
-    return new Builder();
-  }
-
-  public static class Builder {
-
-    private final List<ProguardMapConsumer> proguardMapConsumers = new ArrayList<>();
-
-    public Builder addProguardMapConsumer(ProguardMapConsumer consumer) {
-      proguardMapConsumers.add(consumer);
-      return this;
-    }
-
-    public MultiProguardMapConsumer build() {
-      return new MultiProguardMapConsumer(proguardMapConsumers);
-    }
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/naming/ProguardMapPartitionConsumer.java b/src/main/java/com/android/tools/r8/naming/ProguardMapPartitionConsumer.java
deleted file mode 100644
index 451fd6c..0000000
--- a/src/main/java/com/android/tools/r8/naming/ProguardMapPartitionConsumer.java
+++ /dev/null
@@ -1,96 +0,0 @@
-// 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/naming/ProguardMapStringConsumer.java b/src/main/java/com/android/tools/r8/naming/ProguardMapStringConsumer.java
index 95bbd38..be3d3ba 100644
--- a/src/main/java/com/android/tools/r8/naming/ProguardMapStringConsumer.java
+++ b/src/main/java/com/android/tools/r8/naming/ProguardMapStringConsumer.java
@@ -5,36 +5,36 @@
 package com.android.tools.r8.naming;
 
 import com.android.tools.r8.DiagnosticsHandler;
-import com.android.tools.r8.ProguardMapConsumer;
 import com.android.tools.r8.StringConsumer;
 import com.android.tools.r8.utils.ChainableStringConsumer;
 
 /***
- * Default implementation of a ProguardMapConsumer that wraps around a string consumer for streamed
- * string output.
+ * Default implementation of a MapConsumer that wraps around a string consumer for streamed string
+ * output.
  */
-public class ProguardMapStringConsumer extends ProguardMapConsumer
-    implements ChainableStringConsumer {
+public class ProguardMapStringConsumer implements MapConsumer, ChainableStringConsumer {
 
   private final StringConsumer stringConsumer;
-  private final DiagnosticsHandler diagnosticsHandler;
+  private DiagnosticsHandler diagnosticsHandler;
 
-  private ProguardMapStringConsumer(
-      StringConsumer stringConsumer, DiagnosticsHandler diagnosticsHandler) {
+  private ProguardMapStringConsumer(StringConsumer stringConsumer) {
     assert stringConsumer != null;
-    assert diagnosticsHandler != null;
     this.stringConsumer = stringConsumer;
-    this.diagnosticsHandler = diagnosticsHandler;
   }
 
   @Override
-  public void accept(ProguardMapMarkerInfo markerInfo, ClassNameMapper classNameMapper) {
+  public void accept(
+      DiagnosticsHandler diagnosticsHandler,
+      ProguardMapMarkerInfo markerInfo,
+      ClassNameMapper classNameMapper) {
+    this.diagnosticsHandler = diagnosticsHandler;
     accept(markerInfo.serializeToString());
     classNameMapper.write(this);
   }
 
   @Override
   public ChainableStringConsumer accept(String string) {
+    assert diagnosticsHandler != null;
     stringConsumer.accept(string, diagnosticsHandler);
     return this;
   }
@@ -55,20 +55,14 @@
   public static class Builder {
 
     private StringConsumer stringConsumer;
-    private DiagnosticsHandler diagnosticsHandler;
 
     public Builder setStringConsumer(StringConsumer stringConsumer) {
       this.stringConsumer = stringConsumer;
       return this;
     }
 
-    public Builder setDiagnosticsHandler(DiagnosticsHandler diagnosticsHandler) {
-      this.diagnosticsHandler = diagnosticsHandler;
-      return this;
-    }
-
     public ProguardMapStringConsumer build() {
-      return new ProguardMapStringConsumer(stringConsumer, diagnosticsHandler);
+      return new ProguardMapStringConsumer(stringConsumer);
     }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/naming/ProguardMapSupplier.java b/src/main/java/com/android/tools/r8/naming/ProguardMapSupplier.java
index 1d92c11..0637406 100644
--- a/src/main/java/com/android/tools/r8/naming/ProguardMapSupplier.java
+++ b/src/main/java/com/android/tools/r8/naming/ProguardMapSupplier.java
@@ -5,7 +5,6 @@
 
 import com.android.tools.r8.MapIdEnvironment;
 import com.android.tools.r8.MapIdProvider;
-import com.android.tools.r8.ProguardMapConsumer;
 import com.android.tools.r8.dex.Marker.Tool;
 import com.android.tools.r8.utils.ChainableStringConsumer;
 import com.android.tools.r8.utils.ExceptionUtils;
@@ -44,7 +43,7 @@
 
   private final ClassNameMapper classNameMapper;
   private final InternalOptions options;
-  private final ProguardMapConsumer consumer;
+  private final MapConsumer consumer;
   private final Reporter reporter;
   private final Tool compiler;
 
@@ -53,7 +52,7 @@
     this.classNameMapper = classNameMapper.sorted();
     // TODO(b/217111432): Validate Proguard using ProguardMapChecker without building the entire
     //  Proguard map in memory.
-    this.consumer = options.proguardMapConsumer;
+    this.consumer = options.mapConsumer;
     this.options = options;
     this.reporter = options.reporter;
     this.compiler = tool;
@@ -68,6 +67,7 @@
   public ProguardMapId writeProguardMap() {
     ProguardMapId proguardMapId = computeProguardMapId();
     consumer.accept(
+        reporter,
         ProguardMapMarkerInfo.builder()
             .setCompilerName(compiler.name())
             .setProguardMapId(proguardMapId)
diff --git a/src/main/java/com/android/tools/r8/utils/AndroidAppConsumers.java b/src/main/java/com/android/tools/r8/utils/AndroidAppConsumers.java
index 653e567..f010dd8 100644
--- a/src/main/java/com/android/tools/r8/utils/AndroidAppConsumers.java
+++ b/src/main/java/com/android/tools/r8/utils/AndroidAppConsumers.java
@@ -3,6 +3,8 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.utils;
 
+import static com.android.tools.r8.utils.MapConsumerUtils.wrapExistingMapConsumer;
+
 import com.android.tools.r8.BaseCompilerCommand;
 import com.android.tools.r8.ByteDataView;
 import com.android.tools.r8.ClassFileConsumer;
@@ -14,10 +16,9 @@
 import com.android.tools.r8.DexIndexedConsumer.ForwardingConsumer;
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.ProgramConsumer;
-import com.android.tools.r8.ProguardMapConsumer;
 import com.android.tools.r8.ResourceException;
 import com.android.tools.r8.StringConsumer;
-import com.android.tools.r8.naming.MultiProguardMapConsumer;
+import com.android.tools.r8.naming.MapConsumer;
 import com.android.tools.r8.naming.ProguardMapStringConsumer;
 import com.android.tools.r8.origin.Origin;
 import com.google.common.io.ByteStreams;
@@ -36,20 +37,19 @@
   private boolean closed = false;
 
   private ProgramConsumer programConsumer = null;
-  private ProguardMapConsumer proguardMapConsumer = null;
+  private MapConsumer mapConsumer = null;
 
   public AndroidAppConsumers() {
     // Nothing to do.
   }
 
-  public AndroidAppConsumers(BaseCompilerCommand.Builder builder) {
+  public AndroidAppConsumers(BaseCompilerCommand.Builder<?, ?> builder) {
     builder.setProgramConsumer(wrapProgramConsumer(builder.getProgramConsumer()));
   }
 
   public AndroidAppConsumers(InternalOptions options) {
     options.programConsumer = wrapProgramConsumer(options.programConsumer);
-    options.proguardMapConsumer =
-        wrapProguardMapConsumer(options.proguardMapConsumer, options.reporter);
+    options.mapConsumer = wrapProguardMapConsumer(options.mapConsumer);
   }
 
   public ProgramConsumer wrapProgramConsumer(ProgramConsumer consumer) {
@@ -69,39 +69,33 @@
     return programConsumer;
   }
 
-  public ProguardMapConsumer wrapProguardMapConsumer(
-      ProguardMapConsumer consumer, DiagnosticsHandler diagnosticsHandler) {
-    assert proguardMapConsumer == null;
-    if (consumer != null) {
-      proguardMapConsumer =
-          MultiProguardMapConsumer.builder()
-              .addProguardMapConsumer(consumer)
-              .addProguardMapConsumer(
-                  ProguardMapStringConsumer.builder()
-                      .setStringConsumer(
-                          new StringConsumer() {
-                            StringBuilder stringBuilder = null;
+  public MapConsumer wrapProguardMapConsumer(MapConsumer consumer) {
+    assert mapConsumer == null;
+    mapConsumer =
+        wrapExistingMapConsumer(
+            consumer,
+            ProguardMapStringConsumer.builder()
+                .setStringConsumer(
+                    new StringConsumer() {
+                      StringBuilder stringBuilder = null;
 
-                            @Override
-                            public void accept(String string, DiagnosticsHandler handler) {
-                              if (stringBuilder == null) {
-                                stringBuilder = new StringBuilder();
-                              }
-                              stringBuilder.append(string);
-                            }
+                      @Override
+                      public void accept(String string, DiagnosticsHandler handler) {
+                        if (stringBuilder == null) {
+                          stringBuilder = new StringBuilder();
+                        }
+                        stringBuilder.append(string);
+                      }
 
-                            @Override
-                            public void finished(DiagnosticsHandler handler) {
-                              if (stringBuilder != null) {
-                                builder.setProguardMapOutputData(stringBuilder.toString());
-                              }
-                            }
-                          })
-                      .setDiagnosticsHandler(diagnosticsHandler)
-                      .build())
-              .build();
-    }
-    return proguardMapConsumer;
+                      @Override
+                      public void finished(DiagnosticsHandler handler) {
+                        if (stringBuilder != null) {
+                          builder.setProguardMapOutputData(stringBuilder.toString());
+                        }
+                      }
+                    })
+                .build());
+    return mapConsumer;
   }
 
   public DexIndexedConsumer wrapDexIndexedConsumer(DexIndexedConsumer consumer) {
diff --git a/src/main/java/com/android/tools/r8/utils/Box.java b/src/main/java/com/android/tools/r8/utils/Box.java
index 00fde15..621d108 100644
--- a/src/main/java/com/android/tools/r8/utils/Box.java
+++ b/src/main/java/com/android/tools/r8/utils/Box.java
@@ -32,6 +32,13 @@
     return super.computeIfAbsent(supplier);
   }
 
+  public <E extends Exception> T computeIfAbsentThrowing(ThrowingSupplier<T, E> supplier) throws E {
+    if (!isSet()) {
+      set(supplier.get());
+    }
+    return get();
+  }
+
   @Override
   public T get() {
     return super.get();
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index 67f4fda..cd3c1c5 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -17,7 +17,6 @@
 import com.android.tools.r8.GlobalSyntheticsConsumer;
 import com.android.tools.r8.MapIdProvider;
 import com.android.tools.r8.ProgramConsumer;
-import com.android.tools.r8.ProguardMapConsumer;
 import com.android.tools.r8.SourceFileProvider;
 import com.android.tools.r8.StringConsumer;
 import com.android.tools.r8.SyntheticInfoConsumer;
@@ -72,6 +71,7 @@
 import com.android.tools.r8.ir.optimize.Inliner.Reason;
 import com.android.tools.r8.ir.optimize.enums.EnumDataMap;
 import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.naming.MapConsumer;
 import com.android.tools.r8.naming.MapVersion;
 import com.android.tools.r8.optimize.argumentpropagation.ArgumentPropagatorEventConsumer;
 import com.android.tools.r8.optimize.redundantbridgeremoval.RedundantBridgeRemovalOptions;
@@ -1051,7 +1051,7 @@
 
   // If null, no proguard map needs to be computed.
   // If non null it must be and passed to the consumer.
-  public ProguardMapConsumer proguardMapConsumer = null;
+  public MapConsumer mapConsumer = null;
 
   // If null, no usage information needs to be computed.
   // If non-null, it must be and is passed to the consumer.
diff --git a/src/main/java/com/android/tools/r8/utils/MapConsumerUtils.java b/src/main/java/com/android/tools/r8/utils/MapConsumerUtils.java
new file mode 100644
index 0000000..5f8d295
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/utils/MapConsumerUtils.java
@@ -0,0 +1,90 @@
+// 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.utils;
+
+import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.PartitionMapConsumer;
+import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.naming.MapConsumer;
+import com.android.tools.r8.naming.ProguardMapMarkerInfo;
+import com.android.tools.r8.retrace.MappingPartition;
+import com.android.tools.r8.retrace.MappingPartitionMetadata;
+import com.android.tools.r8.utils.ZipUtils.ZipBuilder;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.function.Function;
+
+public class MapConsumerUtils {
+
+  public static MapConsumer wrapExistingMapConsumer(
+      MapConsumer existingMapConsumer, MapConsumer newConsumer) {
+    if (existingMapConsumer == null) {
+      return newConsumer;
+    }
+    return new MapConsumer() {
+      @Override
+      public void accept(
+          DiagnosticsHandler diagnosticsHandler,
+          ProguardMapMarkerInfo makerInfo,
+          ClassNameMapper classNameMapper) {
+        existingMapConsumer.accept(diagnosticsHandler, makerInfo, classNameMapper);
+        newConsumer.accept(diagnosticsHandler, makerInfo, classNameMapper);
+      }
+
+      @Override
+      public void finished(DiagnosticsHandler handler) {
+        existingMapConsumer.finished(handler);
+        newConsumer.finished(handler);
+      }
+    };
+  }
+
+  public static <T> MapConsumer wrapExistingMapConsumerIfNotNull(
+      MapConsumer existingMapConsumer, T object, Function<T, MapConsumer> producer) {
+    if (object == null) {
+      return existingMapConsumer;
+    }
+    return wrapExistingMapConsumer(existingMapConsumer, producer.apply(object));
+  }
+
+  public static PartitionMapConsumer createZipConsumer(Path path) {
+    return new PartitionMapConsumer() {
+
+      private final Box<ZipBuilder> zipBuilderBox = new Box<>();
+
+      @Override
+      public void acceptMappingPartition(MappingPartition mappingPartition) {
+        try {
+          zipBuilderBox
+              .computeIfAbsentThrowing(() -> ZipBuilder.builder(path))
+              .addBytes(mappingPartition.getKey(), mappingPartition.getPayload());
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
+
+      @Override
+      public void acceptMappingPartitionMetadata(
+          MappingPartitionMetadata mappingPartitionMetadata) {
+        try {
+          zipBuilderBox
+              .computeIfAbsentThrowing(() -> ZipBuilder.builder(path))
+              .addBytes("METADATA", mappingPartitionMetadata.getBytes());
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
+
+      @Override
+      public void finished(DiagnosticsHandler handler) {
+        try {
+          zipBuilderBox.computeIfAbsentThrowing(() -> ZipBuilder.builder(path)).build();
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
+    };
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/utils/positions/LineNumberOptimizer.java b/src/main/java/com/android/tools/r8/utils/positions/LineNumberOptimizer.java
index 0181e06..881fbe2 100644
--- a/src/main/java/com/android/tools/r8/utils/positions/LineNumberOptimizer.java
+++ b/src/main/java/com/android/tools/r8/utils/positions/LineNumberOptimizer.java
@@ -46,7 +46,7 @@
       Timing timing,
       OriginalSourceFiles originalSourceFiles,
       DebugRepresentationPredicate representation) {
-    assert appView.options().proguardMapConsumer != null;
+    assert appView.options().mapConsumer != null;
     if (shouldEmitOriginalMappingFile(appView)) {
       appView.options().reporter.warning(new NotSupportedMapVersionForMappingComposeDiagnostic());
       timing.begin("Write proguard map");
diff --git a/src/test/java/com/android/tools/r8/L8CommandTest.java b/src/test/java/com/android/tools/r8/L8CommandTest.java
index 6f5f226..a5ec18b 100644
--- a/src/test/java/com/android/tools/r8/L8CommandTest.java
+++ b/src/test/java/com/android/tools/r8/L8CommandTest.java
@@ -192,9 +192,9 @@
     assertNotNull(parsedCommand.getR8Command());
     InternalOptions internalOptions = parsedCommand.getR8Command().getInternalOptions();
     assertNotNull(internalOptions);
-    assertTrue(internalOptions.proguardMapConsumer instanceof ProguardMapStringConsumer);
+    assertTrue(internalOptions.mapConsumer instanceof ProguardMapStringConsumer);
     ProguardMapStringConsumer mapStringConsumer =
-        (ProguardMapStringConsumer) internalOptions.proguardMapConsumer;
+        (ProguardMapStringConsumer) internalOptions.mapConsumer;
     FileConsumer proguardMapConsumer = (FileConsumer) mapStringConsumer.getStringConsumer();
     assertEquals(pgMap, proguardMapConsumer.getOutputPath());
   }
diff --git a/src/test/java/com/android/tools/r8/R8TestBuilder.java b/src/test/java/com/android/tools/r8/R8TestBuilder.java
index 7f5b220..b9c6c781 100644
--- a/src/test/java/com/android/tools/r8/R8TestBuilder.java
+++ b/src/test/java/com/android/tools/r8/R8TestBuilder.java
@@ -79,6 +79,7 @@
   private final List<Path> mainDexRulesFiles = new ArrayList<>();
   private final List<String> applyMappingMaps = new ArrayList<>();
   private final List<Path> features = new ArrayList<>();
+  private PartitionMapConsumer partitionMapConsumer = null;
 
   @Override
   public boolean isR8TestBuilder() {
@@ -120,6 +121,7 @@
             // Nothing to do.
           }
         });
+    builder.setPartitionMapConsumer(partitionMapConsumer);
 
     if (!applyMappingMaps.isEmpty()) {
       try {
@@ -829,4 +831,9 @@
     getBuilder().setFakeCompilerVersion(version);
     return self();
   }
+
+  public T setPartitionMapConsumer(PartitionMapConsumer partitionMapConsumer) {
+    this.partitionMapConsumer = partitionMapConsumer;
+    return self();
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/canonicalization/ConstClassCanonicalizationMonitorTest.java b/src/test/java/com/android/tools/r8/ir/optimize/canonicalization/ConstClassCanonicalizationMonitorTest.java
index 5c769f0..a1c1662 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/canonicalization/ConstClassCanonicalizationMonitorTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/canonicalization/ConstClassCanonicalizationMonitorTest.java
@@ -42,10 +42,7 @@
         .enableInliningAnnotations()
         .addKeepAttributeLineNumberTable()
         .addKeepAttributeSourceFile()
-        .addOptionsModification(
-            options -> {
-              options.proguardMapConsumer = null;
-            })
+        .addOptionsModification(options -> options.mapConsumer = null)
         .run(parameters.getRuntime(), Main.class)
         .assertSuccessWithOutputLines("Hello World!")
         .inspect(
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
index 211121c..7c64527 100644
--- a/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/partition/R8ZipContainerMappingFileTest.java
@@ -9,14 +9,15 @@
 import static org.junit.Assert.assertTrue;
 
 import com.android.tools.r8.DiagnosticsHandler;
-import com.android.tools.r8.ProguardMapConsumer;
+import com.android.tools.r8.PartitionMapConsumer;
 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.MappingPartition;
+import com.android.tools.r8.retrace.MappingPartitionMetadata;
 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;
@@ -79,8 +80,6 @@
   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)
@@ -88,7 +87,7 @@
             .addKeepMainRule(Main.class)
             .addKeepAttributeSourceFile()
             .addKeepAttributeLineNumberTable()
-            .addOptionsModification(options -> options.proguardMapConsumer = partitionZipConsumer)
+            .setPartitionMapConsumer(createPartitionZipConsumer(pgMapFile))
             .run(parameters.getRuntime(), Main.class)
             .assertFailureWithErrorThatThrows(RuntimeException.class)
             .getOriginalStackTrace();
@@ -99,36 +98,37 @@
         isSame(EXPECTED));
   }
 
-  private ProguardMapConsumer createPartitionZipConsumer(
-      Path pgMapFile, DiagnosticsHandler diagnosticsHandler) throws IOException {
+  private PartitionMapConsumer createPartitionZipConsumer(Path pgMapFile) 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();
+    return new PartitionMapConsumer() {
+      @Override
+      public void acceptMappingPartition(MappingPartition mappingPartition) {
+        try {
+          zipBuilder.addBytes(mappingPartition.getKey(), mappingPartition.getPayload());
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
+
+      @Override
+      public void acceptMappingPartitionMetadata(
+          MappingPartitionMetadata mappingPartitionMetadata) {
+        try {
+          zipBuilder.addBytes("METADATA", mappingPartitionMetadata.getBytes());
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
+
+      @Override
+      public void finished(DiagnosticsHandler handler) {
+        try {
+          zipBuilder.build();
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
+    };
   }
 
   private PartitionMappingSupplier createMappingSupplierFromPartitionZip(Path pgMapFile)
diff --git a/third_party/retrace/binary_compatibility.tar.gz.sha1 b/third_party/retrace/binary_compatibility.tar.gz.sha1
index f4480aa..5678440 100644
--- a/third_party/retrace/binary_compatibility.tar.gz.sha1
+++ b/third_party/retrace/binary_compatibility.tar.gz.sha1
@@ -1 +1 @@
-fa425f3c589eede5fcbe0d2ce5c513b6a71cb4b9
\ No newline at end of file
+5c0b406a52a08af4318a104660091c61d6bb8446
\ No newline at end of file