From 4d719e1b1a9db31c2ea18dc966dbe8be05550b05 Mon Sep 17 00:00:00 2001
From: Andrew Sinclair <asinclair@google.com>
Date: Tue, 29 Oct 2024 15:16:06 -0700
Subject: [PATCH] Merge the desugar globals classes into the final APK in a
 standalone dex shard.

This treats it similarly to the Java 8 legacy dex. Adding the globals during dexmerging was incorrect. It attempted to merge the globals into each shard. The shards are carefully crafted by the DexFileSplitter to avoid going over dex method limits. This led to build failures where we exceeded the limits. It also meant we were merging the globals into every shard instead of just once per APK.

A nice side effect is that this removes the dependency on Bazel 7.4.0+ since we do not need the additional parameter in the android_common.create_dex_merger_actions call.

PiperOrigin-RevId: 691179644
Change-Id: Idd994c1a59286febbfd4398d1e1cd5721198793d
---
 rules/android_binary/impl.bzl |  8 ++++++--
 rules/dex.bzl                 | 31 ++++++++++---------------------
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/rules/android_binary/impl.bzl b/rules/android_binary/impl.bzl
index 21f35051..50d915b5 100644
--- a/rules/android_binary/impl.bzl
+++ b/rules/android_binary/impl.bzl
@@ -406,11 +406,15 @@ def _process_dex(ctx, validation_ctx, packaged_resources_ctx, deploy_ctx, bp_ctx
                 toolchain_type = ANDROID_TOOLCHAIN_TYPE,
             )
 
-        _dex.append_java8_legacy_dex(
+        dexes_to_append = []
+        if acls.in_record_desugaring_rollout(str(ctx.label)) and not is_binary_optimized:
+            dexes_to_append.append(utils.only(get_android_toolchain(ctx).desugar_globals_dex_archive.files.to_list()))
+        dexes_to_append.append(java8_legacy_dex)
+        _dex.append_desugar_dexes(
             ctx,
             output = final_classes_dex_zip,
             input = classes_dex_zip,
-            java8_legacy_dex = java8_legacy_dex,
+            dexes = dexes_to_append,
             dex_zips_merger = get_android_toolchain(ctx).dex_zips_merger.files_to_run,
         )
     else:
diff --git a/rules/dex.bzl b/rules/dex.bzl
index 35b8a3c2..9925c727 100644
--- a/rules/dex.bzl
+++ b/rules/dex.bzl
@@ -189,22 +189,16 @@ def _process_incremental_dexing(
             toolchain_type = toolchain_type,
         )
 
-        # TODO(asinclair): Remove this once the rollout is complete.
-        dex_merger_actions_args = dict(
+        # TODO(b/130571505): Implement this after SpawnActionTemplate is supported in Starlark
+        android_common.create_dex_merger_actions(
+            ctx,
             output = dexes,
             input = shards,
             dexopts = dexopts,
             dexmerger = dexmerger,
             min_sdk_version = min_sdk_version,
         )
-        if acls.in_record_desugaring_rollout(str(ctx.label)):
-            dex_merger_actions_args["desugar_globals"] = utils.only(get_android_toolchain(ctx).desugar_globals.files.to_list())
 
-        # TODO(b/130571505): Implement this after SpawnActionTemplate is supported in Starlark
-        android_common.create_dex_merger_actions(
-            ctx,
-            **dex_merger_actions_args
-        )
         _java.singlejar(
             ctx,
             output = output,
@@ -379,27 +373,22 @@ def _shard_dexes(
 
     return output
 
-def _append_java8_legacy_dex(
-        ctx,
-        output = None,
-        input = None,
-        java8_legacy_dex = None,
-        dex_zips_merger = None):
+def _append_desugar_dexes(ctx, output = None, input = None, dexes = None, dex_zips_merger = None):
     args = ctx.actions.args()
 
-    # Order matters here: we want java8_legacy_dex to be the highest-numbered classesN.dex
+    # Order matters here: we want the additional dex(s) to be the highest-numbered classesN.dex
     args.add("--input_zip", input)
-    args.add("--input_zip", java8_legacy_dex)
+    args.add_all(dexes, before_each = "--input_zip")
     args.add("--output_zip", output)
 
     ctx.actions.run(
         executable = dex_zips_merger,
-        inputs = [input, java8_legacy_dex],
+        inputs = [input] + dexes,
         outputs = [output],
         arguments = [args],
-        mnemonic = "AppendJava8LegacyDex",
         use_default_shell_env = True,
-        progress_message = "Adding Java8 legacy library for %s" % ctx.label,
+        mnemonic = "AppendDesugarDexes",
+        progress_message = "Adding Desugar dex file(s) for %s" % ctx.label,
         toolchain = ANDROID_TOOLCHAIN_TYPE,
     )
 
@@ -842,7 +831,7 @@ def _transform_dex_list_through_proguard_map(
     return obfuscated_main_dex_list
 
 dex = struct(
-    append_java8_legacy_dex = _append_java8_legacy_dex,
+    append_desugar_dexes = _append_desugar_dexes,
     dex = _dex,
     dex_merge = _dex_merge,
     generate_main_dex_list = _generate_main_dex_list,