[Retrace] Add inlined source file names to partitions

Bug: b/211603371
Change-Id: I500e60409b70663ab5c65e489d58c0d3fe04b4ae
diff --git a/src/main/java/com/android/tools/r8/naming/ClassNamingForNameMapper.java b/src/main/java/com/android/tools/r8/naming/ClassNamingForNameMapper.java
index 352c555..aae3f42 100644
--- a/src/main/java/com/android/tools/r8/naming/ClassNamingForNameMapper.java
+++ b/src/main/java/com/android/tools/r8/naming/ClassNamingForNameMapper.java
@@ -365,6 +365,30 @@
     return fieldMembers.values();
   }
 
+  public void visitAllFullyQualifiedReferences(Consumer<String> consumer) {
+    mappedFieldNamingsByName
+        .values()
+        .forEach(
+            mappedFields ->
+                mappedFields.forEach(
+                    mappedField -> {
+                      if (mappedField.renamedSignature.isQualified()) {
+                        consumer.accept(mappedField.renamedSignature.toHolderFromQualified());
+                      }
+                    }));
+    mappedRangesByRenamedName
+        .values()
+        .forEach(
+            mappedRanges -> {
+              mappedRanges.mappedRanges.forEach(
+                  mappedRange -> {
+                    if (mappedRange.signature.isQualified()) {
+                      consumer.accept(mappedRange.signature.toHolderFromQualified());
+                    }
+                  });
+            });
+  }
+
   @Override
   public <T extends Throwable> void forAllMethodNaming(
       ThrowingConsumer<MemberNaming, T> consumer) throws T {
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 dc0f565..7ef3d1c 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
@@ -8,8 +8,10 @@
 
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.naming.ClassNamingForNameMapper;
 import com.android.tools.r8.naming.LineReader;
 import com.android.tools.r8.naming.MapVersion;
+import com.android.tools.r8.naming.mappinginformation.FileNameInformation;
 import com.android.tools.r8.naming.mappinginformation.MapVersionMappingInformation;
 import com.android.tools.r8.retrace.MappingPartition;
 import com.android.tools.r8.retrace.MappingPartitionMetadata;
@@ -24,9 +26,12 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 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;
 
@@ -63,18 +68,30 @@
     reader.forEachClassMapping(
         (classMapping, entries) -> {
           try {
-            String payLoad = StringUtils.joinLines(entries);
-            ClassNameMapper classNameMapper = ClassNameMapper.mapperFromString(payLoad);
+            String payload = StringUtils.joinLines(entries);
+            ClassNameMapper classNameMapper = ClassNameMapper.mapperFromString(payload);
             if (classNameMapper.getClassNameMappings().size() != 1) {
               diagnosticsHandler.error(
-                  new StringDiagnostic("Multiple class names in payload\n: " + payLoad));
+                  new StringDiagnostic("Multiple class names in payload\n: " + payload));
               return;
             }
-            String obfuscatedName =
-                classNameMapper.getClassNameMappings().keySet().iterator().next();
-            // TODO(b/211603371): Handle inline source file names by adding to partition.
+            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(obfuscatedName, payLoad.getBytes(StandardCharsets.UTF_8)));
+                new MappingPartitionImpl(
+                    currentClassMapping.getKey(),
+                    payloadWithClassReferences.toString().getBytes(StandardCharsets.UTF_8)));
           } catch (IOException e) {
             diagnosticsHandler.error(new ExceptionDiagnostic(e));
           }
@@ -87,6 +104,19 @@
     return MappingPartitionMetadataInternal.obfuscatedTypeNameAsKey(mapVersion);
   }
 
+  private String getSourceFileMapping(String className, String sourceFile) {
+    if (sourceFile == null) {
+      return "";
+    }
+    return className
+        + " -> "
+        + className
+        + ":"
+        + "\n  # "
+        + FileNameInformation.build(sourceFile).serialize()
+        + "\n";
+  }
+
   public static class PartitionLineReader implements LineReader {
 
     private final ProguardMapReaderWithFiltering lineReader;
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java
index ffb522b..e1e73b8 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java
@@ -42,7 +42,8 @@
           RetracePartitionStringTest.ApiTest.class,
           RetracePartitionRoundTripTest.ApiTest.class,
           RetracePartitionJoinNoMetadataTest.ApiTest.class,
-          RetracePartitionSerializedObfuscatedKeyTest.ApiTest.class);
+          RetracePartitionSerializedObfuscatedKeyTest.ApiTest.class,
+          RetracePartitionRoundTripInlineTest.ApiTest.class);
 
   public static List<Class<? extends RetraceApiBinaryTest>> CLASSES_PENDING_BINARY_COMPATIBILITY =
       ImmutableList.of();
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetracePartitionRoundTripInlineTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetracePartitionRoundTripInlineTest.java
new file mode 100644
index 0000000..3bc2eb2
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetracePartitionRoundTripInlineTest.java
@@ -0,0 +1,117 @@
+// Copyright (c) 2022, 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.api;
+
+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.TestParameters;
+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.PartitionMappingSupplier;
+import com.android.tools.r8.retrace.ProguardMapPartitioner;
+import com.android.tools.r8.retrace.ProguardMapProducer;
+import com.android.tools.r8.retrace.RetraceFrameElement;
+import com.android.tools.r8.retrace.RetraceStackTraceContext;
+import com.android.tools.r8.retrace.Retracer;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalInt;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class RetracePartitionRoundTripInlineTest extends RetraceApiTestBase {
+
+  public RetracePartitionRoundTripInlineTest(TestParameters parameters) {
+    super(parameters);
+  }
+
+  @Override
+  protected Class<? extends RetraceApiBinaryTest> binaryTestClass() {
+    return ApiTest.class;
+  }
+
+  public static class ApiTest implements RetraceApiBinaryTest {
+
+    private final ClassReference callerOriginal = Reference.classFromTypeName("some.Class");
+    private final ClassReference callerRenamed = Reference.classFromTypeName("b");
+
+    private final String mapping =
+        "# { id: 'com.android.tools.r8.mapping', version: '2.0' }\n"
+            + "inlinee.Class -> inlinee.Class:\n"
+            + " # {'id':'sourceFile','fileName':'InlineeClass.kt'}\n"
+            + callerOriginal.getTypeName()
+            + " -> "
+            + callerRenamed.getTypeName()
+            + ":\n"
+            + " # {'id':'sourceFile','fileName':'CallerClass.kt'}\n"
+            + "  1:1:void inlinee.Class.bar():42:42 -> a\n"
+            + "  1:1:void foo():43 -> a\n";
+
+    private int prepareCounter = 0;
+
+    @Test
+    public void test() throws IOException {
+      ProguardMapProducer proguardMapProducer = ProguardMapProducer.fromString(mapping);
+      Map<String, byte[]> partitions = new HashMap<>();
+      MappingPartitionMetadata metadataData =
+          ProguardMapPartitioner.builder(new DiagnosticsHandler() {})
+              .setProguardMapProducer(proguardMapProducer)
+              .setPartitionConsumer(
+                  partition -> partitions.put(partition.getKey(), partition.getPayload()))
+              .build()
+              .run();
+      assertNotNull(metadataData);
+      assertEquals(2, partitions.size());
+
+      Set<String> preFetchedKeys = new LinkedHashSet<>();
+      PartitionMappingSupplier mappingSupplier =
+          PartitionMappingSupplier.builder()
+              .setMetadata(metadataData.getBytes())
+              .setRegisterMappingPartitionCallback(preFetchedKeys::add)
+              .setPrepareMappingPartitionsCallback(() -> prepareCounter++)
+              .setMappingPartitionFromKeySupplier(
+                  key -> {
+                    assertTrue(preFetchedKeys.contains(key));
+                    assertTrue(partitions.containsKey(key));
+                    return partitions.get(key);
+                  })
+              .build();
+      assertEquals(0, prepareCounter);
+      Retracer retracer = Retracer.builder().setMappingSupplier(mappingSupplier).build();
+      List<RetraceFrameElement> callerRetraced =
+          retracer
+              .retraceFrame(
+                  RetraceStackTraceContext.empty(),
+                  OptionalInt.of(1),
+                  Reference.methodFromDescriptor(callerRenamed, "a", "()V"))
+              .stream()
+              .collect(Collectors.toList());
+      // The retrace result should not be ambiguous or empty.
+      assertEquals(1, callerRetraced.size());
+      assertEquals(1, preFetchedKeys.size());
+      assertEquals(1, prepareCounter);
+      RetraceFrameElement retraceFrameElement = callerRetraced.get(0);
+
+      // Check that visiting all frames report all source files.
+      List<String> allSourceFiles =
+          retraceFrameElement.stream()
+              .map(x -> x.getSourceFile().getOrInferSourceFile())
+              .collect(Collectors.toList());
+      assertEquals(Arrays.asList("InlineeClass.kt", "CallerClass.kt"), allSourceFiles);
+    }
+  }
+}
diff --git a/third_party/retrace/binary_compatibility.tar.gz.sha1 b/third_party/retrace/binary_compatibility.tar.gz.sha1
index 0ebc3f5..da20d2e 100644
--- a/third_party/retrace/binary_compatibility.tar.gz.sha1
+++ b/third_party/retrace/binary_compatibility.tar.gz.sha1
@@ -1 +1 @@
-f49f174fe4f8ecf0ddc4b47d98bcb293d689eeb9
\ No newline at end of file
+4408d86adfb21f32c2dcb93c2c8f2e392f0e0de0
\ No newline at end of file