Minor updates for relayouting agsa with Pixel 5 profile

Change-Id: I2cb7ef37bc88ea08a7686719aac145ae541bbcf8
diff --git a/src/main/java/com/android/tools/r8/D8Command.java b/src/main/java/com/android/tools/r8/D8Command.java
index e7f27d5..35d22bb 100644
--- a/src/main/java/com/android/tools/r8/D8Command.java
+++ b/src/main/java/com/android/tools/r8/D8Command.java
@@ -702,7 +702,9 @@
         .getArtProfileOptions()
         .setArtProfilesForRewriting(getArtProfilesForRewriting())
         .setPassthrough(true);
-    internal.getStartupOptions().setStartupProfileProviders(getStartupProfileProviders());
+    if (!getStartupProfileProviders().isEmpty()) {
+      internal.getStartupOptions().setStartupProfileProviders(getStartupProfileProviders());
+    }
 
     internal.programClassConflictResolver =
         ProgramClassCollection.wrappedConflictResolver(
diff --git a/src/main/java/com/android/tools/r8/R8Command.java b/src/main/java/com/android/tools/r8/R8Command.java
index 81efb90..e579889 100644
--- a/src/main/java/com/android/tools/r8/R8Command.java
+++ b/src/main/java/com/android/tools/r8/R8Command.java
@@ -1077,7 +1077,9 @@
     internal.configureAndroidPlatformBuild(getAndroidPlatformBuild());
 
     internal.getArtProfileOptions().setArtProfilesForRewriting(getArtProfilesForRewriting());
-    internal.getStartupOptions().setStartupProfileProviders(getStartupProfileProviders());
+    if (!getStartupProfileProviders().isEmpty()) {
+      internal.getStartupOptions().setStartupProfileProviders(getStartupProfileProviders());
+    }
 
     internal.programClassConflictResolver =
         ProgramClassCollection.wrappedConflictResolver(
diff --git a/src/main/java/com/android/tools/r8/experimental/startup/StartupProfileProviderUtils.java b/src/main/java/com/android/tools/r8/experimental/startup/StartupProfileProviderUtils.java
index 3b5b940..0eef997 100644
--- a/src/main/java/com/android/tools/r8/experimental/startup/StartupProfileProviderUtils.java
+++ b/src/main/java/com/android/tools/r8/experimental/startup/StartupProfileProviderUtils.java
@@ -9,10 +9,14 @@
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.origin.PathOrigin;
 import com.android.tools.r8.profile.art.ArtProfileBuilderUtils.SyntheticToSyntheticContextGeneralization;
+import com.android.tools.r8.profile.art.ArtProfileClassRuleInfo;
+import com.android.tools.r8.profile.art.ArtProfileMethodRuleInfo;
+import com.android.tools.r8.profile.art.ArtProfileRulePredicate;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.MethodReference;
 import com.android.tools.r8.startup.StartupProfileBuilder;
 import com.android.tools.r8.startup.StartupProfileProvider;
 import com.android.tools.r8.startup.diagnostic.MissingStartupProfileItemsDiagnostic;
-import com.android.tools.r8.utils.ConsumerUtils;
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.UTF8TextInputStream;
 import java.io.IOException;
@@ -28,7 +32,23 @@
       public void getStartupProfile(StartupProfileBuilder startupProfileBuilder) {
         try {
           startupProfileBuilder.addHumanReadableArtProfile(
-              new UTF8TextInputStream(path), ConsumerUtils.emptyConsumer());
+              new UTF8TextInputStream(path),
+              profileParserBuilder ->
+                  profileParserBuilder.setRulePredicate(
+                      new ArtProfileRulePredicate() {
+                        @Override
+                        public boolean testClassRule(
+                            ClassReference classReference, ArtProfileClassRuleInfo classRuleInfo) {
+                          return false;
+                        }
+
+                        @Override
+                        public boolean testMethodRule(
+                            MethodReference methodReference,
+                            ArtProfileMethodRuleInfo methodRuleInfo) {
+                          return methodRuleInfo.isStartup();
+                        }
+                      }));
         } catch (IOException e) {
           throw new UncheckedIOException(e);
         }
diff --git a/src/main/java/com/android/tools/r8/profile/art/HumanReadableArtProfileParser.java b/src/main/java/com/android/tools/r8/profile/art/HumanReadableArtProfileParser.java
index 3abbbca..cc39eb9 100644
--- a/src/main/java/com/android/tools/r8/profile/art/HumanReadableArtProfileParser.java
+++ b/src/main/java/com/android/tools/r8/profile/art/HumanReadableArtProfileParser.java
@@ -9,8 +9,9 @@
 import com.android.tools.r8.profile.art.diagnostic.HumanReadableArtProfileParserErrorDiagnostic;
 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.references.TypeReference;
 import com.android.tools.r8.utils.Action;
-import com.android.tools.r8.utils.ClassReferenceUtils;
 import com.android.tools.r8.utils.MethodReferenceUtils;
 import com.android.tools.r8.utils.Reporter;
 import java.io.BufferedReader;
@@ -65,12 +66,16 @@
   }
 
   public boolean parseRule(String rule) {
-    ArtProfileMethodRuleInfoImpl.Builder methodRuleInfoBuilder =
-        ArtProfileMethodRuleInfoImpl.builder();
-    rule = parseFlag(rule, 'H', methodRuleInfoBuilder::setIsHot);
-    rule = parseFlag(rule, 'S', methodRuleInfoBuilder::setIsStartup);
-    rule = parseFlag(rule, 'P', methodRuleInfoBuilder::setIsPostStartup);
-    return parseClassOrMethodDescriptor(rule, methodRuleInfoBuilder.build());
+    try {
+      ArtProfileMethodRuleInfoImpl.Builder methodRuleInfoBuilder =
+          ArtProfileMethodRuleInfoImpl.builder();
+      rule = parseFlag(rule, 'H', methodRuleInfoBuilder::setIsHot);
+      rule = parseFlag(rule, 'S', methodRuleInfoBuilder::setIsStartup);
+      rule = parseFlag(rule, 'P', methodRuleInfoBuilder::setIsPostStartup);
+      return parseClassOrMethodRule(rule, methodRuleInfoBuilder.build());
+    } catch (Throwable t) {
+      return false;
+    }
   }
 
   private static String parseFlag(String rule, char c, Action action) {
@@ -81,23 +86,30 @@
     return rule;
   }
 
-  private boolean parseClassOrMethodDescriptor(
-      String descriptor, ArtProfileMethodRuleInfoImpl methodRuleInfo) {
-    int arrowStartIndex = descriptor.indexOf("->");
+  private boolean parseClassOrMethodRule(String rule, ArtProfileMethodRuleInfoImpl methodRuleInfo) {
+    int arrowStartIndex = rule.indexOf("->");
     if (arrowStartIndex >= 0) {
-      return parseMethodRule(descriptor, methodRuleInfo, arrowStartIndex);
+      return parseMethodRule(rule, methodRuleInfo, arrowStartIndex);
     } else if (methodRuleInfo.isEmpty()) {
-      return parseClassRule(descriptor);
+      return parseClassRule(rule);
     } else {
       return false;
     }
   }
 
   private boolean parseClassRule(String descriptor) {
-    ClassReference classReference = ClassReferenceUtils.parseClassDescriptor(descriptor);
-    if (classReference == null) {
+    TypeReference typeReference = Reference.typeFromDescriptor(descriptor);
+    if (typeReference == null) {
       return false;
     }
+    if (typeReference.isArray()) {
+      typeReference = typeReference.asArray().getBaseType();
+    }
+    if (typeReference.isPrimitive()) {
+      return false;
+    }
+    assert typeReference.isClass();
+    ClassReference classReference = typeReference.asClass();
     if (rulePredicate.testClassRule(classReference, ArtProfileClassRuleInfoImpl.empty())) {
       profileBuilder.addClassRule(
           classRuleBuilder -> classRuleBuilder.setClassReference(classReference));
@@ -106,7 +118,9 @@
   }
 
   private boolean parseMethodRule(
-      String descriptor, ArtProfileMethodRuleInfoImpl methodRuleInfo, int arrowStartIndex) {
+      String rule, ArtProfileMethodRuleInfoImpl methodRuleInfo, int arrowStartIndex) {
+    int inlineCacheStartIndex = rule.indexOf('+', arrowStartIndex + 2);
+    String descriptor = inlineCacheStartIndex > 0 ? rule.substring(0, inlineCacheStartIndex) : rule;
     MethodReference methodReference =
         MethodReferenceUtils.parseSmaliString(descriptor, arrowStartIndex);
     if (methodReference == null) {
diff --git a/tools/extractmarker.py b/tools/extractmarker.py
index b50a2fd..d8348dd 100755
--- a/tools/extractmarker.py
+++ b/tools/extractmarker.py
@@ -3,8 +3,30 @@
 # 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.
 
+import argparse
 import sys
 import toolhelper
 
+def extractmarker(apk_or_dex, build=True):
+  stdout = toolhelper.run('extractmarker', [apk_or_dex], build=build, return_stdout=True)
+  return stdout
+
+def parse_options(argv):
+  result = argparse.ArgumentParser(
+      description='Relayout a given APK using a startup profile.')
+  result.add_argument('--no-build',
+                      action='store_true',
+                      default=False,
+                      help='To disable building using gradle')
+  options, args = result.parse_known_args(argv)
+  return options, args
+
+def main(argv):
+  options, args = parse_options(argv)
+  build = not options.no_build
+  for apk_or_dex in args:
+    print(extractmarker(apk_or_dex, build=build))
+    build = False
+
 if __name__ == '__main__':
-  sys.exit(toolhelper.run('extractmarker', sys.argv[1:]))
+  sys.exit(main(sys.argv[1:]))
diff --git a/tools/startup/relayout.py b/tools/startup/relayout.py
index 6220336..0ace441 100755
--- a/tools/startup/relayout.py
+++ b/tools/startup/relayout.py
@@ -11,8 +11,10 @@
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 import apk_masseur
+import extractmarker
 import toolhelper
 import utils
+import zip_utils
 
 def parse_options(argv):
   result = argparse.ArgumentParser(
@@ -20,6 +22,15 @@
   result.add_argument('--apk',
                       help='Path to the .apk',
                       required=True)
+  result.add_argument('--desugared-library',
+                      choices=['auto', 'true', 'false'],
+                      default='auto',
+                      help='Whether the last dex file of the app is desugared '
+                           'library')
+  result.add_argument('--no-build',
+                      action='store_true',
+                      default=False,
+                      help='To disable building using gradle')
   result.add_argument('--out',
                       help='Destination of resulting apk',
                       required=True)
@@ -29,6 +40,29 @@
   options, args = result.parse_known_args(argv)
   return options, args
 
+def get_dex_to_relayout(options, temp):
+  marker = extractmarker.extractmarker(options.apk, build=not options.no_build)
+  if '~~L8' not in marker:
+    return [options.apk], None
+  dex_dir = os.path.join(temp, 'dex')
+  dex_predicate = \
+      lambda name : name.startswith('classes') and name.endswith('.dex')
+  extracted_dex_files = \
+      zip_utils.extract_all_that_matches(options.apk, dex_dir, dex_predicate)
+  desugared_library_dex = 'classes%s.dex' % len(extracted_dex_files)
+  assert desugared_library_dex in extracted_dex_files
+  return [
+      os.path.join(dex_dir, name) \
+          for name in extracted_dex_files if name != desugared_library_dex], \
+      os.path.join(dex_dir, desugared_library_dex)
+
+def has_desugared_library_dex(options):
+  if options.desugared_library == 'auto':
+    marker = extractmarker.extractmarker(
+        options.apk, build=not options.no_build)
+    return '~~L8' in marker
+  return options.desugared_library == 'true'
+
 def get_min_api(apk):
   aapt = os.path.join(utils.getAndroidBuildTools(), 'aapt')
   cmd = [aapt, 'dump', 'badging', apk]
@@ -46,14 +80,21 @@
         '--min-api', str(get_min_api(options.apk)),
         '--output', dex,
         '--no-desugaring',
-        '--release',
-        options.apk]
+        '--release']
+    dex_to_relayout, desugared_library_dex = get_dex_to_relayout(options, temp)
+    d8_args.extend(dex_to_relayout)
     extra_args = ['-Dcom.android.tools.r8.startup.profile=%s' % options.profile]
     toolhelper.run(
         'd8',
         d8_args,
+        build=not options.no_build,
         extra_args=extra_args,
         main='com.android.tools.r8.D8')
+    if desugared_library_dex is not None:
+      dex_files = [name for name in \
+          zip_utils.get_names_that_matches(dex, lambda x : True)]
+      zip_utils.add_file_to_zip(
+          desugared_library_dex, 'classes%s.dex' % str(len(dex_files) + 1), dex)
     apk_masseur.masseur(options.apk, dex=dex, out=options.out)
 
 if __name__ == '__main__':
diff --git a/tools/toolhelper.py b/tools/toolhelper.py
index c728953..881f6dc 100644
--- a/tools/toolhelper.py
+++ b/tools/toolhelper.py
@@ -63,10 +63,10 @@
       stdout, stderr = proc.communicate()
     finally:
       timer.cancel()
-    result = stdout if return_stdout else proc.returncode
+    result = stdout.decode('utf-8') if return_stdout else proc.returncode
   else:
     result = (
-        subprocess.check_output(cmd)
+        subprocess.check_output(cmd).decode('utf-8')
         if return_stdout
         else subprocess.call(cmd, stdout=stdout, stderr=stderr))
   duration = int((time.time() - start) * 1000)
diff --git a/tools/zip_utils.py b/tools/zip_utils.py
index b0571f2..795119b 100644
--- a/tools/zip_utils.py
+++ b/tools/zip_utils.py
@@ -8,3 +8,13 @@
 def add_file_to_zip(file, destination, zip_file):
   with zipfile.ZipFile(zip_file, 'a') as zip:
     zip.write(file, destination)
+
+def extract_all_that_matches(zip_file, destination, predicate):
+  with zipfile.ZipFile(zip_file) as zip:
+    names_to_extract = [name for name in zip.namelist() if predicate(name)]
+    zip.extractall(path=destination, members=names_to_extract)
+    return names_to_extract
+
+def get_names_that_matches(zip_file, predicate):
+  with zipfile.ZipFile(zip_file) as zip:
+    return [name for name in zip.namelist() if predicate(name)]