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

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.

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.

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.

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%";

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,

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;

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();
   }