BEGIN_PUBLIC

Expose CreateTemplatedMergerActions method to Starlark
END_PUBLIC

PiperOrigin-RevId: 554579309
Change-Id: I0cfb447259f3209e519caa08c5468866eacf65fa

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

@@ -1666,7 +1666,11 @@ private static void createIncrementalDexingActions(
     } else {
       SpecialArtifact shardsToMerge =
           createSharderAction(ruleContext, dexArchives, mainDexList, dexopts, inclusionFilterJar);
-      Artifact multidexShards = createTemplatedMergerActions(ruleContext, shardsToMerge, dexopts);
+      SpecialArtifact multidexShards =
+          ruleContext.getTreeArtifact(
+              ruleContext.getUniqueDirectory("dexfiles"), ruleContext.getBinOrGenfilesDirectory());
+      FilesToRunProvider dexMerger = ruleContext.getExecutablePrerequisite("$dexmerger");
+      createTemplatedMergerActions(ruleContext, multidexShards, shardsToMerge, dexopts, dexMerger);
       // TODO(b/69431301): avoid this action and give the files to apk build action directly
       createZipMergeAction(ruleContext, multidexShards, classesDex);
     }
@@ -1790,18 +1794,17 @@ private static SpecialArtifact createSharderAction(
 
   /**
    * Sets up a monodex {@code $dexmerger} actions for each dex archive in the given tree artifact
-   * and returns the output tree artifact.
-   *
-   * @return Tree artifact containing zips with final dex files named for inclusion in an APK.
+   * and puts the outputs in a tree artifact.
    */
-  private static Artifact createTemplatedMergerActions(
-      RuleContext ruleContext, SpecialArtifact inputTree, Collection<String> dexopts) {
-    SpecialArtifact outputTree =
-        ruleContext.getTreeArtifact(
-            ruleContext.getUniqueDirectory("dexfiles"), ruleContext.getBinOrGenfilesDirectory());
+  public static void createTemplatedMergerActions(
+      RuleContext ruleContext,
+      SpecialArtifact outputTree,
+      SpecialArtifact inputTree,
+      List<String> dexopts,
+      FilesToRunProvider executable) {
     SpawnActionTemplate.Builder dexmerger =
         new SpawnActionTemplate.Builder(inputTree, outputTree)
-            .setExecutable(ruleContext.getExecutablePrerequisite("$dexmerger"))
+            .setExecutable(executable)
             .setMnemonics("DexShardsToMerge", "DexMerger")
             .setOutputPathMapper(
                 (OutputPathMapper & Serializable) TreeFileArtifact::getParentRelativePath);
@@ -1817,8 +1820,6 @@ private static Artifact createTemplatedMergerActions(
                         dexopts, Predicates.not(Predicates.equalTo(DX_MINIMAL_MAIN_DEX_OPTION)))));
     dexmerger.setCommandLineTemplate(commandLine.build());
     ruleContext.registerAction(dexmerger.build(ruleContext.getActionOwner()));
-
-    return outputTree;
   }
 
   private static void createZipMergeAction(

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

@@ -14,16 +14,24 @@
 package com.google.devtools.build.lib.rules.android;
 
 import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
+import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.config.transitions.StarlarkExposedRuleTransitionFactory;
+import com.google.devtools.build.lib.analysis.platform.ConstraintValueInfo;
+import com.google.devtools.build.lib.analysis.starlark.StarlarkRuleContext;
 import com.google.devtools.build.lib.packages.Info;
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
 import com.google.devtools.build.lib.rules.java.JavaCompilationArgsProvider;
 import com.google.devtools.build.lib.rules.java.JavaInfo;
 import com.google.devtools.build.lib.starlarkbuildapi.android.AndroidSplitTransitionApi;
 import com.google.devtools.build.lib.starlarkbuildapi.android.AndroidStarlarkCommonApi;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
 
 /** Common utilities for Starlark rules related to Android. */
-public class AndroidStarlarkCommon implements AndroidStarlarkCommonApi<Artifact, JavaInfo> {
+public class AndroidStarlarkCommon
+    implements AndroidStarlarkCommonApi<
+        Artifact, JavaInfo, FilesToRunProvider, ConstraintValueInfo, StarlarkRuleContext> {
 
   @Override
   public AndroidDeviceBrokerInfo createDeviceBrokerInfo(String deviceBrokerType) {
@@ -70,4 +78,20 @@ public JavaInfo enableImplicitSourcelessDepsExportsCompatibility(Info javaInfo,
                     JavaCompilationArgsProvider.builder().addExports(args, type).build()));
     return builder.setNeverlink(neverlink).build();
   }
+
+  @Override
+  public void createDexMergerActions(
+      StarlarkRuleContext starlarkRuleContext,
+      Artifact output,
+      Artifact input,
+      Sequence<?> dexopts, // <String> expected.
+      FilesToRunProvider dexmerger)
+      throws EvalException, RuleErrorException {
+    AndroidBinary.createTemplatedMergerActions(
+        starlarkRuleContext.getRuleContext(),
+        (SpecialArtifact) output,
+        (SpecialArtifact) input,
+        Sequence.cast(dexopts, String.class, "dexopts"),
+        dexmerger);
+  }
 }

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidBootstrap.java

@@ -37,11 +37,11 @@ public class AndroidBootstrap implements Bootstrap {
           PackageIdentifier.createUnchecked("rules_android", ""),
           PackageIdentifier.createUnchecked("", "tools/build_defs/android"));
 
-  private final AndroidStarlarkCommonApi<?, ?> androidCommon;
+  private final AndroidStarlarkCommonApi<?, ?, ?, ?, ?> androidCommon;
   private final ImmutableMap<String, Object> providers;
 
   public AndroidBootstrap(
-      AndroidStarlarkCommonApi<?, ?> androidCommon,
+      AndroidStarlarkCommonApi<?, ?, ?, ?, ?> androidCommon,
       ApkInfoApiProvider apkInfoProvider,
       AndroidInstrumentationInfoApiProvider<?> androidInstrumentationInfoProvider,
       AndroidDeviceBrokerInfoApiProvider androidDeviceBrokerInfoProvider,

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidStarlarkCommonApi.java

@@ -18,11 +18,17 @@
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.starlarkbuildapi.FileApi;
+import com.google.devtools.build.lib.starlarkbuildapi.FilesToRunProviderApi;
+import com.google.devtools.build.lib.starlarkbuildapi.StarlarkRuleContextApi;
 import com.google.devtools.build.lib.starlarkbuildapi.java.JavaInfoApi;
+import com.google.devtools.build.lib.starlarkbuildapi.platform.ConstraintValueInfoApi;
 import javax.annotation.Nullable;
 import net.starlark.java.annot.Param;
+import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.StarlarkValue;
 
 /** Common utilities for Starlark rules related to Android. */
@@ -34,7 +40,11 @@
             + "Common utilities and functionality related to Android rules.",
     documented = false)
 public interface AndroidStarlarkCommonApi<
-        FileT extends FileApi, JavaInfoT extends JavaInfoApi<?, ?, ?>>
+        FileT extends FileApi,
+        JavaInfoT extends JavaInfoApi<?, ?, ?>,
+        FilesToRunProviderT extends FilesToRunProviderApi<FileT>,
+        ConstraintValueT extends ConstraintValueInfoApi,
+        StarlarkRuleContextT extends StarlarkRuleContextApi<ConstraintValueT>>
     extends StarlarkValue {
 
   @StarlarkMethod(
@@ -102,4 +112,48 @@ public interface AndroidStarlarkCommonApi<
       })
   JavaInfoT enableImplicitSourcelessDepsExportsCompatibility(Info javaInfo, boolean neverlink)
       throws RuleErrorException;
+
+  @StarlarkMethod(
+      name = "create_dex_merger_actions",
+      doc =
+          "Creates a list of DexMerger actions to be run in parallel, each action taking one shard"
+              + " from the input directory, merging all the dex archives inside the shard to a"
+              + " single dexarchive under the output directory.",
+      documented = false,
+      enableOnlyWithFlag = BuildLanguageOptions.EXPERIMENTAL_ENABLE_ANDROID_MIGRATION_APIS,
+      parameters = {
+        @Param(name = "ctx", doc = "The rule context.", positional = true, named = false),
+        @Param(
+            name = "output",
+            doc = "The output directory.",
+            positional = false,
+            named = true,
+            allowedTypes = {@ParamType(type = FileApi.class)}),
+        @Param(
+            name = "input",
+            doc = "The input directory.",
+            positional = false,
+            named = true,
+            allowedTypes = {@ParamType(type = FileApi.class)}),
+        @Param(
+            name = "dexopts",
+            doc = "A list of additional command-line flags for the dx tool. Optional",
+            positional = false,
+            named = true,
+            allowedTypes = {@ParamType(type = Sequence.class, generic1 = String.class)},
+            defaultValue = "[]"),
+        @Param(
+            name = "dexmerger",
+            doc = "A FilesToRunProvider to be used for dex merging.",
+            positional = false,
+            named = true,
+            allowedTypes = {@ParamType(type = FilesToRunProviderApi.class)})
+      })
+  void createDexMergerActions(
+      StarlarkRuleContextT starlarkRuleContext,
+      FileT output,
+      FileT input,
+      Sequence<?> dexopts, // <String> expected.
+      FilesToRunProviderT dexmerger)
+      throws EvalException, RuleErrorException;
 }

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/BUILD

@@ -33,6 +33,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/core",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/cpp",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/java",
+        "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/platform",
         "//src/main/java/net/starlark/java/annot",
         "//src/main/java/net/starlark/java/eval",
         "//third_party:guava",