Wire up AndroidOptimizationInfo in Starlark

PiperOrigin-RevId: 570897571
Change-Id: I82932e67d790caf8db65608d7115b9ab8ee1f8aa

2 files changed, 36 insertions, 38 deletions

diff --git a/rules/android_binary_internal/impl.bzl b/rules/android_binary_internal/impl.bzl
index 3b3516e..e49b628 100644
--- a/rules/android_binary_internal/impl.bzl
+++ b/rules/android_binary_internal/impl.bzl
@@ -31,7 +31,7 @@ load(
     "ProviderInfo",
     "processing_pipeline",
 )
-load("//rules:proguard.bzl", "proguard", proguard_testing = "testing")
+load("//rules:proguard.bzl", "proguard")
 load("//rules:providers.bzl", "StarlarkAndroidDexInfo", "StarlarkApkInfo")
 load("//rules:resources.bzl", _resources = "resources")
 load(
@@ -389,6 +389,7 @@ def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, dep
 
         dex_info = AndroidDexInfo(
             deploy_jar = deploy_jar,
+            filtered_deploy_jar = deploy_ctx.filtered_deploy_jar,
             final_classes_dex_zip = final_classes_dex_zip,
             final_proguard_output_map = final_proguard_output_map,
             java_resource_jar = binary_jar if ctx.fragments.android.get_java_resources_from_optimized_jar else deploy_jar,
@@ -409,7 +410,7 @@ def _process_dex(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, proto_ctx, dep
     )
 
 def _process_deploy_jar(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, build_info_ctx, proto_ctx, **_unused_ctxs):
-    deploy_jar, desugar_dict = None, {}
+    deploy_jar, filtered_deploy_jar, desugar_dict = None, None, {}
 
     if acls.in_android_binary_starlark_dex_desugar_proguard(str(ctx.label)):
         java_toolchain = common.get_java_toolchain(ctx)
@@ -460,7 +461,9 @@ def _process_deploy_jar(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, build_i
 
         if _is_instrumentation(ctx):
             filtered_deploy_jar = ctx.actions.declare_file(ctx.label.name + "_migrated_filtered.jar")
-            filter_jar = ctx.attr.instruments[AndroidPreDexJarInfo].pre_dex_jar
+
+            # TODO(b/303286042): Use AndroidPreDexInfo.pre_dex_jar to be the filter_jar
+            filter_jar = ctx.attr.instruments[ApkInfo].deploy_jar
             common.filter_zip_exclude(
                 ctx,
                 output = filtered_deploy_jar,
@@ -478,6 +481,7 @@ def _process_deploy_jar(ctx, stamp_ctx, packaged_resources_ctx, jvm_ctx, build_i
         value = struct(
             deploy_jar = deploy_jar,
             desugar_dict = desugar_dict,
+            filtered_deploy_jar = filtered_deploy_jar,
             providers = [],
         ),
     )
@@ -709,20 +713,6 @@ def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, bp_ctx, **_unused
         proguard_tool = get_android_sdk(ctx).proguard,
     )
 
-    providers = []
-    if proguard_output:
-        providers.append(proguard_testing.ProguardOutputInfo(
-            input_jar = deploy_ctx.deploy_jar,
-            output_jar = proguard_output.output_jar,
-            mapping = proguard_output.mapping,
-            seeds = proguard_output.seeds,
-            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
     shrunk_resource_output = None
     if use_resource_shrinking:
@@ -737,7 +727,6 @@ def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, bp_ctx, **_unused
             busybox = get_android_toolchain(ctx).android_resources_busybox.files_to_run,
             host_javabase = common.get_host_javabase(ctx),
         )
-        providers.append(shrunk_resource_output)
 
     optimized_resource_output = _resources.optimize(
         ctx,
@@ -748,7 +737,26 @@ def _process_optimize(ctx, deploy_ctx, packaged_resources_ctx, bp_ctx, **_unused
         busybox = get_android_toolchain(ctx).android_resources_busybox.files_to_run,
         host_javabase = common.get_host_javabase(ctx),
     )
-    providers.append(optimized_resource_output)
+
+    providers = []
+    providers.append(
+        AndroidOptimizationInfo(
+            optimized_jar = proguard_output.output_jar if proguard_output else None,
+            mapping = proguard_output.mapping if proguard_output else None,
+            seeds = proguard_output.seeds if proguard_output else None,
+            library_jar = proguard_output.library_jar if proguard_output else None,
+            config = proguard_output.config if proguard_output else None,
+            proto_mapping = proguard_output.proto_mapping if proguard_output else None,
+            rewritten_startup_profile = proguard_output.startup_profile_rewritten if proguard_output else None,
+            rewriten_merged_baseline_profile = proguard_output.baseline_profile_rewritten if proguard_output else None,
+            optimized_resource_apk = optimized_resource_output.resources_apk,
+            shrunk_resource_apk = shrunk_resource_output.resources_apk if shrunk_resource_output else None,
+            shrunk_resource_zip = shrunk_resource_output.resources_zip if shrunk_resource_output else None,
+            resource_shrinker_log = shrunk_resource_output.shrinker_log if shrunk_resource_output else None,
+            resource_optimization_config = shrunk_resource_output.optimization_config if shrunk_resource_output else None,
+            resource_path_shortening_map = optimized_resource_output.path_shortening_map,
+        ),
+    )
 
     return ProviderInfo(
         name = "optimize_ctx",
diff --git a/rules/proguard.bzl b/rules/proguard.bzl
index 5a8c420..374e2f8 100644
--- a/rules/proguard.bzl
+++ b/rules/proguard.bzl
@@ -30,23 +30,6 @@ _ProguardSpecContextInfo = provider(
     ),
 )
 
-_ProguardOutputInfo = provider(
-    doc = "Temporary provider to hold all proguard outputs. Will be replaced by a native  " +
-          "provider. Useful for testing.",
-    fields = dict(
-        input_jar = "The input program jar, unoptimized",
-        output_jar = "The optimized output jar",
-        mapping = "Output proguard map",
-        proto_mapping = "Output proto mapping",
-        seeds = "Output seeds",
-        usage = "Output usage",
-        library_jar = "Merged library jar",
-        config = "Output config",
-        baseline_profile_rewritten = "Optimized baseline profile",
-        startup_profile_rewritten = "Optimized startup profile",
-    ),
-)
-
 def _validate_proguard_spec(
         ctx,
         out_validated_proguard_spec,
@@ -369,7 +352,7 @@ def _apply_proguard(
       proguard_tool: FilesToRun. The proguard executable.
 
     Returns:
-      A struct of proguard outputs, corresponding to the fields in ProguardOutputInfo.
+      A struct of proguard outputs.
     """
     if not proguard_specs:
         # Fail at execution time if these artifacts are requested, to avoid issue where outputs are
@@ -416,11 +399,19 @@ def _get_proguard_output(
         startup_profile_rewritten,
         baseline_profile_rewritten):
     """Helper method to get a struct of all proguard outputs."""
+
+    # Proto Output Map is currently empty from ProGuard.
+    proguard_output_proto_map = None
+    if proguard_output_map:
+        proguard_output_proto_map = _get_proguard_temp_artifact(ctx, "_proguard.pbmap")
+        ctx.actions.write(proguard_output_proto_map, content = "")
+
     config_output = _get_proguard_temp_artifact(ctx, "_proguard.config")
 
     return struct(
         output_jar = proguard_output_jar,
         mapping = proguard_output_map,
+        proto_mapping = proguard_output_proto_map,
         seeds = proguard_seeds,
         usage = proguard_usage,
         library_jar = combined_library_jar,
@@ -678,5 +669,4 @@ testing = struct(
     collect_transitive_proguard_specs = _collect_transitive_proguard_specs,
     optimization_action = _optimization_action,
     ProguardSpecContextInfo = _ProguardSpecContextInfo,
-    ProguardOutputInfo = _ProguardOutputInfo,
 )