[Partition] Extend partitioning with strategies

Having different strategies will allow the partitioning api to evolve going forward. The current need is to have all partition keys mapped in the metadata. This will allow us to check that the partitioning has been completed correctly and allow one to reconstruct an ordinary proguard map from a partitioned set.

Bug: b/274735214
Change-Id: I0abe2b17891014c675bc2a83f5d468c8e79787f6
diff --git a/src/main/java/com/android/tools/r8/retrace/RetracePartitionException.java b/src/main/java/com/android/tools/r8/retrace/RetracePartitionException.java
new file mode 100644
index 0000000..7751302
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/retrace/RetracePartitionException.java
@@ -0,0 +1,15 @@
+// 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;
+
+import com.android.tools.r8.Keep;
+
+@Keep
+public class RetracePartitionException extends RuntimeException {
+
+  public RetracePartitionException(String message) {
+    super(message);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionKeyStrategy.java b/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionKeyStrategy.java
index ad53968..1ee7fa3 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionKeyStrategy.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionKeyStrategy.java
@@ -5,11 +5,23 @@
 package com.android.tools.r8.retrace.internal;
 
 public enum MappingPartitionKeyStrategy {
-  OBFUSCATED_TYPE_NAME_AS_KEY(0);
+  OBFUSCATED_TYPE_NAME_AS_KEY(0),
+  OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS(1);
 
-  final int serializedKey;
+  private static final MappingPartitionKeyStrategy DEFAULT_STRATEGY =
+      OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS;
+
+  private final int serializedKey;
 
   MappingPartitionKeyStrategy(int serializedKey) {
     this.serializedKey = serializedKey;
   }
+
+  public int getSerializedKey() {
+    return serializedKey;
+  }
+
+  public static MappingPartitionKeyStrategy getDefaultStrategy() {
+    return DEFAULT_STRATEGY;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionMetadataInternal.java b/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionMetadataInternal.java
index 6c2e3ca..ec9449a 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionMetadataInternal.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/MappingPartitionMetadataInternal.java
@@ -5,11 +5,14 @@
 package com.android.tools.r8.retrace.internal;
 
 import static com.android.tools.r8.retrace.internal.MappingPartitionKeyStrategy.OBFUSCATED_TYPE_NAME_AS_KEY;
+import static com.android.tools.r8.retrace.internal.MappingPartitionKeyStrategy.OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS;
 
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.naming.MapVersion;
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.retrace.MappingPartitionMetadata;
+import com.android.tools.r8.retrace.RetracePartitionException;
+import com.android.tools.r8.retrace.internal.MetadataPartitionCollection.LazyMetadataPartitionCollection;
 import com.android.tools.r8.utils.ExceptionDiagnostic;
 import com.google.common.primitives.Ints;
 import java.io.ByteArrayOutputStream;
@@ -22,26 +25,57 @@
 
   MapVersion getMapVersion();
 
-  byte ZERO_BYTE = (byte) 0;
+  default boolean isObfuscatedTypeNameAsKeyMetadataWithPartitionNames() {
+    return false;
+  }
 
-  static MappingPartitionMetadataInternal createFromBytes(
+  default ObfuscatedTypeNameAsKeyMetadataWithPartitionNames
+      asObfuscatedTypeNameAsKeyMetadataWithPartitionNames() {
+    return null;
+  }
+
+  byte ZERO_BYTE = (byte) 0;
+  // Magic byte put into the metadata
+  byte[] MAGIC = new byte[] {(byte) 0xAA, (byte) 0xA8};
+
+  static int magicOffset() {
+    return MAGIC.length;
+  }
+
+  static MappingPartitionMetadataInternal deserialize(
       byte[] bytes, MapVersion fallBackMapVersion, DiagnosticsHandler diagnosticsHandler) {
     if (bytes == null) {
-      return obfuscatedTypeNameAsKey(fallBackMapVersion);
-    } else if (bytes.length > 2) {
-      int serializedStrategyId = Ints.fromBytes(ZERO_BYTE, ZERO_BYTE, bytes[0], bytes[1]);
-      MapVersion mapVersion = MapVersion.fromName(new String(bytes, 2, bytes.length - 2));
-      if (serializedStrategyId == OBFUSCATED_TYPE_NAME_AS_KEY.serializedKey) {
-        return obfuscatedTypeNameAsKey(mapVersion);
+      return ObfuscatedTypeNameAsKeyMetadata.create(fallBackMapVersion);
+    }
+    if (bytes.length > 2) {
+      if (startsWithMagic(bytes)) {
+        int serializedKey =
+            Ints.fromBytes(ZERO_BYTE, ZERO_BYTE, bytes[magicOffset()], bytes[magicOffset() + 1]);
+        if (serializedKey == OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS.getSerializedKey()) {
+          return ObfuscatedTypeNameAsKeyMetadataWithPartitionNames.deserialize(bytes);
+        }
+      } else if (OBFUSCATED_TYPE_NAME_AS_KEY.getSerializedKey()
+          == Ints.fromBytes(ZERO_BYTE, ZERO_BYTE, bytes[0], bytes[1])) {
+        return ObfuscatedTypeNameAsKeyMetadata.deserialize(bytes);
       }
     }
-    RuntimeException exception = new RuntimeException("Unable to build key strategy from metadata");
+    // If we arrived here then we could not deserialize the metadata.
+    RetracePartitionException exception =
+        new RetracePartitionException("Unknown map partition strategy for metadata");
     diagnosticsHandler.error(new ExceptionDiagnostic(exception));
     throw exception;
   }
 
-  static ObfuscatedTypeNameAsKeyMetadata obfuscatedTypeNameAsKey(MapVersion mapVersion) {
-    return new ObfuscatedTypeNameAsKeyMetadata(mapVersion);
+  private static boolean startsWithMagic(byte[] bytes) {
+    if (bytes.length < MAGIC.length) {
+      return false;
+    }
+    for (int i = 0; i < MAGIC.length; i++) {
+      if (bytes[i] != MAGIC[i]) {
+        return false;
+      }
+    }
+    return true;
   }
 
   class ObfuscatedTypeNameAsKeyMetadata implements MappingPartitionMetadataInternal {
@@ -62,12 +96,14 @@
       return mapVersion;
     }
 
+    // The format is:
+    // <type:short><map-version>
     @Override
     public byte[] getBytes() {
       try {
         ByteArrayOutputStream temp = new ByteArrayOutputStream();
         DataOutputStream dataOutputStream = new DataOutputStream(temp);
-        dataOutputStream.writeShort(OBFUSCATED_TYPE_NAME_AS_KEY.serializedKey);
+        dataOutputStream.writeShort(OBFUSCATED_TYPE_NAME_AS_KEY.getSerializedKey());
         dataOutputStream.writeBytes(mapVersion.getName());
         dataOutputStream.close();
         return temp.toByteArray();
@@ -75,5 +111,88 @@
         throw new RuntimeException(e);
       }
     }
+
+    public static ObfuscatedTypeNameAsKeyMetadata deserialize(byte[] bytes) {
+      MapVersion mapVersion = MapVersion.fromName(new String(bytes, 2, bytes.length - 2));
+      return create(mapVersion);
+    }
+
+    public static ObfuscatedTypeNameAsKeyMetadata create(MapVersion mapVersion) {
+      return new ObfuscatedTypeNameAsKeyMetadata(mapVersion);
+    }
+  }
+
+  class ObfuscatedTypeNameAsKeyMetadataWithPartitionNames
+      implements MappingPartitionMetadataInternal {
+
+    private final MapVersion mapVersion;
+    private final MetadataPartitionCollection metadataPartitionCollection;
+
+    private ObfuscatedTypeNameAsKeyMetadataWithPartitionNames(
+        MapVersion mapVersion, MetadataPartitionCollection metadataPartitionCollection) {
+      this.mapVersion = mapVersion;
+      this.metadataPartitionCollection = metadataPartitionCollection;
+    }
+
+    public static ObfuscatedTypeNameAsKeyMetadataWithPartitionNames create(
+        MapVersion mapVersion, MetadataPartitionCollection metadataPartitionCollection) {
+      return new ObfuscatedTypeNameAsKeyMetadataWithPartitionNames(
+          mapVersion, metadataPartitionCollection);
+    }
+
+    @Override
+    public String getKey(ClassReference classReference) {
+      return classReference.getTypeName();
+    }
+
+    @Override
+    public MapVersion getMapVersion() {
+      return mapVersion;
+    }
+
+    @Override
+    public boolean isObfuscatedTypeNameAsKeyMetadataWithPartitionNames() {
+      return true;
+    }
+
+    @Override
+    public ObfuscatedTypeNameAsKeyMetadataWithPartitionNames
+        asObfuscatedTypeNameAsKeyMetadataWithPartitionNames() {
+      return this;
+    }
+
+    // The format is:
+    // <type:short><map-version-length:short><map-version>[<partition_key>]
+    @Override
+    public byte[] getBytes() {
+      try {
+        ByteArrayOutputStream temp = new ByteArrayOutputStream();
+        DataOutputStream dataOutputStream = new DataOutputStream(temp);
+        dataOutputStream.write(MAGIC);
+        dataOutputStream.writeShort(OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS.getSerializedKey());
+        String name = mapVersion.getName();
+        dataOutputStream.writeShort(name.length());
+        dataOutputStream.writeBytes(name);
+        dataOutputStream.write(metadataPartitionCollection.serialize());
+        dataOutputStream.close();
+        return temp.toByteArray();
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    public static ObfuscatedTypeNameAsKeyMetadataWithPartitionNames deserialize(byte[] bytes) {
+      int start = magicOffset();
+      int length = Ints.fromBytes(ZERO_BYTE, ZERO_BYTE, bytes[start + 2], bytes[start + 3]);
+      MapVersion mapVersion = MapVersion.fromName(new String(bytes, start + 4, length));
+      int partitionCollectionStartIndex = start + 4 + length;
+      return ObfuscatedTypeNameAsKeyMetadataWithPartitionNames.create(
+          mapVersion,
+          new LazyMetadataPartitionCollection(bytes, partitionCollectionStartIndex, bytes.length));
+    }
+
+    public MetadataPartitionCollection getMetadataPartitionCollection() {
+      return metadataPartitionCollection;
+    }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/MetadataPartitionCollection.java b/src/main/java/com/android/tools/r8/retrace/internal/MetadataPartitionCollection.java
new file mode 100644
index 0000000..aef50f0
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/retrace/internal/MetadataPartitionCollection.java
@@ -0,0 +1,72 @@
+// 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.utils.StringUtils;
+import java.nio.charset.StandardCharsets;
+import java.util.Collection;
+import java.util.Collections;
+
+public class MetadataPartitionCollection {
+
+  private static final char SEPARATOR = ';';
+
+  private final Collection<String> partitionKeys;
+
+  private MetadataPartitionCollection(Collection<String> partitionKeys) {
+    this.partitionKeys = partitionKeys;
+  }
+
+  public Collection<String> getPartitionKeys() {
+    return partitionKeys;
+  }
+
+  public byte[] serialize() {
+    return StringUtils.join(SEPARATOR + "", partitionKeys).getBytes(StandardCharsets.UTF_8);
+  }
+
+  public static MetadataPartitionCollection create(Collection<String> partitionKeys) {
+    return new MetadataPartitionCollection(partitionKeys);
+  }
+
+  public static MetadataPartitionCollection createLazy(
+      byte[] bytes, int partitionCollectionStartIndex, int partitionCollectionEndIndex) {
+    return new LazyMetadataPartitionCollection(
+        bytes, partitionCollectionStartIndex, partitionCollectionEndIndex);
+  }
+
+  public static class LazyMetadataPartitionCollection extends MetadataPartitionCollection {
+
+    private final byte[] bytes;
+    private final int partitionCollectionStartIndex;
+    private final int partitionCollectionEndIndex;
+    private MetadataPartitionCollection metadataPartitionCollection = null;
+
+    public LazyMetadataPartitionCollection(
+        byte[] bytes, int partitionCollectionStartIndex, int partitionCollectionEndIndex) {
+      super(Collections.emptyList());
+      this.bytes = bytes;
+      this.partitionCollectionStartIndex = partitionCollectionStartIndex;
+      this.partitionCollectionEndIndex = partitionCollectionEndIndex;
+    }
+
+    @Override
+    public Collection<String> getPartitionKeys() {
+      if (metadataPartitionCollection == null) {
+        metadataPartitionCollection = deserialize();
+      }
+      return metadataPartitionCollection.getPartitionKeys();
+    }
+
+    private MetadataPartitionCollection deserialize() {
+      String allKeys =
+          new String(
+              bytes,
+              partitionCollectionStartIndex,
+              partitionCollectionEndIndex - partitionCollectionStartIndex);
+      return create(StringUtils.split(allKeys, SEPARATOR));
+    }
+  }
+}
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 d37d45c..915fc81 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
@@ -65,7 +65,7 @@
       return mappingPartitionMetadataCache;
     }
     return mappingPartitionMetadataCache =
-        MappingPartitionMetadataInternal.createFromBytes(
+        MappingPartitionMetadataInternal.deserialize(
             metadata, fallbackMapVersion, diagnosticsHandler);
   }
 
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 5fd87ef..98ac266 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
@@ -18,6 +18,9 @@
 import com.android.tools.r8.retrace.ProguardMapPartitioner;
 import com.android.tools.r8.retrace.ProguardMapPartitionerBuilder;
 import com.android.tools.r8.retrace.ProguardMapProducer;
+import com.android.tools.r8.retrace.RetracePartitionException;
+import com.android.tools.r8.retrace.internal.MappingPartitionMetadataInternal.ObfuscatedTypeNameAsKeyMetadata;
+import com.android.tools.r8.retrace.internal.MappingPartitionMetadataInternal.ObfuscatedTypeNameAsKeyMetadataWithPartitionNames;
 import com.android.tools.r8.retrace.internal.ProguardMapReaderWithFiltering.ProguardMapReaderWithFilteringInputBuffer;
 import com.android.tools.r8.retrace.internal.ProguardMapReaderWithFiltering.ProguardMapReaderWithFilteringMappedBuffer;
 import com.android.tools.r8.utils.ExceptionDiagnostic;
@@ -28,6 +31,7 @@
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -42,18 +46,21 @@
   private final DiagnosticsHandler diagnosticsHandler;
   private final boolean allowEmptyMappedRanges;
   private final boolean allowExperimentalMapping;
+  private final MappingPartitionKeyStrategy mappingPartitionKeyStrategy;
 
   private ProguardMapPartitionerOnClassNameToText(
       ProguardMapProducer proguardMapProducer,
       Consumer<MappingPartition> mappingPartitionConsumer,
       DiagnosticsHandler diagnosticsHandler,
       boolean allowEmptyMappedRanges,
-      boolean allowExperimentalMapping) {
+      boolean allowExperimentalMapping,
+      MappingPartitionKeyStrategy mappingPartitionKeyStrategy) {
     this.proguardMapProducer = proguardMapProducer;
     this.mappingPartitionConsumer = mappingPartitionConsumer;
     this.diagnosticsHandler = diagnosticsHandler;
     this.allowEmptyMappedRanges = allowEmptyMappedRanges;
     this.allowExperimentalMapping = allowExperimentalMapping;
+    this.mappingPartitionKeyStrategy = mappingPartitionKeyStrategy;
   }
 
   @Override
@@ -70,6 +77,7 @@
     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.forEachClassMapping(
         (classMapping, entries) -> {
@@ -100,6 +108,7 @@
                 new MappingPartitionImpl(
                     currentClassMapping.getKey(),
                     payloadWithClassReferences.toString().getBytes(StandardCharsets.UTF_8)));
+            keys.add(currentClassMapping.getKey());
           } catch (IOException e) {
             diagnosticsHandler.error(new ExceptionDiagnostic(e));
           }
@@ -109,7 +118,18 @@
     if (mapVersionInfo != null) {
       mapVersion = mapVersionInfo.getMapVersion();
     }
-    return MappingPartitionMetadataInternal.obfuscatedTypeNameAsKey(mapVersion);
+    if (mappingPartitionKeyStrategy == MappingPartitionKeyStrategy.OBFUSCATED_TYPE_NAME_AS_KEY) {
+      return ObfuscatedTypeNameAsKeyMetadata.create(mapVersion);
+    } else if (mappingPartitionKeyStrategy
+        == MappingPartitionKeyStrategy.OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS) {
+      return ObfuscatedTypeNameAsKeyMetadataWithPartitionNames.create(
+          mapVersion, MetadataPartitionCollection.create(keys));
+    } else {
+      RetracePartitionException retraceError =
+          new RetracePartitionException("Unknown mapping partitioning strategy");
+      diagnosticsHandler.error(new ExceptionDiagnostic(retraceError));
+      throw retraceError;
+    }
   }
 
   private String getSourceFileMapping(String className, String sourceFile) {
@@ -129,13 +149,11 @@
 
     private final ProguardMapReaderWithFiltering lineReader;
     private final Map<String, List<String>> readSections = new LinkedHashMap<>();
-    private final List<String> preamble;
     private List<String> currentList;
 
     public PartitionLineReader(ProguardMapReaderWithFiltering lineReader) {
       this.lineReader = lineReader;
       currentList = new ArrayList<>();
-      preamble = currentList;
     }
 
     @Override
@@ -166,11 +184,11 @@
       implements ProguardMapPartitionerBuilder<
           ProguardMapPartitionerBuilderImpl, ProguardMapPartitionerOnClassNameToText> {
 
-    private ProguardMapProducer proguardMapProducer;
-    private Consumer<MappingPartition> mappingPartitionConsumer;
-    private final DiagnosticsHandler diagnosticsHandler;
-    private boolean allowEmptyMappedRanges = false;
-    private boolean allowExperimentalMapping = false;
+    protected ProguardMapProducer proguardMapProducer;
+    protected Consumer<MappingPartition> mappingPartitionConsumer;
+    protected final DiagnosticsHandler diagnosticsHandler;
+    protected boolean allowEmptyMappedRanges = false;
+    protected boolean allowExperimentalMapping = false;
 
     public ProguardMapPartitionerBuilderImpl(DiagnosticsHandler diagnosticsHandler) {
       this.diagnosticsHandler = diagnosticsHandler;
@@ -211,7 +229,38 @@
           mappingPartitionConsumer,
           diagnosticsHandler,
           allowEmptyMappedRanges,
-          allowExperimentalMapping);
+          allowExperimentalMapping,
+          MappingPartitionKeyStrategy.getDefaultStrategy());
+    }
+  }
+
+  // This class should not be exposed to clients and is only used from tests to control the
+  // partitioning strategy.
+  public static class ProguardMapPartitionerBuilderImplInternal
+      extends ProguardMapPartitionerBuilderImpl {
+
+    private MappingPartitionKeyStrategy mappingPartitionKeyStrategy =
+        MappingPartitionKeyStrategy.OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS;
+
+    public ProguardMapPartitionerBuilderImplInternal(DiagnosticsHandler diagnosticsHandler) {
+      super(diagnosticsHandler);
+    }
+
+    public ProguardMapPartitionerBuilderImplInternal setMappingPartitionKeyStrategy(
+        MappingPartitionKeyStrategy mappingPartitionKeyStrategy) {
+      this.mappingPartitionKeyStrategy = mappingPartitionKeyStrategy;
+      return this;
+    }
+
+    @Override
+    public ProguardMapPartitionerOnClassNameToText build() {
+      return new ProguardMapPartitionerOnClassNameToText(
+          proguardMapProducer,
+          mappingPartitionConsumer,
+          diagnosticsHandler,
+          allowEmptyMappedRanges,
+          allowExperimentalMapping,
+          mappingPartitionKeyStrategy);
     }
   }
 }
diff --git a/src/test/java/com/android/tools/r8/retrace/partition/RetracePartitionMetadataPartitionNamesTest.java b/src/test/java/com/android/tools/r8/retrace/partition/RetracePartitionMetadataPartitionNamesTest.java
new file mode 100644
index 0000000..8d40215
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/partition/RetracePartitionMetadataPartitionNamesTest.java
@@ -0,0 +1,94 @@
+// 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 org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import com.android.tools.r8.naming.MapVersion;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.retrace.MappingPartitionMetadata;
+import com.android.tools.r8.retrace.ProguardMapProducer;
+import com.android.tools.r8.retrace.internal.MappingPartitionKeyStrategy;
+import com.android.tools.r8.retrace.internal.MappingPartitionMetadataInternal;
+import com.android.tools.r8.retrace.internal.MappingPartitionMetadataInternal.ObfuscatedTypeNameAsKeyMetadataWithPartitionNames;
+import com.android.tools.r8.retrace.internal.ProguardMapPartitionerOnClassNameToText.ProguardMapPartitionerBuilderImplInternal;
+import com.android.tools.r8.utils.StringUtils;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class RetracePartitionMetadataPartitionNamesTest extends TestBase {
+
+  @Parameters(name = "{0}")
+  public static TestParametersCollection data() {
+    return getTestParameters().withNoneRuntime().build();
+  }
+
+  public RetracePartitionMetadataPartitionNamesTest(TestParameters parameters) {
+    parameters.assertNoneRuntime();
+  }
+
+  public ClassReference obfuscatedClass1 = Reference.classFromTypeName("a.a");
+  public ClassReference obfuscatedClass2 = Reference.classFromTypeName("b.b");
+  public ClassReference obfuscatedClass3 = Reference.classFromTypeName("c.c");
+
+  public String mapping =
+      StringUtils.unixLines(
+          "# { id: 'com.android.tools.r8.mapping', version: '2.0' }",
+          "some.class1 -> " + obfuscatedClass1.getTypeName() + ":",
+          "  void field -> a",
+          "some.class2 -> " + obfuscatedClass2.getTypeName() + ":",
+          "  void field -> a",
+          "some.class2 -> " + obfuscatedClass3.getTypeName() + ":",
+          "  void field -> a");
+
+  @Test
+  public void test() throws Exception {
+    List<String> expectedPartitionKeys = new ArrayList<>();
+    expectedPartitionKeys.add(obfuscatedClass1.getTypeName());
+    expectedPartitionKeys.add(obfuscatedClass2.getTypeName());
+    expectedPartitionKeys.add(obfuscatedClass3.getTypeName());
+
+    ProguardMapProducer proguardMapProducer = ProguardMapProducer.fromString(mapping);
+    DiagnosticsHandler diagnosticsHandler = new DiagnosticsHandler() {};
+    Map<String, byte[]> partitions = new HashMap<>();
+    MappingPartitionMetadata metadataData =
+        new ProguardMapPartitionerBuilderImplInternal(diagnosticsHandler)
+            .setMappingPartitionKeyStrategy(
+                MappingPartitionKeyStrategy.OBFUSCATED_TYPE_NAME_AS_KEY_WITH_PARTITIONS)
+            .setProguardMapProducer(proguardMapProducer)
+            .setPartitionConsumer(
+                partition -> partitions.put(partition.getKey(), partition.getPayload()))
+            .build()
+            .run();
+    assertNotNull(metadataData);
+    assertEquals(new HashSet<>(expectedPartitionKeys), partitions.keySet());
+
+    byte[] bytes = metadataData.getBytes();
+    MappingPartitionMetadataInternal mappingPartitionMetadata =
+        MappingPartitionMetadataInternal.deserialize(
+            bytes, MapVersion.MAP_VERSION_NONE, diagnosticsHandler);
+    assertTrue(mappingPartitionMetadata.isObfuscatedTypeNameAsKeyMetadataWithPartitionNames());
+    ObfuscatedTypeNameAsKeyMetadataWithPartitionNames obfuscatedTypeNameMetadata =
+        mappingPartitionMetadata.asObfuscatedTypeNameAsKeyMetadataWithPartitionNames();
+    assertEquals(
+        expectedPartitionKeys,
+        obfuscatedTypeNameMetadata.getMetadataPartitionCollection().getPartitionKeys());
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/retrace/partition/RetracePartitionMetadataUnknownTest.java b/src/test/java/com/android/tools/r8/retrace/partition/RetracePartitionMetadataUnknownTest.java
new file mode 100644
index 0000000..345566b
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/partition/RetracePartitionMetadataUnknownTest.java
@@ -0,0 +1,53 @@
+// 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 org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import com.android.tools.r8.naming.MapVersion;
+import com.android.tools.r8.retrace.RetracePartitionException;
+import com.android.tools.r8.retrace.internal.MappingPartitionMetadataInternal;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class RetracePartitionMetadataUnknownTest extends TestBase {
+
+  @Parameters(name = "{0}")
+  public static TestParametersCollection data() {
+    return getTestParameters().withNoneRuntime().build();
+  }
+
+  public RetracePartitionMetadataUnknownTest(TestParameters parameters) {
+    parameters.assertNoneRuntime();
+  }
+
+  @Test
+  public void test() throws Exception {
+    ByteArrayOutputStream temp = new ByteArrayOutputStream();
+    DataOutputStream dataOutputStream = new DataOutputStream(temp);
+    dataOutputStream.writeShort(Short.MAX_VALUE);
+    dataOutputStream.close();
+    byte[] bytes = temp.toByteArray();
+    DiagnosticsHandler diagnosticsHandler = new DiagnosticsHandler() {};
+    RetracePartitionException retracePartitionException =
+        assertThrows(
+            RetracePartitionException.class,
+            () ->
+                MappingPartitionMetadataInternal.deserialize(
+                    bytes, MapVersion.MAP_VERSION_NONE, diagnosticsHandler));
+    assertEquals(
+        "Unknown map partition strategy for metadata", retracePartitionException.getMessage());
+  }
+}