From 76fac817e4dbe5aaece9ed6a962e1cc94815d4d6 Mon Sep 17 00:00:00 2001
From: Googler <noreply@google.com>
Date: Thu, 14 Apr 2022 11:40:13 -0700
Subject: [PATCH] Add Blaze support for "local" optimization of individual Java
 libraries.

This should allow optimization on a per-library level within Blaze of every
artifact generated by javac. The purpose of this is for improved test coverage
by the AppReduce team and is not intended for general use. See this thread for
more context:

[]

This change should only add a dependency on the bytecode optimizer when
--experimental_local_java_optimizations is true, so it should be a no-op for
all standard builds.

I implemented this in Java instead of Starlark because it allowed me to modify
the produced JavaInfo much more simply. If this was implemented in Starlark,
I'd have to be able to replace the JavaInfo object produced by a java_* target
since the JavaInfo from javac is currently all still wired up in Java.  My
attempts at structuring the change that way seemed much more invasive to
JavaInfo and more brittle.

PiperOrigin-RevId: 441814811
---
 .../lib/rules/java/JavaCompilationHelper.java | 70 +++++++++++++++++--
 .../lib/rules/java/JavaConfiguration.java     | 25 ++++++-
 .../build/lib/rules/java/JavaOptions.java     | 25 +++++++
 .../build/lib/rules/java/JavaSemantics.java   | 15 +++-
 .../build/lib/rules/java/JavaToolchain.java   |  8 +++
 .../lib/rules/java/JavaToolchainProvider.java | 19 +++++
 .../lib/rules/java/JavaToolchainRule.java     |  9 +++
 7 files changed, 163 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompilationHelper.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompilationHelper.java
index a4906dc62152f8..7f94317e74677e 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompilationHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompilationHelper.java
@@ -270,6 +270,20 @@ && getJavaConfiguration().experimentalEnableJspecify()
       createResourceJarAction(originalOutput, ImmutableList.copyOf(resourceJars));
     }
 
+    Artifact optimizedJar = null;
+    if (getJavaConfiguration().runLocalJavaOptimizations()) {
+      optimizedJar = outputs.output();
+      outputs =
+          outputs.withOutput(
+              ruleContext.getDerivedArtifact(
+                  FileSystemUtils.replaceExtension(
+                      outputs
+                          .output()
+                          .getOutputDirRelativePath(getConfiguration().isSiblingRepositoryLayout()),
+                      "-pre-optimization.jar"),
+                  outputs.output().getRoot()));
+    }
+
     ImmutableList<String> javacopts = customJavacOpts;
     if (jspecify) {
       plugins =
@@ -294,12 +308,13 @@ && getJavaConfiguration().experimentalEnableJspecify()
     builder.setCoverageArtifact(coverageArtifact);
     BootClassPathInfo bootClassPathInfo = getBootclasspathOrDefault();
     builder.setBootClassPath(bootClassPathInfo);
+    NestedSet<Artifact> classpath =
+        NestedSetBuilder.<Artifact>naiveLinkOrder()
+            .addTransitive(bootClassPathInfo.auxiliary())
+            .addTransitive(attributes.getCompileTimeClassPath())
+            .build();
     if (!bootClassPathInfo.auxiliary().isEmpty()) {
-      builder.setClasspathEntries(
-          NestedSetBuilder.<Artifact>naiveLinkOrder()
-              .addTransitive(bootClassPathInfo.auxiliary())
-              .addTransitive(attributes.getCompileTimeClassPath())
-              .build());
+      builder.setClasspathEntries(classpath);
       builder.setDirectJars(
           NestedSetBuilder.<Artifact>naiveLinkOrder()
               .addTransitive(bootClassPathInfo.auxiliary())
@@ -355,6 +370,20 @@ && getJavaConfiguration().experimentalEnableJspecify()
 
     JavaCompileAction javaCompileAction = builder.build();
     ruleContext.getAnalysisEnvironment().registerAction(javaCompileAction);
+
+    if (optimizedJar != null) {
+      JavaConfiguration.NamedLabel optimizerLabel = getJavaConfiguration().getBytecodeOptimizer();
+      createLocalOptimizationAction(
+          outputs.output(),
+          optimizedJar,
+          NestedSetBuilder.<Artifact>naiveLinkOrder()
+              .addTransitive(bootClassPathInfo.bootclasspath())
+              .addTransitive(classpath)
+              .build(),
+          javaToolchain.getLocalJavaOptimizationConfiguration(),
+          javaToolchain.getBytecodeOptimizer().tool(),
+          optimizerLabel.name());
+    }
   }
 
   /**
@@ -705,6 +734,37 @@ public Artifact createCompileTimeJarAction(
     return jar;
   }
 
+  public void createLocalOptimizationAction(
+      Artifact unoptimizedOutputJar,
+      Artifact optimizedOutputJar,
+      NestedSet<Artifact> classpath,
+      List<Artifact> configs,
+      FilesToRunProvider optimizer,
+      String mnemonic) {
+    CustomCommandLine.Builder command =
+        CustomCommandLine.builder()
+            .add("-runtype", "LOCAL_ONLY")
+            .addExecPath("-injars", unoptimizedOutputJar)
+            .addExecPath("-outjars", optimizedOutputJar)
+            .addExecPaths("-libraryjars", CustomCommandLine.VectorArg.join(":").each(classpath));
+    for (Artifact config : configs) {
+      command.addPrefixedExecPath("@", config);
+    }
+
+    getRuleContext()
+        .registerAction(
+            new SpawnAction.Builder()
+                .addInput(unoptimizedOutputJar)
+                .addTransitiveInputs(classpath)
+                .addInputs(configs)
+                .addOutput(optimizedOutputJar)
+                .setExecutable(optimizer)
+                .addCommandLine(command.build())
+                .setProgressMessage("Optimizing jar %{label}")
+                .setMnemonic(mnemonic)
+                .build(getRuleContext()));
+  }
+
   private void addArgsAndJarsToAttributes(
       JavaCompilationArgsProvider args, NestedSet<Artifact> directJars) {
     // Can only be non-null when isStrict() returns true.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaConfiguration.java
index 0543e88b0a8e17..57020b7d1ee1c1 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaConfiguration.java
@@ -91,6 +91,8 @@ public enum ImportDepsCheckingLevel {
   private final Label proguardBinary;
   private final ImmutableList<Label> extraProguardSpecs;
   private final NamedLabel bytecodeOptimizer;
+  private final boolean runLocalJavaOptimizations;
+  private final ImmutableList<Label> localJavaOptimizationConfiguration;
   private final boolean splitBytecodeOptimizationPass;
   private final boolean enforceProguardFileExtension;
   private final boolean runAndroidLint;
@@ -125,6 +127,9 @@ public JavaConfiguration(BuildOptions buildOptions) throws InvalidConfigurationE
     this.fixDepsTool = javaOptions.fixDepsTool;
     this.proguardBinary = javaOptions.proguard;
     this.extraProguardSpecs = ImmutableList.copyOf(javaOptions.extraProguardSpecs);
+    this.runLocalJavaOptimizations = javaOptions.runLocalJavaOptimizations;
+    this.localJavaOptimizationConfiguration =
+        ImmutableList.copyOf(javaOptions.localJavaOptimizationConfiguration);
     this.splitBytecodeOptimizationPass = javaOptions.splitBytecodeOptimizationPass;
     this.enforceProguardFileExtension = javaOptions.enforceProguardFileExtension;
     this.useLegacyBazelJavaTest = javaOptions.legacyBazelJavaTest;
@@ -148,7 +153,7 @@ public JavaConfiguration(BuildOptions buildOptions) throws InvalidConfigurationE
     if (optimizers.size() > 1) {
       throw new InvalidConfigurationException(
           String.format(
-              "--experimental_bytecode_optimizers can only accept up to one mapping, but %s"
+              "--experimental_bytecode_optimizers can only accept up to one mapping, but %d"
                   + " mappings were provided.",
               optimizers.size()));
     }
@@ -159,6 +164,11 @@ public JavaConfiguration(BuildOptions buildOptions) throws InvalidConfigurationE
       throw new InvalidConfigurationException("Must supply label for optimizer " + mnemonic);
     }
     this.bytecodeOptimizer = NamedLabel.create(mnemonic, Optional.fromNullable(optimizerLabel));
+    if (runLocalJavaOptimizations && optimizerLabel == null) {
+      throw new InvalidConfigurationException(
+          "--experimental_local_java_optimizations cannot be provided without "
+              + "--experimental_bytecode_optimizers.");
+    }
 
     this.pluginList = ImmutableList.copyOf(javaOptions.pluginList);
     this.experimentalTurbineAnnotationProcessing =
@@ -312,11 +322,22 @@ public static NamedLabel create(String name, Optional<Label> label) {
     public abstract Optional<Label> label();
   }
 
-  /** Returns ordered list of optimizers to run. */
+  /** Returns bytecode optimizer to run. */
+  @Nullable
   public NamedLabel getBytecodeOptimizer() {
     return bytecodeOptimizer;
   }
 
+  /** Returns true if the bytecode optimizer should incrementally optimize all Java artifacts. */
+  public boolean runLocalJavaOptimizations() {
+    return runLocalJavaOptimizations;
+  }
+
+  /** Returns the optimization configuration for local Java optimizations if they are enabled. */
+  public ImmutableList<Label> getLocalJavaOptimizationConfiguration() {
+    return localJavaOptimizationConfiguration;
+  }
+
   /**
    * Returns true if java_test in Bazel should behave in legacy mode that existed before we
    * open-sourced our test runner.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaOptions.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaOptions.java
index 1475c135e53d80..375f3694452592 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaOptions.java
@@ -320,6 +320,31 @@ public ImportDepsCheckingLevelConverter() {
       help = "Do not use.")
   public Map<String, Label> bytecodeOptimizers;
 
+  /**
+   * If true, the bytecode optimizer will be used to incrementally optimize each compiled Java
+   * artifact.
+   */
+  @Option(
+      name = "experimental_local_java_optimizations",
+      defaultValue = "false",
+      documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help = "Do not use.")
+  public boolean runLocalJavaOptimizations;
+
+  /**
+   * Configuration for the bytecode optimizer if --experimental_local_java_optimizations is enabled.
+   */
+  @Option(
+      name = "experimental_local_java_optimization_configuration",
+      allowMultiple = true,
+      defaultValue = "null",
+      converter = LabelConverter.class,
+      documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help = "Do not use.")
+  public List<Label> localJavaOptimizationConfiguration;
+
   /**
    * If true, the OPTIMIZATION stage of the bytecode optimizer will be split across multiple
    * actions.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java
index 7bc0ad76b7f312..2598032b9a740a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaSemantics.java
@@ -184,12 +184,25 @@ static Label hostJdkAttribute(RuleDefinitionEnvironment env) {
                 attributes.has("proguard_specs")
                     && !attributes.get("proguard_specs", LABEL_LIST).isEmpty();
             JavaConfiguration.NamedLabel optimizer = javaConfig.getBytecodeOptimizer();
-            if (!hasProguardSpecs || !optimizer.label().isPresent()) {
+            if ((!hasProguardSpecs && !javaConfig.runLocalJavaOptimizations())
+                || !optimizer.label().isPresent()) {
               return null;
             }
             return optimizer.label().get();
           });
 
+  @SerializationConstant
+  LabelListLateBoundDefault<JavaConfiguration> LOCAL_JAVA_OPTIMIZATION_CONFIGURATION =
+      LabelListLateBoundDefault.fromTargetConfiguration(
+          JavaConfiguration.class,
+          (rule, attributes, javaConfig) -> {
+            // Don't bother adding the configuration dep if we're not going to use it.
+            if (!javaConfig.runLocalJavaOptimizations()) {
+              return ImmutableList.of();
+            }
+            return javaConfig.getLocalJavaOptimizationConfiguration();
+          });
+
   String JACOCO_METADATA_PLACEHOLDER = "%set_jacoco_metadata%";
   String JACOCO_MAIN_CLASS_PLACEHOLDER = "%set_jacoco_main_class%";
   String JACOCO_JAVA_RUNFILES_ROOT_PLACEHOLDER = "%set_jacoco_java_runfiles_root%";
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchain.java
index dd162ec228b408..155e75ffc08f38 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchain.java
@@ -141,6 +141,12 @@ public ConfiguredTarget create(RuleContext ruleContext)
               jspecifyProcessor, jspecifyImplicitDeps, jspecifyJavacopts.build(), jspecifyPackages);
     }
 
+    JavaToolchainTool bytecodeOptimizer =
+        JavaToolchainTool.fromFilesToRunProvider(
+            ruleContext.getExecutablePrerequisite(":bytecode_optimizer"));
+    ImmutableList<Artifact> localJavaOptimizationConfiguration =
+        ruleContext.getPrerequisiteArtifacts(":local_java_optimization_configuration").list();
+
     AndroidLintTool androidLint = AndroidLintTool.fromRuleContext(ruleContext);
 
     ImmutableList<JavaPackageConfigurationProvider> packageConfiguration =
@@ -168,6 +174,8 @@ public ConfiguredTarget create(RuleContext ruleContext)
             headerCompilerDirect,
             androidLint,
             jspecifyInfo,
+            bytecodeOptimizer,
+            localJavaOptimizationConfiguration,
             headerCompilerBuiltinProcessors,
             reducedClasspathIncompatibleProcessors,
             forciblyDisableHeaderCompilation,
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainProvider.java
index ebd6001def5876..342fa25389c837 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainProvider.java
@@ -100,6 +100,8 @@ public static JavaToolchainProvider create(
       @Nullable JavaToolchainTool headerCompilerDirect,
       @Nullable AndroidLintTool androidLint,
       JspecifyInfo jspecifyInfo,
+      @Nullable JavaToolchainTool bytecodeOptimizer,
+      ImmutableList<Artifact> localJavaOptimizationConfiguration,
       ImmutableSet<String> headerCompilerBuiltinProcessors,
       ImmutableSet<String> reducedClasspathIncompatibleProcessors,
       boolean forciblyDisableHeaderCompilation,
@@ -126,6 +128,8 @@ public static JavaToolchainProvider create(
         headerCompilerDirect,
         androidLint,
         jspecifyInfo,
+        bytecodeOptimizer,
+        localJavaOptimizationConfiguration,
         headerCompilerBuiltinProcessors,
         reducedClasspathIncompatibleProcessors,
         forciblyDisableHeaderCompilation,
@@ -158,6 +162,8 @@ public static JavaToolchainProvider create(
   @Nullable private final JavaToolchainTool headerCompilerDirect;
   @Nullable private final AndroidLintTool androidLint;
   @Nullable private final JspecifyInfo jspecifyInfo;
+  @Nullable private final JavaToolchainTool bytecodeOptimizer;
+  private final ImmutableList<Artifact> localJavaOptimizationConfiguration;
   private final ImmutableSet<String> headerCompilerBuiltinProcessors;
   private final ImmutableSet<String> reducedClasspathIncompatibleProcessors;
   private final boolean forciblyDisableHeaderCompilation;
@@ -190,6 +196,8 @@ private JavaToolchainProvider(
       @Nullable JavaToolchainTool headerCompilerDirect,
       @Nullable AndroidLintTool androidLint,
       @Nullable JspecifyInfo jspecifyInfo,
+      @Nullable JavaToolchainTool bytecodeOptimizer,
+      ImmutableList<Artifact> localJavaOptimizationConfiguration,
       ImmutableSet<String> headerCompilerBuiltinProcessors,
       ImmutableSet<String> reducedClasspathIncompatibleProcessors,
       boolean forciblyDisableHeaderCompilation,
@@ -221,6 +229,8 @@ private JavaToolchainProvider(
     this.headerCompilerDirect = headerCompilerDirect;
     this.androidLint = androidLint;
     this.jspecifyInfo = jspecifyInfo;
+    this.bytecodeOptimizer = bytecodeOptimizer;
+    this.localJavaOptimizationConfiguration = localJavaOptimizationConfiguration;
     this.headerCompilerBuiltinProcessors = headerCompilerBuiltinProcessors;
     this.reducedClasspathIncompatibleProcessors = reducedClasspathIncompatibleProcessors;
     this.forciblyDisableHeaderCompilation = forciblyDisableHeaderCompilation;
@@ -290,6 +300,15 @@ public JspecifyInfo jspecifyInfo() {
     return jspecifyInfo;
   }
 
+  @Nullable
+  public JavaToolchainTool getBytecodeOptimizer() {
+    return bytecodeOptimizer;
+  }
+
+  public ImmutableList<Artifact> getLocalJavaOptimizationConfiguration() {
+    return localJavaOptimizationConfiguration;
+  }
+
   /** Returns class names of annotation processors that are built in to the header compiler. */
   public ImmutableSet<String> getHeaderCompilerBuiltinProcessors() {
     return headerCompilerBuiltinProcessors;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainRule.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainRule.java
index 3c28fda4bb24a8..883b30aa7d2334 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainRule.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainRule.java
@@ -369,6 +369,15 @@ The Java target version (e.g., '6' or '7'). It specifies for which Java runtime
                 .allowedFileTypes()
                 .mandatoryProviders(ImmutableList.of(PackageGroupConfiguredTarget.PROVIDER.id()))
                 .undocumented("experimental"))
+        .add(
+            attr(":bytecode_optimizer", LABEL)
+                .cfg(ExecutionTransitionFactory.create())
+                .value(JavaSemantics.BYTECODE_OPTIMIZER)
+                .exec())
+        .add(
+            attr(":local_java_optimization_configuration", LABEL_LIST)
+                .cfg(ExecutionTransitionFactory.create())
+                .value(JavaSemantics.LOCAL_JAVA_OPTIMIZATION_CONFIGURATION))
         .build();
   }