Reland "Include art profiles for rewriting in dumps"

This reverts commit 63d40ecbafa356ccb3c228cf8f564202c275711b.

Change-Id: Iab426c4ccbe08f429fb5f5dcba16c54fd84e30a2
diff --git a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
index 4d1269c..f251ecb 100644
--- a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
+++ b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
@@ -21,6 +21,7 @@
 import com.android.tools.r8.utils.DumpInputFlags;
 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.ProgramConsumerUtils;
 import com.android.tools.r8.utils.Reporter;
 import com.android.tools.r8.utils.ThreadUtils;
@@ -154,6 +155,9 @@
         .setOptimizeMultidexForLinearAlloc(isOptimizeMultidexForLinearAlloc())
         .setThreadCount(getThreadCount())
         .setDesugarState(getDesugarState())
+        .setArtProfileProviders(
+            ListUtils.map(
+                getArtProfilesForRewriting(), ArtProfileForRewriting::getArtProfileProvider))
         .setStartupProfileProviders(getStartupProfileProviders());
     if (getAndroidPlatformBuild()) {
       builder.setAndroidPlatformBuild(true);
diff --git a/src/main/java/com/android/tools/r8/L8Command.java b/src/main/java/com/android/tools/r8/L8Command.java
index ebe84b2..6bac5df 100644
--- a/src/main/java/com/android/tools/r8/L8Command.java
+++ b/src/main/java/com/android/tools/r8/L8Command.java
@@ -164,7 +164,13 @@
 
   @Override
   List<ArtProfileForRewriting> getArtProfilesForRewriting() {
-    return getD8Command().getArtProfilesForRewriting();
+    if (getD8Command() != null) {
+      return getD8Command().getArtProfilesForRewriting();
+    }
+    if (getR8Command() != null) {
+      return getR8Command().getArtProfilesForRewriting();
+    }
+    return Collections.emptyList();
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/dump/DumpOptions.java b/src/main/java/com/android/tools/r8/dump/DumpOptions.java
index 2ecf1d3..0c538fa 100644
--- a/src/main/java/com/android/tools/r8/dump/DumpOptions.java
+++ b/src/main/java/com/android/tools/r8/dump/DumpOptions.java
@@ -9,6 +9,7 @@
 import com.android.tools.r8.dex.Marker.Tool;
 import com.android.tools.r8.features.FeatureSplitConfiguration;
 import com.android.tools.r8.ir.desugar.desugaredlibrary.DesugaredLibrarySpecification;
+import com.android.tools.r8.profile.art.ArtProfileProvider;
 import com.android.tools.r8.shaking.ProguardConfiguration;
 import com.android.tools.r8.shaking.ProguardConfigurationRule;
 import com.android.tools.r8.startup.StartupProfileProvider;
@@ -69,6 +70,7 @@
   private final FeatureSplitConfiguration featureSplitConfiguration;
   private final ProguardConfiguration proguardConfiguration;
   private final List<ProguardConfigurationRule> mainDexKeepRules;
+  private final Collection<ArtProfileProvider> artProfileProviders;
   private final Collection<StartupProfileProvider> startupProfileProviders;
   private final boolean enableMissingLibraryApiModeling;
   private final boolean isAndroidPlatformBuild;
@@ -98,6 +100,7 @@
       FeatureSplitConfiguration featureSplitConfiguration,
       ProguardConfiguration proguardConfiguration,
       List<ProguardConfigurationRule> mainDexKeepRules,
+      Collection<ArtProfileProvider> artProfileProviders,
       Collection<StartupProfileProvider> startupProfileProviders,
       boolean enableMissingLibraryApiModeling,
       boolean isAndroidPlatformBuild,
@@ -120,6 +123,7 @@
     this.featureSplitConfiguration = featureSplitConfiguration;
     this.proguardConfiguration = proguardConfiguration;
     this.mainDexKeepRules = mainDexKeepRules;
+    this.artProfileProviders = artProfileProviders;
     this.startupProfileProviders = startupProfileProviders;
     this.enableMissingLibraryApiModeling = enableMissingLibraryApiModeling;
     this.isAndroidPlatformBuild = isAndroidPlatformBuild;
@@ -296,6 +300,14 @@
     return mainDexKeepRules;
   }
 
+  public boolean hasArtProfileProviders() {
+    return artProfileProviders != null && !artProfileProviders.isEmpty();
+  }
+
+  public Collection<ArtProfileProvider> getArtProfileProviders() {
+    return artProfileProviders;
+  }
+
   public boolean hasStartupProfileProviders() {
     return startupProfileProviders != null && !startupProfileProviders.isEmpty();
   }
@@ -331,6 +343,7 @@
     private FeatureSplitConfiguration featureSplitConfiguration;
     private ProguardConfiguration proguardConfiguration;
     private List<ProguardConfigurationRule> mainDexKeepRules;
+    private Collection<ArtProfileProvider> artProfileProviders;
     private Collection<StartupProfileProvider> startupProfileProviders;
 
     private boolean enableMissingLibraryApiModeling = false;
@@ -437,6 +450,11 @@
       return this;
     }
 
+    public Builder setArtProfileProviders(Collection<ArtProfileProvider> artProfileProviders) {
+      this.artProfileProviders = artProfileProviders;
+      return this;
+    }
+
     public Builder setStartupProfileProviders(
         Collection<StartupProfileProvider> startupProfileProviders) {
       this.startupProfileProviders = startupProfileProviders;
@@ -497,6 +515,7 @@
           featureSplitConfiguration,
           proguardConfiguration,
           mainDexKeepRules,
+          artProfileProviders,
           startupProfileProviders,
           enableMissingLibraryApiModeling,
           isAndroidPlatformBuild,
diff --git a/src/main/java/com/android/tools/r8/profile/art/ArtProfileProviderUtils.java b/src/main/java/com/android/tools/r8/profile/art/ArtProfileProviderUtils.java
index bc3cbb5..e09b79c 100644
--- a/src/main/java/com/android/tools/r8/profile/art/ArtProfileProviderUtils.java
+++ b/src/main/java/com/android/tools/r8/profile/art/ArtProfileProviderUtils.java
@@ -6,12 +6,23 @@
 
 import static com.android.tools.r8.utils.ConsumerUtils.emptyConsumer;
 
+import com.android.tools.r8.TextInputStream;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.origin.PathOrigin;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.MethodReference;
+import com.android.tools.r8.utils.Box;
+import com.android.tools.r8.utils.ClassReferenceUtils;
+import com.android.tools.r8.utils.MethodReferenceUtils;
 import com.android.tools.r8.utils.UTF8TextInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
 import java.io.UncheckedIOException;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
+import java.util.function.Consumer;
 
 public class ArtProfileProviderUtils {
 
@@ -33,4 +44,101 @@
       }
     };
   }
+
+  /** Serialize the given {@param artProfileProvider} to a string for writing it to a dump. */
+  public static String serializeToString(ArtProfileProvider artProfileProvider) throws IOException {
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    try (OutputStreamWriter outputStreamWriter =
+        new OutputStreamWriter(baos, StandardCharsets.UTF_8)) {
+      artProfileProvider.getArtProfile(
+          new ArtProfileBuilder() {
+
+            @Override
+            public ArtProfileBuilder addClassRule(
+                Consumer<ArtProfileClassRuleBuilder> classRuleBuilderConsumer) {
+              classRuleBuilderConsumer.accept(
+                  new ArtProfileClassRuleBuilder() {
+
+                    @Override
+                    public ArtProfileClassRuleBuilder setClassReference(
+                        ClassReference classReference) {
+                      writeLine(
+                          outputStreamWriter, ClassReferenceUtils.toSmaliString(classReference));
+                      return this;
+                    }
+                  });
+              return this;
+            }
+
+            @Override
+            public ArtProfileBuilder addMethodRule(
+                Consumer<ArtProfileMethodRuleBuilder> methodRuleBuilderConsumer) {
+              Box<MethodReference> methodReferenceBox = new Box<>();
+              methodRuleBuilderConsumer.accept(
+                  new ArtProfileMethodRuleBuilder() {
+
+                    @Override
+                    public ArtProfileMethodRuleBuilder setMethodReference(
+                        MethodReference methodReference) {
+                      methodReferenceBox.set(methodReference);
+                      return this;
+                    }
+
+                    @Override
+                    public ArtProfileMethodRuleBuilder setMethodRuleInfo(
+                        Consumer<ArtProfileMethodRuleInfoBuilder> methodRuleInfoBuilderConsumer) {
+                      ArtProfileMethodRuleInfoImpl.Builder artProfileMethodRuleInfoBuilder =
+                          ArtProfileMethodRuleInfoImpl.builder();
+                      methodRuleInfoBuilderConsumer.accept(artProfileMethodRuleInfoBuilder);
+                      ArtProfileMethodRuleInfoImpl artProfileMethodRuleInfo =
+                          artProfileMethodRuleInfoBuilder.build();
+                      try {
+                        artProfileMethodRuleInfo.writeHumanReadableFlags(outputStreamWriter);
+                      } catch (IOException e) {
+                        throw new UncheckedIOException(e);
+                      }
+                      return this;
+                    }
+                  });
+              writeLine(
+                  outputStreamWriter, MethodReferenceUtils.toSmaliString(methodReferenceBox.get()));
+              return this;
+            }
+
+            @Override
+            public ArtProfileBuilder addHumanReadableArtProfile(
+                TextInputStream textInputStream,
+                Consumer<HumanReadableArtProfileParserBuilder> parserBuilderConsumer) {
+              try (InputStreamReader inputStreamReader =
+                  new InputStreamReader(
+                      textInputStream.getInputStream(), textInputStream.getCharset())) {
+                char[] buffer = new char[1024];
+                int len = inputStreamReader.read(buffer);
+                while (len != -1) {
+                  outputStreamWriter.write(buffer, 0, len);
+                  len = inputStreamReader.read(buffer);
+                }
+                writeLine(outputStreamWriter);
+              } catch (IOException e) {
+                throw new UncheckedIOException(e);
+              }
+              return this;
+            }
+          });
+    }
+    return baos.toString();
+  }
+
+  private static void writeLine(OutputStreamWriter outputStreamWriter) {
+    writeLine(outputStreamWriter, "");
+  }
+
+  private static void writeLine(OutputStreamWriter outputStreamWriter, String string) {
+    try {
+      outputStreamWriter.write(string);
+      outputStreamWriter.write('\n');
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/utils/AndroidApp.java b/src/main/java/com/android/tools/r8/utils/AndroidApp.java
index 8ea5c72..cf3047d 100644
--- a/src/main/java/com/android/tools/r8/utils/AndroidApp.java
+++ b/src/main/java/com/android/tools/r8/utils/AndroidApp.java
@@ -42,6 +42,8 @@
 import com.android.tools.r8.origin.ArchiveEntryOrigin;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.origin.PathOrigin;
+import com.android.tools.r8.profile.art.ArtProfileProvider;
+import com.android.tools.r8.profile.art.ArtProfileProviderUtils;
 import com.android.tools.r8.shaking.FilteredClassPath;
 import com.android.tools.r8.startup.StartupProfileProvider;
 import com.android.tools.r8.synthesis.SyntheticItems;
@@ -539,6 +541,9 @@
             StringUtils.joinLines(dumpOptions.getMainDexKeepRules()).getBytes(),
             ZipEntry.DEFLATED);
       }
+      if (dumpOptions.hasArtProfileProviders()) {
+        dumpArtProfileProviders(dumpOptions.getArtProfileProviders(), options, out);
+      }
       if (dumpOptions.hasStartupProfileProviders()) {
         dumpStartupProfileProviders(dumpOptions.getStartupProfileProviders(), options, out);
       }
@@ -571,6 +576,23 @@
     return nextDexIndex;
   }
 
+  private void dumpArtProfileProviders(
+      Collection<ArtProfileProvider> artProfileProviders,
+      InternalOptions options,
+      ZipOutputStream out)
+      throws IOException {
+    int artProfileProviderIndex = 1;
+    for (ArtProfileProvider artProfileProvider : artProfileProviders) {
+      String artProfileFileName = "art-profile-" + artProfileProviderIndex + ".txt";
+      writeToZipStream(
+          out,
+          artProfileFileName,
+          ArtProfileProviderUtils.serializeToString(artProfileProvider).getBytes(),
+          ZipEntry.DEFLATED);
+      artProfileProviderIndex++;
+    }
+  }
+
   private void dumpStartupProfileProviders(
       Collection<StartupProfileProvider> startupProfileProviders,
       InternalOptions options,
@@ -578,11 +600,10 @@
       throws IOException {
     int startupProfileProviderIndex = 1;
     for (StartupProfileProvider startupProfileProvider : startupProfileProviders) {
-      String startupProfileProviderFileName =
-          "startup-profile-" + startupProfileProviderIndex + ".txt";
+      String startupProfileFileName = "startup-profile-" + startupProfileProviderIndex + ".txt";
       writeToZipStream(
           out,
-          startupProfileProviderFileName,
+          startupProfileFileName,
           StartupProfileProviderUtils.serializeToString(options, startupProfileProvider).getBytes(),
           ZipEntry.DEFLATED);
       startupProfileProviderIndex++;
diff --git a/src/main/java/com/android/tools/r8/utils/UTF8TextInputStream.java b/src/main/java/com/android/tools/r8/utils/UTF8TextInputStream.java
index 20252e7..709df16 100644
--- a/src/main/java/com/android/tools/r8/utils/UTF8TextInputStream.java
+++ b/src/main/java/com/android/tools/r8/utils/UTF8TextInputStream.java
@@ -5,6 +5,7 @@
 package com.android.tools.r8.utils;
 
 import com.android.tools.r8.TextInputStream;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.Charset;
@@ -20,6 +21,10 @@
     this(Files.newInputStream(path));
   }
 
+  public UTF8TextInputStream(String string) {
+    this(new ByteArrayInputStream(string.getBytes()));
+  }
+
   public UTF8TextInputStream(InputStream inputStream) {
     this.inputStream = inputStream;
   }
diff --git a/src/test/java/com/android/tools/r8/profile/art/dump/DumpArtProfileProvidersTest.java b/src/test/java/com/android/tools/r8/profile/art/dump/DumpArtProfileProvidersTest.java
new file mode 100644
index 0000000..71a58e7
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/profile/art/dump/DumpArtProfileProvidersTest.java
@@ -0,0 +1,192 @@
+// 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.profile.art.dump;
+
+import static com.android.tools.r8.DiagnosticsMatcher.diagnosticMessage;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import com.android.tools.r8.CompilationFailedException;
+import com.android.tools.r8.R8FullTestBuilder;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.origin.Origin;
+import com.android.tools.r8.profile.art.ArtProfileBuilder;
+import com.android.tools.r8.profile.art.ArtProfileProvider;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.MethodReference;
+import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.utils.AndroidApiLevel;
+import com.android.tools.r8.utils.DumpInputFlags;
+import com.android.tools.r8.utils.FileUtils;
+import com.android.tools.r8.utils.MethodReferenceUtils;
+import com.android.tools.r8.utils.StringUtils;
+import com.android.tools.r8.utils.UTF8TextInputStream;
+import com.android.tools.r8.utils.ZipUtils;
+import com.google.common.collect.Lists;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+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 DumpArtProfileProvidersTest extends TestBase {
+
+  private enum DumpStrategy {
+    DIRECTORY,
+    FILE;
+
+    DumpInputFlags createDumpInputFlags(Path dump) {
+      if (this == DIRECTORY) {
+        return DumpInputFlags.dumpToDirectory(dump);
+      }
+      assert this == FILE;
+      return DumpInputFlags.dumpToFile(dump);
+    }
+
+    Path createDumpPath(TemporaryFolder temp) throws IOException {
+      if (this == DIRECTORY) {
+        return temp.newFolder().toPath();
+      }
+      assert this == FILE;
+      return temp.newFile("dump.zip").toPath();
+    }
+  }
+
+  @Parameter(0)
+  public DumpStrategy dumpStrategy;
+
+  @Parameter(1)
+  public TestParameters parameters;
+
+  @Parameters(name = "{1}, {0}")
+  public static List<Object[]> data() {
+    return buildParameters(DumpStrategy.values(), getTestParameters().withNoneRuntime().build());
+  }
+
+  @Test
+  public void test() throws Exception {
+    Path dump = dumpStrategy.createDumpPath(temp);
+    DumpInputFlags dumpInputFlags = dumpStrategy.createDumpInputFlags(dump);
+    try {
+      testForR8(Backend.DEX)
+          .addProgramClasses(Main.class)
+          .addKeepMainRule(Main.class)
+          .addOptionsModification(options -> options.setDumpInputFlags(dumpInputFlags))
+          .allowDiagnosticInfoMessages()
+          .apply(this::addArtProfileProviders)
+          .setMinApi(AndroidApiLevel.LATEST)
+          .compileWithExpectedDiagnostics(
+              diagnostics -> {
+                if (dumpInputFlags.shouldFailCompilation()) {
+                  diagnostics.assertErrorsMatch(
+                      diagnosticMessage(containsString("Dumped compilation inputs to:")));
+                } else {
+                  diagnostics.assertInfosMatch(
+                      diagnosticMessage(containsString("Dumped compilation inputs to:")));
+                }
+              });
+      assertFalse("Expected compilation to fail", dumpInputFlags.shouldFailCompilation());
+    } catch (CompilationFailedException e) {
+      assertTrue("Expected compilation to succeed", dumpInputFlags.shouldFailCompilation());
+    }
+    verifyDump(dump);
+  }
+
+  private void addArtProfileProviders(R8FullTestBuilder testBuilder) {
+    testBuilder.addArtProfileForRewriting(
+        new ArtProfileProvider() {
+
+          @Override
+          public void getArtProfile(ArtProfileBuilder profileBuilder) {
+            profileBuilder.addHumanReadableArtProfile(
+                new UTF8TextInputStream(StringUtils.joinLines("# Comment", "Lfoo/Bar;")),
+                parserBuilder -> {});
+            ClassReference bazClassReference = Reference.classFromDescriptor("Lfoo/Baz;");
+            MethodReference bazMainMethodReference =
+                MethodReferenceUtils.mainMethod(bazClassReference);
+            profileBuilder.addClassRule(
+                classRuleBuilder -> classRuleBuilder.setClassReference(bazClassReference));
+            profileBuilder.addMethodRule(
+                methodRuleBuilder ->
+                    methodRuleBuilder
+                        .setMethodReference(bazMainMethodReference)
+                        .setMethodRuleInfo(
+                            methodRuleInfoBuilder -> methodRuleInfoBuilder.setIsHot(true)));
+          }
+
+          @Override
+          public Origin getOrigin() {
+            return Origin.unknown();
+          }
+        });
+    testBuilder.addArtProfileForRewriting(
+        new ArtProfileProvider() {
+
+          @Override
+          public void getArtProfile(ArtProfileBuilder profileBuilder) {
+            ClassReference bazClassReference = Reference.classFromDescriptor("Lfoo/Baz;");
+            MethodReference bazMainMethodReference =
+                MethodReferenceUtils.mainMethod(bazClassReference);
+            profileBuilder.addClassRule(
+                classRuleBuilder -> classRuleBuilder.setClassReference(bazClassReference));
+            profileBuilder.addHumanReadableArtProfile(
+                new UTF8TextInputStream(StringUtils.joinLines("# Comment", "Lfoo/Bar;")),
+                parserBuilder -> {});
+            profileBuilder.addMethodRule(
+                methodRuleBuilder ->
+                    methodRuleBuilder
+                        .setMethodReference(bazMainMethodReference)
+                        .setMethodRuleInfo(
+                            methodRuleInfoBuilder -> methodRuleInfoBuilder.setIsHot(true)));
+          }
+
+          @Override
+          public Origin getOrigin() {
+            return Origin.unknown();
+          }
+        });
+  }
+
+  private void verifyDump(Path dump) throws IOException {
+    if (dumpStrategy == DumpStrategy.DIRECTORY) {
+      List<Path> dumps =
+          Files.walk(dump, 1).filter(path -> path.toFile().isFile()).collect(Collectors.toList());
+      assertEquals(1, dumps.size());
+      dump = dumps.get(0);
+    }
+
+    assertTrue(Files.exists(dump));
+    Path unzipped = temp.newFolder().toPath();
+    ZipUtils.unzip(dump.toString(), unzipped.toFile());
+
+    Path artProfile1 = unzipped.resolve("art-profile-1.txt");
+    assertTrue(Files.exists(artProfile1));
+    assertEquals(
+        Lists.newArrayList(
+            "# Comment", "Lfoo/Bar;", "Lfoo/Baz;", "HLfoo/Baz;->main([Ljava/lang/String;)V"),
+        FileUtils.readAllLines(artProfile1));
+
+    Path artProfile2 = unzipped.resolve("art-profile-2.txt");
+    assertTrue(Files.exists(artProfile2));
+    assertEquals(
+        Lists.newArrayList(
+            "Lfoo/Baz;", "# Comment", "Lfoo/Bar;", "HLfoo/Baz;->main([Ljava/lang/String;)V"),
+        FileUtils.readAllLines(artProfile2));
+  }
+
+  static class Main {
+
+    public static void main(String[] args) {}
+  }
+}