Allow tweaking jvm flags per mnemonic

rules/aar_import/attrs.bzl

@@ -73,6 +73,12 @@ ATTRS = _attrs.add(
         _manifest_merge_order = attr.label(
             default = "//rules/flags:manifest_merge_order",
         ),
+        _mnemonic_jvm_flags = attr.label(
+            default = "//rules/flags:mnemonic_jvm_flags",
+            doc = "Additional JVM flags to set depending on the mnemonic " +
+                  "E.g --//rules/flags:mnemonic_jvm_flags=CompileAndroidResources=Xms1G " +
+                  " will add -Xms1G to all invocations of the CompileAndroidResources mnemonic",
+        ),
         _cpu_constraints = attr.label_keyed_string_dict(
             default = {
                 # The keys are labels to constraint_value targets representing the CPU for Android

rules/android_library/attrs.bzl

@@ -179,6 +179,12 @@ ATTRS = _attrs.add(
                 "Add these flags to the AIDL compiler command."
             ),
         ),
+        _mnemonic_jvm_flags = attr.label(
+            default = "//rules/flags:mnemonic_jvm_flags",
+            doc = "Additional JVM flags to set depending on the mnemonic " +
+                  "E.g --//rules/flags:mnemonic_jvm_flags=CompileAndroidResources=Xms1G " +
+                  " will add -Xms1G to all invocations of the CompileAndroidResources mnemonic",
+        ),
         neverlink = attr.bool(
             default = False,
             doc = (

rules/busybox.bzl

@@ -840,11 +840,20 @@ def _merge_compiled(
         jvm_flags = _C1_ONLY_FLAGS,
     )
 
-def _java_run(ctx, mnemonic = None, *args, **kwargs):
+def _java_run(ctx, mnemonic = None, jvm_flags = [], *args, **kwargs):
     enable_workers = ctx.fragments.android.persistent_busybox_tools
     multiplex_workers = ctx.fragments.android.persistent_multiplex_busybox_tools
 
-    _java.run(ctx, mnemonic = mnemonic, supports_workers = enable_workers, supports_multiplex_workers = multiplex_workers, *args, **kwargs)
+    extra_jvm_flags = ctx.attr._jvm_args[BuildSettingInfo].value
+    _java.run(
+        ctx,
+        mnemonic = mnemonic,
+        supports_workers = enable_workers,
+        supports_multiplex_workers = multiplex_workers,
+        jvm_flags = jvm_flags + extra_jvm_flags,
+        *args,
+        **kwargs
+    )
 
 def _escape_mv(s):
     """Escapes `:` and `,` in manifest values so they can be used as a busybox flag."""

rules/flags/BUILD

@@ -6,6 +6,7 @@ load("//rules/flags:additional_flags.bzl", "additional_flags")
 load("//rules/flags:configurations.bzl", "configurations")
 load("//rules/flags:flag_defs.bzl", "define_flags")
 load("//rules/flags:flags.bzl", "flags")
+load("//rules/flags:mnemonic_flags.bzl", "mnemonic_flag")
 
 licenses(["notice"])
 
@@ -209,6 +210,12 @@ int_flag(
     visibility = ["//visibility:public"],
 )
 
+mnemonic_flag(
+    name = "mnemonic_jvm_flags",
+    build_setting_default = "",
+    visibility = ["//visibility:public"],
+)
+
 configurations()
 
 define_flags()

rules/flags/mnemonic_flags.bzl

@@ -0,0 +1,18 @@
+load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
+
+def _mnemonic_flag_impl(ctx):
+    return [BuildSettingInfo(value = ctx.build_setting_value)]
+
+def extract_values_for_mnemonic(ctx, mnemonic):
+    flag_values = []
+    for value in ctx.attr._mnemonic_jvm_flags[BuildSettingInfo].value:
+        key, jvm_flag = value.split("=", 1)
+        if key != mnemonic:
+            continue
+        flag_values.append(jvm_flag)
+    return flag_values
+
+mnemonic_flag = rule(
+    implementation = _mnemonic_flag_impl,
+    build_setting = config.string(flag = True, allow_multiple = True)
+)

rules/java.bzl

@@ -14,6 +14,7 @@
 """Bazel Java APIs for the Android rules."""
 
 load("//rules:visibility.bzl", "PROJECT_VISIBILITY")
+load("//rules/flags:mnemonic_flags.bzl", "extract_values_for_mnemonic")
 load("@rules_java//java/common:java_common.bzl", "java_common")
 load("@rules_java//java/private:android_support.bzl", "android_support")  # buildifier: disable=bzl-visibility
 load(":path.bzl", _path = "path")
@@ -466,6 +467,12 @@ def _run(
     # Can still be overridden by callers of this method.
     jvm_flags = ["-Xms3G", "-Xmx3G", "-XX:+ExitOnOutOfMemoryError"] + jvm_flags
 
+    # Additional JVM flags depending on the mnemonic
+    if mnemonic and getattr(ctx.attr, "_mnemonic_jvm_flags", None):
+        extra_jvm_flags = extract_values_for_mnemonic(ctx, mnemonic)
+        if extra_jvm_flags:
+            jvm_flags += extra_jvm_flags
+
     # executable should be a File or a FilesToRunProvider
     jar = args.get("executable")
     if type(jar) == "FilesToRunProvider":