BEGIN_PUBLIC

Basic wiring for optimizer tools which support baseline profile rewriting.
END_PUBLIC

Support rewriting of profiles in optimizer tools.

This implements the blaze portion of Phase 1 of what is described here:
https://docs.google.com/document/d/14ZXSjMK8V1slSCxigGp7B9tCRO_YeXXdNSNSeDzXSxU/edit?usp=sharing

RELNOTES: Support for optimizers rewriting baseline profiles.
PiperOrigin-RevId: 547910576
Change-Id: I4a66a045f9a34e8fbd6484426286a6005b6c9fe8

Author: Googler

src/main/java/com/google/devtools/build/lib/bazel/rules/android/BazelAndroidSemantics.java

@@ -122,7 +122,27 @@ public void registerMigrationRuleError(RuleContext ruleContext) throws RuleError
   /* Bazel does not currently support baseline profiles in the final apk.  */
   @Override
   public Artifact getArtProfileForApk(
-      RuleContext ruleContext, Artifact finalClassesDex, Artifact proguardOutputMap) {
+      RuleContext ruleContext,
+      Artifact finalClassesDex,
+      Artifact proguardOutputMap,
+      String baselineProfileDir) {
+    return null;
+  }
+
+  /* Bazel does not currently support baseline profiles in the final apk.  */
+  @Override
+  public Artifact compileBaselineProfile(
+      RuleContext ruleContext,
+      Artifact finalClassesDex,
+      Artifact proguardOutputMap,
+      Artifact mergedStaticProfile,
+      String baselineProfileDir) {
+    return null;
+  }
+
+  /* Bazel does not currently support baseline profiles in the final apk.  */
+  @Override
+  public Artifact mergeBaselineProfiles(RuleContext ruleContext, String baselineProfileDir) {
     return null;
   }
 }

src/main/java/com/google/devtools/build/lib/rules/android/AndroidBinary.java

@@ -531,6 +531,17 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
       }
     }
 
+    BaselineProfileProvider baselineprofileProvider =
+        ruleContext.getPrerequisite("application_resources", BaselineProfileProvider.PROVIDER);
+
+    Artifact baselineProfile = null;
+    String baselineProfileDir = ruleContext.getLabel().getName() + "-baseline-profile/";
+    if (baselineprofileProvider == null
+        && Allowlist.hasAllowlist(ruleContext, "allow_baseline_profiles_optimizer_integration")
+        && Allowlist.isAvailable(ruleContext, "allow_baseline_profiles_optimizer_integration")) {
+      baselineProfile = androidSemantics.mergeBaselineProfiles(ruleContext, baselineProfileDir);
+    }
+
     ProguardOutput proguardOutput =
         applyProguard(
             ruleContext,
@@ -539,7 +550,9 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
             binaryJar,
             proguardSpecs,
             proguardMapping,
-            proguardOutputMap);
+            proguardOutputMap,
+            baselineProfile,
+            baselineProfileDir);
 
     // Determine the outputs for the proguard map transition through post-processing and adding of
     // desugared library map.
@@ -682,13 +695,28 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
       proguardOutput.addAllToSet(filesBuilder, finalProguardOutputMap);
     }
 
-    BaselineProfileProvider baselineprofileProvider =
-        ruleContext.getPrerequisite("application_resources", BaselineProfileProvider.PROVIDER);
-    Artifact artProfileZip =
-        (baselineprofileProvider != null)
-            ? baselineprofileProvider.getArtProfileZip()
-            : androidSemantics.getArtProfileForApk(
-                ruleContext, finalClassesDex, finalProguardOutputMap);
+    Artifact artProfileZip = null;
+    if (baselineprofileProvider != null) {
+      // This happens when baseline profiles are provided via starlark.
+      artProfileZip = baselineprofileProvider.getArtProfileZip();
+    } else if (baselineProfile == null) {
+      // This happens when optimizer profile rewriting isn't enabled.
+      artProfileZip =
+          androidSemantics.getArtProfileForApk(
+              ruleContext, finalClassesDex, finalProguardOutputMap, baselineProfileDir);
+    } else {
+      // This happens when optimizer profile rewriting is enabled.
+      artProfileZip =
+          androidSemantics.compileBaselineProfile(
+              ruleContext,
+              finalClassesDex,
+              // Minified symbols are emitted when rewriting, so map isn't needed.
+              /* proguardOutputMap= */ null,
+              proguardOutput.getBaselineProfileOut() == null
+                  ? baselineProfile
+                  : proguardOutput.getBaselineProfileOut(),
+              baselineProfileDir);
+    }
 
     Artifact unsignedApk =
         ruleContext.getImplicitOutputArtifact(AndroidRuleClasses.ANDROID_BINARY_UNSIGNED_APK);
@@ -1054,7 +1082,9 @@ private static ProguardOutput applyProguard(
       Artifact deployJarArtifact,
       ImmutableList<Artifact> proguardSpecs,
       Artifact proguardMapping,
-      @Nullable Artifact proguardOutputMap)
+      @Nullable Artifact proguardOutputMap,
+      @Nullable Artifact baselineProfile,
+      String baselineProfileDir)
       throws InterruptedException {
     Artifact proguardOutputJar =
         ruleContext.getImplicitOutputArtifact(AndroidRuleClasses.ANDROID_BINARY_PROGUARD_JAR);
@@ -1101,7 +1131,9 @@ private static ProguardOutput applyProguard(
         proguardOutputJar,
         javaSemantics,
         getProguardOptimizationPasses(ruleContext),
-        proguardOutputMap);
+        proguardOutputMap,
+        baselineProfile,
+        baselineProfileDir);
   }
 
   @Nullable
@@ -1131,7 +1163,8 @@ private static ProguardOutput createEmptyProguardAction(
             ruleContext,
             semantics,
             proguardOutputMap,
-            null);
+            /* libraryJar= */ null,
+            /* baselineProfileOutput= */ null);
     outputs.addAllToSet(failures);
     ruleContext.registerAction(
         new FailAction(

src/main/java/com/google/devtools/build/lib/rules/android/AndroidSemantics.java

@@ -113,7 +113,21 @@ default void validateAndroidLibraryRuleContext(RuleContext ruleContext)
 
   /** The artifact for ART profile information. */
   Artifact getArtProfileForApk(
-      RuleContext ruleContext, Artifact finalClassesDex, Artifact proguardOutputMap);
+      RuleContext ruleContext,
+      Artifact finalClassesDex,
+      Artifact proguardOutputMap,
+      String baselineProfileDir);
+
+  /** The merged baseline profiles from the {@code baseline_profiles} attribute. */
+  Artifact mergeBaselineProfiles(RuleContext ruleContext, String baselineProfileDir);
+
+  /** The artifact for ART profile information, given a particular merged profile. */
+  Artifact compileBaselineProfile(
+      RuleContext ruleContext,
+      Artifact finalClassesDex,
+      Artifact proguardOutputMap,
+      Artifact mergedStaticProfile,
+      String baselineProfileDir);
 
   boolean postprocessClassesRewritesMap(RuleContext ruleContext);
 

src/main/java/com/google/devtools/build/lib/rules/android/ProguardHelper.java

@@ -63,6 +63,7 @@ public static final class ProguardOutput {
     @Nullable private final Artifact constantStringObfuscatedMapping;
     @Nullable private final Artifact libraryJar;
     private final Artifact config;
+    @Nullable private final Artifact baselineProfileOut;
 
     public ProguardOutput(
         Artifact outputJar,
@@ -72,7 +73,8 @@ public ProguardOutput(
         @Nullable Artifact usage,
         @Nullable Artifact constantStringObfuscatedMapping,
         @Nullable Artifact libraryJar,
-        Artifact config) {
+        Artifact config,
+        @Nullable Artifact baselineProfileOut) {
       this.outputJar = checkNotNull(outputJar);
       this.mapping = mapping;
       this.protoMapping = protoMapping;
@@ -81,10 +83,11 @@ public ProguardOutput(
       this.constantStringObfuscatedMapping = constantStringObfuscatedMapping;
       this.libraryJar = libraryJar;
       this.config = config;
+      this.baselineProfileOut = baselineProfileOut;
     }
 
     public static ProguardOutput createEmpty(Artifact outputJar) {
-      return new ProguardOutput(outputJar, null, null, null, null, null, null, null);
+      return new ProguardOutput(outputJar, null, null, null, null, null, null, null, null);
     }
 
     public Artifact getOutputJar() {
@@ -129,6 +132,10 @@ public Artifact getConfig() {
       return config;
     }
 
+    public Artifact getBaselineProfileOut() {
+      return baselineProfileOut;
+    }
+
     /** Adds the output artifacts to the given set builder. */
     public void addAllToSet(NestedSetBuilder<Artifact> filesBuilder) {
       addAllToSet(filesBuilder, null);
@@ -256,7 +263,8 @@ public static ProguardOutput getProguardOutputs(
       RuleContext ruleContext,
       JavaSemantics semantics,
       @Nullable Artifact proguardOutputMap,
-      @Nullable Artifact libraryJar)
+      @Nullable Artifact libraryJar,
+      @Nullable Artifact baselineProfileOutput)
       throws InterruptedException {
     boolean mappingRequested = genProguardMapping(ruleContext.attributes());
 
@@ -284,7 +292,8 @@ public static ProguardOutput getProguardOutputs(
         proguardUsage,
         proguardConstantStringMap,
         libraryJar,
-        proguardConfigOutput);
+        proguardConfigOutput,
+        baselineProfileOutput);
   }
 
   /**
@@ -304,6 +313,10 @@ public static ProguardOutput getProguardOutputs(
    * @param optimizationPasses if not null specifies to break proguard up into multiple passes with
    *     the given number of optimization passes.
    * @param proguardOutputMap mapping generated by Proguard if requested. could be null.
+   * @param baselineProfileIn Profile to pass the optimizer to perform feedback-directed
+   *     optimization.
+   * @param baselineProfileDir Directory to write baseline profile artifacts if baseline profile is
+   *     provided.
    */
   public static ProguardOutput createOptimizationActions(
       RuleContext ruleContext,
@@ -318,10 +331,11 @@ public static ProguardOutput createOptimizationActions(
       Artifact proguardOutputJar,
       JavaSemantics semantics,
       @Nullable Integer optimizationPasses,
-      @Nullable Artifact proguardOutputMap)
+      @Nullable Artifact proguardOutputMap,
+      @Nullable Artifact baselineProfileIn,
+      String baselineProfileDir)
       throws InterruptedException {
     Preconditions.checkArgument(!proguardSpecs.isEmpty());
-
     Artifact libraryJar = null;
 
     if (!libraryJars.isEmpty() && !libraryJars.isSingleton()) {
@@ -359,7 +373,10 @@ public static ProguardOutput createOptimizationActions(
       libraryJars = NestedSetBuilder.create(Order.STABLE_ORDER, filteredLibraryJar);
       libraryJar = filteredLibraryJar;
     }
-
+    Artifact baselineProfileOut = null;
+    if (baselineProfileIn != null) {
+      baselineProfileOut = ruleContext.getBinArtifact(baselineProfileDir + "rewritten-prof.txt");
+    }
     ProguardOutput output =
         getProguardOutputs(
             proguardOutputJar,
@@ -368,7 +385,8 @@ public static ProguardOutput createOptimizationActions(
             ruleContext,
             semantics,
             proguardOutputMap,
-            libraryJar);
+            libraryJar,
+            baselineProfileOut);
 
     JavaConfiguration javaConfiguration =
         ruleContext.getConfiguration().getFragment(JavaConfiguration.class);
@@ -394,6 +412,8 @@ public static ProguardOutput createOptimizationActions(
           output.getUsage(),
           output.getConstantStringObfuscatedMapping(),
           output.getConfig(),
+          baselineProfileIn,
+          baselineProfileOut,
           mnemonic);
       proguardAction
           .setProgressMessage("Trimming binary with %s: %s", mnemonic, ruleContext.getLabel())
@@ -435,6 +455,8 @@ public static ProguardOutput createOptimizationActions(
           /* proguardUsage */ null,
           /* constantStringObfuscatedMapping */ null,
           /* proguardConfigOutput */ null,
+          baselineProfileIn,
+          /* baselineProfileOutput */ null,
           mnemonic);
       initialAction
           .setProgressMessage("Trimming binary with %s: Verification/Shrinking Pass", mnemonic)
@@ -515,6 +537,8 @@ public static ProguardOutput createOptimizationActions(
           output.getUsage(),
           output.getConstantStringObfuscatedMapping(),
           output.getConfig(),
+          /* baselineProfileInput */ null,
+          baselineProfileOut,
           mnemonic);
       finalAction
           .setProgressMessage(
@@ -576,6 +600,8 @@ private static Artifact createSingleOptimizationAction(
         /* proguardUsage */ null,
         /* constantStringObfuscatedMapping */ null,
         /* proguardConfigOutput */ null,
+        /* baselineProfileInput */ null,
+        /* baselineProfileOutput */ null,
         mnemonic);
     optimizationAction
         .setProgressMessage(
@@ -609,6 +635,8 @@ private static void defaultAction(
       @Nullable Artifact proguardUsage,
       @Nullable Artifact constantStringObfuscatedMapping,
       @Nullable Artifact proguardConfigOutput,
+      @Nullable Artifact baselineProfileInput,
+      @Nullable Artifact baselineProfileOutput,
       String mnemonic) {
 
     builder
@@ -680,6 +708,14 @@ private static void defaultAction(
       builder.addOutput(proguardConfigOutput);
       commandLine.addExecPath("-printconfiguration", proguardConfigOutput);
     }
+    if (baselineProfileInput != null) {
+      builder.addInput(baselineProfileInput);
+      commandLine.addExecPath("-baselineprofile", baselineProfileInput);
+    }
+    if (baselineProfileOutput != null) {
+      builder.addOutput(baselineProfileOutput);
+      commandLine.addExecPath("-printbaselineprofile", baselineProfileOutput);
+    }
   }
 
   /** Returns an intermediate artifact used to run Proguard. */