Android: Use d8's desugar dependencies

Based on Søren's sample code in http://b/176818423.

D8's java API allows us to determine the exact set of dependencies that
a class in the current target depends on to desugar correctly (generally
its superclass(es) and interface(s) it implements). The class would only
need to be re-desugared with D8 if one of these dependencies changes.

This CL makes use of this information in a less granular way. If any jar
containing one of these dependencies is changed, then we re-dex all the
classes. Even so, this limited reduction in dependencies from using D8
for desugaring reduced the previous regression from +50% to +10% for the
base_java_sig benchmark.

Follow-up CLs will attempt to make use of the more granular desugar
dependency information to further improve the timing of D8 desugar. For
the time being, we are continuing to use bazel desugar.

Bug: 1015559
Change-Id: If507a4013af7bed0820758629a4b9db69942727b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2621928
Commit-Queue: Peter Wen <wnwen@chromium.org>
Auto-Submit: Peter Wen <wnwen@chromium.org>
Reviewed-by: Andrew Grieve <agrieve@chromium.org>
Cr-Commit-Position: refs/heads/master@{#844201}

BUILD.gn

@@ -353,6 +353,7 @@ group("gn_all") {
       "//services:services_junit_tests",
       "//testing/android/junit:junit_unit_tests",
       "//third_party/catapult/devil",
+      "//third_party/r8:custom_d8_java",
       "//third_party/smhasher:murmurhash3",
       "//tools/android:android_tools",
       "//tools/android:memconsumer",

build/android/gyp/dex.py

@@ -70,6 +70,11 @@ def _ParseArgs(args):
                       action='store_true',
                       help='Allow numerous dex files within output.')
   parser.add_argument('--r8-jar-path', required=True, help='Path to R8 jar.')
+  parser.add_argument('--custom-d8-jar-path',
+                      required=True,
+                      help='Path to our customized d8 jar.')
+  parser.add_argument('--desugar-dependencies',
+                      help='Path to store desugar dependencies.')
   parser.add_argument('--desugar', action='store_true')
   parser.add_argument(
       '--bootclasspath',
@@ -371,6 +376,7 @@ def _CreateFinalDex(d8_inputs, output, tmp_dir, dex_cmd, options=None):
 
     tmp_dex_dir = os.path.join(tmp_dir, 'tmp_dex_dir')
     os.mkdir(tmp_dex_dir)
+
     _RunD8(dex_cmd, d8_inputs, tmp_dex_dir,
            (not options or options.warnings_as_errors),
            (options and options.show_desugar_default_interface_warnings))
@@ -498,6 +504,7 @@ def main(args):
   if options.multi_dex and options.main_dex_list_path:
     input_paths.append(options.main_dex_list_path)
   input_paths.append(options.r8_jar_path)
+  input_paths.append(options.custom_d8_jar_path)
 
   depfile_deps = options.class_inputs_filearg + options.dex_inputs_filearg
 
@@ -515,8 +522,8 @@ def main(args):
 
   dex_cmd = build_utils.JavaCmd(options.warnings_as_errors) + [
       '-cp',
-      options.r8_jar_path,
-      'com.android.tools.r8.D8',
+      '{}:{}'.format(options.r8_jar_path, options.custom_d8_jar_path),
+      'org.chromium.build.CustomD8',
   ]
   if options.release:
     dex_cmd += ['--release']
@@ -526,17 +533,30 @@ def main(args):
   if not options.desugar:
     dex_cmd += ['--no-desugaring']
   elif options.classpath:
-    # Don't pass classpath when Desugar.jar is doing interface desugaring.
+    # The classpath is used by D8 to for interface desugaring.
+    classpath_paths = options.classpath
+    if options.desugar_dependencies:
+      dex_cmd += ['--desugar-dependencies', options.desugar_dependencies]
+      if os.path.exists(options.desugar_dependencies):
+        with open(options.desugar_dependencies, 'r') as f:
+          lines = [line.strip() for line in f.readlines()]
+          # Use a set to deduplicate entries.
+          desugar_dependencies = set(dep for dep in lines if dep)
+        # Desugar dependencies are a subset of classpath.
+        classpath_paths = list(desugar_dependencies)
+    depfile_deps += classpath_paths
+    input_paths += classpath_paths
     dex_cmd += ['--lib', build_utils.JAVA_HOME]
     for path in options.bootclasspath:
       dex_cmd += ['--lib', path]
+    # Still pass the entire classpath in case a new dependency is needed by
+    # desugar, so that desugar_dependencies will be updated for the next build.
     for path in options.classpath:
       dex_cmd += ['--classpath', path]
-    depfile_deps += options.classpath
     depfile_deps += options.bootclasspath
-    input_paths += options.classpath
     input_paths += options.bootclasspath
 
+
   if options.desugar_jdk_libs_json:
     dex_cmd += ['--desugared-lib', options.desugar_jdk_libs_json]
   if options.force_enable_assertions:

build/config/android/internal_rules.gni

@@ -82,6 +82,7 @@ _java_library_patterns = [
 java_target_patterns = _java_library_patterns + _java_resource_patterns
 
 _r8_path = "//third_party/r8/lib/r8.jar"
+_custom_d8_path = "//third_party/r8/custom_d8.jar"
 _desugar_jdk_libs_json = "//third_party/r8/desugar_jdk_libs.json"
 _desugar_jdk_libs_jar = "//third_party/android_deps/libs/com_android_tools_desugar_jdk_libs/desugar_jdk_libs-1.0.10.jar"
 _desugar_jdk_libs_configuration_jar = "//third_party/android_deps/libs/com_android_tools_desugar_jdk_libs_configuration/desugar_jdk_libs_configuration-1.0.10.jar"
@@ -1666,7 +1667,10 @@ if (enable_java_templates) {
         deps = _deps
         depfile = "$target_gen_dir/$target_name.d"
         outputs = [ invoker.output ]
-        inputs = [ _r8_path ]
+        inputs = [
+          _r8_path,
+          _custom_d8_path,
+        ]
 
         if (!_is_library) {
           # http://crbug.com/725224. Fix for bots running out of memory.
@@ -1685,6 +1689,8 @@ if (enable_java_templates) {
           "--min-api=$_min_sdk_version",
           "--r8-jar-path",
           rebase_path(_r8_path, root_build_dir),
+          "--custom-d8-jar-path",
+          rebase_path(_custom_d8_path, root_build_dir),
         ]
         if (treat_warnings_as_errors) {
           args += [ "--warnings-as-errors" ]
@@ -1773,8 +1779,15 @@ if (enable_java_templates) {
           }
         }
         if (_desugar_needs_classpath) {
+          _desugar_dependencies_path =
+              "$target_gen_dir/$target_name.desugardeps"
           args += [
+            "--desugar-dependencies",
+            rebase_path(_desugar_dependencies_path, root_build_dir),
             "--bootclasspath=@FileArg($_rebased_build_config:android:sdk_jars)",
+
+            # Pass the full classpath to find new dependencies that are not in
+            # the .desugardeps file.
             "--classpath=@FileArg($_rebased_build_config:deps_info:javac_full_interface_classpath)",
           ]
           if (defined(invoker.desugar_jars_paths)) {

third_party/r8/BUILD.gn

+# Copyright 2021 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import("//build/config/android/rules.gni")
+
+# The output jar for this target is copied and checked-in to minimize the number
+# of edges in the build graph. This is a class that rarely changes.
+java_library("custom_d8_java") {
+  sources = [ "java/src/org/chromium/build/CustomD8.java" ]
+
+  # Avoid using java_prebuilt() to ensure all uses go through the checked-in
+  # version.
+  input_jars_paths = [ "//third_party/r8/lib/r8.jar" ]
+}

third_party/r8/README.chromium

@@ -23,8 +23,14 @@ Local Modifications:
   * Plus: https://r8-review.googlesource.com/c/r8/+/56708
 * Added "playground" directory for quick "how does this optimize" tests.
 * Removed references to ConcurrentHashMap and TimeZone#getTimeZone in desugar_jdk_libs.json.
+* Added "java/src/org/chromium/build/CustomD8.java", custom_d8.jar, and BUILD.gn.
+  * Used in "build/android/gyp/dex.py" to enable desugar dependencies.
 
 Update Instructions:
+# For the custom d8 jar (required only when CustomD8.java changes):
+autoninja -C out/Debug third_party/r8:custom_d8_java
+cp out/Debug/obj/third_party/r8/custom_d8_java.javac.jar third_party/r8/custom_d8.jar
+
 # Download R8:
 git clone https://r8.googlesource.com/r8
 cd r8

third_party/r8/java/src/org/chromium/build/CustomD8.java

+// Copyright 2021 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.build;
+
+import com.android.tools.r8.CompilationFailedException;
+import com.android.tools.r8.D8;
+import com.android.tools.r8.D8Command;
+import com.android.tools.r8.DesugarGraphConsumer;
+import com.android.tools.r8.origin.Origin;
+
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+
+public class CustomD8 {
+    private static class CommandLineOrigin extends Origin {
+        private CommandLineOrigin() {
+            super(root());
+        }
+
+        @Override
+        public String part() {
+            return "Command line";
+        }
+    }
+
+    // Entry point for D8 compilation with support for --desugar-dependencies option
+    // as well.
+    public static void main(String[] args) throws CompilationFailedException, IOException {
+        String desugarDependenciesOptions = "--desugar-dependencies";
+        String desugarDependenciesPath = null;
+        String[] d8Args = null;
+
+        int desugarDepIdx = Arrays.asList(args).indexOf(desugarDependenciesOptions);
+        if (desugarDepIdx != -1) {
+            int numRemainingArgs = args.length - (desugarDepIdx + 2);
+            if (numRemainingArgs < 0) {
+                throw new CompilationFailedException(
+                        "Missing argument to '" + desugarDependenciesOptions + "'");
+            }
+            desugarDependenciesPath = args[desugarDepIdx + 1];
+            d8Args = new String[args.length - 2];
+            // Copy over all other args before and after the desugar dependencies arg.
+            System.arraycopy(args, 0, d8Args, 0, desugarDepIdx);
+            System.arraycopy(args, desugarDepIdx + 2, d8Args, desugarDepIdx, numRemainingArgs);
+        } else {
+            d8Args = args;
+        }
+
+        // Use D8 command line parser to handle the normal D8 command line.
+        D8Command.Builder builder = D8Command.parse(d8Args, new CommandLineOrigin());
+        // If additional options was passed amend the D8 command builder.
+        if (desugarDependenciesPath != null) {
+            final Path desugarDependencies = Paths.get(desugarDependenciesPath);
+            PrintWriter desugarDependenciesPrintWriter =
+                    new PrintWriter(Files.newOutputStream(desugarDependencies));
+            if (builder.getDesugarGraphConsumer() != null) {
+                throw new CompilationFailedException("Too many desugar graph consumers.");
+            }
+            builder.setDesugarGraphConsumer(new DesugarGraphConsumer() {
+                @Override
+                public void accept(Origin dependent, Origin dependency) {
+                    // The target's class files have root as their parent.
+                    if (dependency.parent().equals(Origin.root())) {
+                        return;
+                    }
+                    desugarDependenciesPrintWriter.println(dependency.parent());
+                }
+
+                @Override
+                public void finished() {
+                    desugarDependenciesPrintWriter.close();
+                }
+            });
+        }
+
+        // Run D8.
+        D8.run(builder.build());
+    }
+}