Migrate baseline profiles rewriting, wildcards expansion and startup …

…profiles support in optimizer tools from native to Starlark.

PiperOrigin-RevId: 565733260
Change-Id: I5f62382792f0e3b83339750337961f56fe218cf4

rules/acls/baseline_profiles_optimizer_integration.bzl

@@ -16,4 +16,5 @@
 
 # keep sorted
 BASELINE_PROFILES_OPTIMIZER_INTEGRATION = [
+    "//test/rules/android_binary_internal:__subpackages__",
 ]

rules/android_binary_internal/attrs.bzl

@@ -81,6 +81,19 @@ ATTRS = _attrs.replace(
                 profile, speeding up app startup time or reducing jank in some circumstances.
                 """,
             ),
+            startup_profiles = attr.label_list(
+                allow_empty = True,
+                allow_files = [".txt"],
+                doc = """
+                List of baseline profiles that were collected at runtime (often from start-up) for
+                this binary. When this is specified, all baseline profiles (including these) are
+                used to inform code optimizations in the build toolchain. This may improve runtime
+                performance at the cost of dex size. If the dex size cost is too large and the
+                performance wins too small, the same profiles can be provided as a dep from an
+                android_library with `baseline_profiles` to avoid the runtime-focused code
+                optimizations that are enabled by `startup_profiles`.
+                """,
+            ),
             proguard_specs = attr.label_list(allow_empty = True, allow_files = True),
             resource_apks = attr.label_list(
                 allow_rules = ["apk_import"],

rules/android_binary_internal/impl.bzl

@@ -219,7 +219,7 @@ def _process_build_info(_unused_ctx, **unused_ctxs):
         ),
     )
 
-def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, deploy_ctx, optimize_ctx, **_unused_ctxs):
+def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, deploy_ctx, bp_ctx, optimize_ctx, **_unused_ctxs):
     providers = []
     classes_dex_zip = None
     dex_info = None
@@ -231,6 +231,7 @@ def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, dep
     main_dex_list = ctx.file.main_dex_list
     multidex = ctx.attr.multidex
     optimizing_dexer = ctx.attr._optimizing_dexer
+    java8_legacy_dex_map = None
 
     if acls.in_android_binary_starlark_dex_desugar_proguard(str(ctx.label)):
         proguarded_jar = optimize_ctx.proguard_output.output_jar if is_binary_optimized else None
@@ -320,8 +321,7 @@ def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, dep
                 min_sdk_version = ctx.attr.min_sdk_version,
                 main_dex_list = main_dex_list,
                 library_jar = optimize_ctx.proguard_output.library_jar,
-                # TODO(b/286955442): Support baseline profiles.
-                startup_profile = None,
+                startup_profile = bp_ctx.baseline_profile_output.startup_profile if bp_ctx.baseline_profile_output else None,
                 optimizing_dexer = optimizing_dexer.files_to_run,
                 toolchain_type = ANDROID_TOOLCHAIN_TYPE,
             )
@@ -398,6 +398,7 @@ def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, dep
         name = "dex_ctx",
         value = struct(
             dex_info = dex_info,
+            java8_legacy_dex_map = java8_legacy_dex_map,
             providers = providers,
         ),
     )
@@ -557,11 +558,20 @@ def _is_instrumentation(ctx):
     """
     return bool(ctx.attr.instruments)
 
-def _process_baseline_profiles(ctx, dex_ctx, **_unused_ctxs):
-    providers = []
-    if (ctx.attr.generate_art_profile and
+def _process_baseline_profiles(ctx, deploy_ctx, **_unused_ctxs):
+    baseline_profile_output = struct()
+    if (ctx.attr.generate_art_profile or
         acls.in_android_binary_starlark_dex_desugar_proguard(str(ctx.label))):
+        enable_optimizer_integration = acls.in_baseline_profiles_optimizer_integration(str(ctx.label))
+        has_proguard_specs = bool(ctx.files.proguard_specs)
+
+        if ctx.files.startup_profiles and not enable_optimizer_integration:
+            fail("Target %s is not allowed to set startup_profiles." % str(ctx.label))
+
+        # Include startup profiles if the optimizer is disabled since profiles won't be merged
+        # in the optimizer.
         transitive_profiles = depset(
+            ctx.files.startup_profiles if enable_optimizer_integration and not has_proguard_specs else [],
             transitive = [
                 profile_provider.files
                 for profile_provider in utils.collect_providers(
@@ -570,20 +580,57 @@ def _process_baseline_profiles(ctx, dex_ctx, **_unused_ctxs):
                 )
             ],
         )
-        if transitive_profiles:
-            providers.append(
-                _baseline_profiles.process(
-                    ctx,
-                    dex_ctx.dex_info.final_classes_dex_zip,
-                    transitive_profiles,
-                ),
-            )
+        baseline_profile_output = _baseline_profiles.process(
+            ctx,
+            transitive_profiles = transitive_profiles,
+            startup_profiles = ctx.files.startup_profiles,
+            deploy_jar = deploy_ctx.deploy_jar,
+            has_proguard_specs = has_proguard_specs,
+            enable_optimizer_integration = enable_optimizer_integration,
+            merge_tool = get_android_toolchain(ctx).merge_baseline_profiles_tool.files_to_run,
+            profgen = get_android_toolchain(ctx).profgen.files_to_run,
+            toolchain_type = ANDROID_TOOLCHAIN_TYPE,
+        )
     return ProviderInfo(
         name = "bp_ctx",
+        value = struct(
+            baseline_profile_output = baseline_profile_output,
+        ),
+    )
+
+def _process_art_profile(ctx, bp_ctx, dex_ctx, optimize_ctx, **_unused_ctxs):
+    providers = []
+    if (ctx.attr.generate_art_profile and
+        acls.in_android_binary_starlark_dex_desugar_proguard(str(ctx.label))):
+        merged_baseline_profile = bp_ctx.baseline_profile_output.baseline_profile
+        merged_baseline_profile_rewritten = \
+            optimize_ctx.proguard_output.baseline_profile_rewritten if optimize_ctx.proguard_output else None
+        proguard_output_map = dex_ctx.dex_info.final_proguard_output_map
+
+        if acls.in_baseline_profiles_optimizer_integration(str(ctx.label)):
+            # Minified symbols are emitted when rewriting, so only use map for symbols which
+            # weren't passed to bytecode optimizer (if it exists).
+            proguard_output_map = dex_ctx.java8_legacy_dex_map
+
+            # At this point, either baseline profile here also contains startup-profiles, if any.
+            if merged_baseline_profile_rewritten:
+                merged_baseline_profile = merged_baseline_profile_rewritten
+        if merged_baseline_profile:
+            providers.append(_baseline_profiles.process_art_profile(
+                ctx,
+                final_classes_dex = dex_ctx.dex_info.final_classes_dex_zip,
+                merged_profile = merged_baseline_profile,
+                proguard_output_map = proguard_output_map,
+                profgen = get_android_toolchain(ctx).profgen.files_to_run,
+                zipper = get_android_toolchain(ctx).zipper.files_to_run,
+                toolchain_type = ANDROID_TOOLCHAIN_TYPE,
+            ))
+    return ProviderInfo(
+        name = "ap_ctx",
         value = struct(providers = providers),
     )
 
-def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, **_unused_ctxs):
+def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, bp_ctx, **_unused_ctxs):
     if not acls.in_android_binary_starlark_dex_desugar_proguard(str(ctx.label)):
         return ProviderInfo(
             name = "optimize_ctx",
@@ -636,6 +683,9 @@ def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, **_unused_ctxs):
     proguard_seeds = ctx.actions.declare_file(ctx.label.name + "_migrated_proguard.seeds")
     proguard_usage = ctx.actions.declare_file(ctx.label.name + "_migrated_proguard.usage")
 
+    startup_profile = bp_ctx.baseline_profile_output.startup_profile if bp_ctx.baseline_profile_output else None
+    baseline_profile = bp_ctx.baseline_profile_output.baseline_profile if bp_ctx.baseline_profile_output else None
+
     proguard_output = proguard.apply_proguard(
         ctx,
         input_jar = deploy_ctx.deploy_jar,
@@ -646,6 +696,8 @@ def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, **_unused_ctxs):
         proguard_output_map = proguard_output_map,
         proguard_seeds = proguard_seeds,
         proguard_usage = proguard_usage,
+        startup_profile = startup_profile,
+        baseline_profile = baseline_profile,
         proguard_tool = get_android_sdk(ctx).proguard,
     )
 
@@ -659,6 +711,8 @@ def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, **_unused_ctxs):
             usage = proguard_output.usage,
             library_jar = proguard_output.library_jar,
             config = proguard_output.config,
+            baseline_profile_rewritten = proguard_output.baseline_profile_rewritten,
+            startup_profile_rewritten = proguard_output.startup_profile_rewritten,
         ))
 
     use_resource_shrinking = is_resource_shrinking_enabled and has_proguard_specs
@@ -712,9 +766,10 @@ PROCESSORS = dict(
     BuildInfoProcessor = _process_build_info,
     ProtoProcessor = _process_proto,
     DeployJarProcessor = _process_deploy_jar,
+    BaselineProfilesProcessor = _process_baseline_profiles,
     OptimizeProcessor = _process_optimize,
     DexProcessor = _process_dex,
-    BaselineProfilesProcessor = _process_baseline_profiles,
+    ArtProfileProcessor = _process_art_profile,
     R8Processor = process_r8,
     ResourecShrinkerR8Processor = process_resource_shrinking_r8,
 )