Add requires_not_none support for cc_args.env

cc/toolchains/args.bzl

@@ -14,7 +14,7 @@
 """All providers for rule-based bazel toolchain config."""
 
 load("@bazel_skylib//rules/directory:providers.bzl", "DirectoryInfo")
-load("//cc/toolchains/impl:args_utils.bzl", "validate_nested_args")
+load("//cc/toolchains/impl:args_utils.bzl", "validate_env_variables", "validate_nested_args")
 load(
     "//cc/toolchains/impl:collect.bzl",
     "collect_action_types",
@@ -33,7 +33,9 @@ load(
     "ArgsInfo",
     "ArgsListInfo",
     "BuiltinVariablesInfo",
+    "EnvInfo",
     "FeatureConstraintInfo",
+    "VariableInfo",
 )
 
 visibility("public")
@@ -68,12 +70,24 @@ def _cc_args_impl(ctx):
 
     requires = collect_provider(ctx.attr.requires_any_of, FeatureConstraintInfo)
 
+    env = EnvInfo(
+        label = ctx.label,
+        entries = formatted_env,
+        requires_not_none = ctx.attr.requires_not_none[VariableInfo].name if ctx.attr.requires_not_none else None,
+    )
+    validate_env_variables(
+        actions = actions.to_list(),
+        env = env,
+        variables = ctx.attr._variables[BuiltinVariablesInfo].variables,
+        used_format_vars = used_format_vars,
+    )
+
     args = ArgsInfo(
         label = ctx.label,
         actions = actions,
         requires_any_of = tuple(requires),
         nested = nested,
-        env = formatted_env,
+        env = env,
         files = files,
         allowlist_include_directories = depset(
             direct = [d[DirectoryInfo] for d in ctx.attr.allowlist_include_directories],

cc/toolchains/cc_toolchain_info.bzl

@@ -97,6 +97,16 @@ NestedArgsInfo = provider(
     },
 )
 
+EnvInfo = provider(
+    doc = "A set of environment variables to be added to the environment for specific actions",
+    # @unsorted-dict-items
+    fields = {
+        "label": "(Label) The label defining this provider. Place in error messages to simplify debugging",
+        "entries": "(dict[str, str]) A mapping from environment variable name to value",
+        "requires_not_none": "(Optional[str]) The variable that must be not None to apply these environment variables",
+    },
+)
+
 ArgsInfo = provider(
     doc = "A set of arguments to be added to the command line for specific actions",
     # @unsorted-dict-items
@@ -106,7 +116,7 @@ ArgsInfo = provider(
         "requires_any_of": "(Sequence[FeatureConstraintInfo]) This will be enabled if any of the listed predicates are met. Equivalent to with_features",
         "nested": "(Optional[NestedArgsInfo]) The args expand. Equivalent to a flag group.",
         "files": "(depset[File]) Files required for the args",
-        "env": "(dict[str, str]) Environment variables to apply",
+        "env": "(EnvInfo) Environment variables to apply",
         "allowlist_include_directories": "(depset[DirectoryInfo]) Include directories implied by these arguments that should be allowlisted in Bazel's include checker",
         "allowlist_absolute_include_directories": "(depset[str]) Absolute include directories implied by these arguments that should be allowlisted in Bazel's include checker",
     },

cc/toolchains/impl/args_utils.bzl

@@ -122,3 +122,60 @@ def validate_nested_args(*, nested_args, variables, actions, label, fail = fail)
 
         for child in nested_args.nested:
             stack.append((child, overrides))
+
+def validate_env_variables(*, env, variables, actions, used_format_vars, fail = fail):
+    """Validates the typing for environment variables with requires_not_none and format variables.
+
+    Args:
+        env: (EnvInfo) The env info to validate
+        variables: (Dict[str, VariableInfo]) A mapping from variable name to
+          the VariableInfo.
+        actions: (List[ActionTypeInfo]) The actions we require these variables
+          to be valid for.
+        used_format_vars: (List[str]) The list of variables used in env format strings.
+        fail: The fail function. Use for testing only.
+    """
+    if not env.entries:
+        return
+
+    env_requires_types = {}
+    if env.requires_not_none:
+        env_requires_types[env.requires_not_none] = struct(
+            msg = "requires_not_none for env variables only works on optional string, file, or directory types",
+            valid_types = ["option"],
+            valid_elements = ["string", "file", "directory"],
+        )
+
+    for var_name in used_format_vars:
+        env_requires_types[var_name] = struct(
+            msg = "format only works on string, file, or directory type variables",
+            valid_types = ["string", "file", "directory", "option"],
+            valid_elements = ["string", "file", "directory"],
+        )
+
+    for var_name, requirement in env_requires_types.items():
+        type = get_type(
+            name = var_name,
+            variables = variables,
+            overrides = {},
+            actions = actions,
+            args_label = env.label,
+            nested_label = env.label,
+            fail = fail,
+        )
+
+        if type["name"] not in requirement.valid_types:
+            fail("{msg}, but {var_name} has type {type}".format(
+                var_name = var_name,
+                msg = requirement.msg,
+                type = type["repr"],
+            ))
+
+        if type["name"] == "option":
+            underlying_type = type["elements"]
+            if underlying_type["name"] not in requirement.valid_elements:
+                fail("{msg}, but {var_name} has type {type}".format(
+                    var_name = var_name,
+                    msg = requirement.msg,
+                    type = underlying_type["repr"],
+                ))

cc/toolchains/impl/legacy_converter.bzl

@@ -70,16 +70,21 @@ def convert_args(args, strip_actions = False):
         ))
 
     env_sets = []
-    if args.env:
+    if args.env.entries:
+        # NOTE: Use kwargs to support older bazel versions
+        kwargs = {}
+        if args.env.requires_not_none:
+            kwargs["expand_if_available"] = args.env.requires_not_none
         env_sets.append(legacy_env_set(
             actions = actions,
             with_features = with_features,
             env_entries = [
                 legacy_env_entry(
                     key = key,
                     value = value,
+                    **kwargs
                 )
-                for key, value in args.env.items()
+                for key, value in args.env.entries.items()
             ],
         ))
     return struct(

cc/toolchains/impl/nested_args.bzl

@@ -291,10 +291,8 @@ def nested_args_provider(
 def _escape(s):
     return s.replace("%", "%%")
 
-def _format_target(target, arg, allow_variables, fail = fail):
+def _format_target(target, fail = fail):
     if VariableInfo in target:
-        if not allow_variables:
-            fail("Unsupported cc_variable substitution %s in %r." % (target.label, arg))
         return "%%{%s}" % target[VariableInfo].name
     elif DirectoryInfo in target:
         return _escape(target[DirectoryInfo].path)
@@ -305,7 +303,7 @@ def _format_target(target, arg, allow_variables, fail = fail):
 
     fail("%s should be either a variable, a directory, or a single file." % target.label)
 
-def _format_string(arg, format, used_vars, allow_variables, fail = fail):
+def _format_string(arg, format, used_vars, fail = fail):
     upto = 0
     out = []
     has_format = False
@@ -333,7 +331,7 @@ def _format_string(arg, format, used_vars, allow_variables, fail = fail):
             else:
                 used_vars[variable] = None
                 has_format = True
-                out.append(_format_target(format[variable], arg, allow_variables, fail = fail))
+                out.append(_format_target(format[variable], fail = fail))
                 upto += len(variable) + 2
 
         elif arg[upto] == "}":
@@ -363,7 +361,7 @@ def format_list(args, format, must_use = [], fail = fail):
     used_vars = {}
 
     for arg in args:
-        formatted.append(_format_string(arg, format, used_vars, True, fail))
+        formatted.append(_format_string(arg, format, used_vars, fail))
 
     unused_vars = [var for var in must_use if var not in used_vars]
     if unused_vars:
@@ -390,7 +388,7 @@ def format_dict_values(env, format, must_use = [], fail = fail):
     used_vars = {}
 
     for key, value in env.items():
-        formatted[key] = _format_string(value, format, used_vars, False, fail)
+        formatted[key] = _format_string(value, format, used_vars, fail)
 
     unused_vars = [var for var in must_use if var not in used_vars]
     if unused_vars:

tests/rule_based_toolchain/args/BUILD

@@ -2,6 +2,7 @@ load("@rules_testing//lib:util.bzl", "util")
 load("//cc/toolchains:args.bzl", "cc_args")
 load("//cc/toolchains/impl:variables.bzl", "cc_variable", "types")
 load("//tests/rule_based_toolchain:analysis_test_suite.bzl", "analysis_test_suite")
+load("//tests/rule_based_toolchain:testing_rules.bzl", "expect_failure_test")
 load(":args_test.bzl", "TARGETS", "TESTS")
 
 cc_variable(
@@ -35,6 +36,15 @@ util.helper_target(
     env = {"BAR": "bar"},
 )
 
+util.helper_target(
+    cc_args,
+    name = "env_only_requires",
+    actions = ["//cc/toolchains/actions:compile_actions"],
+    env = {"BAR": "{dependency_file}"},
+    format = {"dependency_file": "//cc/toolchains/variables:dependency_file"},
+    requires_not_none = "//cc/toolchains/variables:dependency_file",
+)
+
 util.helper_target(
     cc_args,
     name = "iterate_over_optional",
@@ -53,6 +63,68 @@ util.helper_target(
     args = ["--secret-builtin-include-dir"],
 )
 
+util.helper_target(
+    cc_args,
+    name = "good_env_format",
+    actions = ["//cc/toolchains/actions:compile_actions"],
+    env = {"FOO": "{gcov_gcno_file}"},
+    format = {"gcov_gcno_file": "//cc/toolchains/variables:gcov_gcno_file"},
+)
+
+util.helper_target(
+    cc_args,
+    name = "good_env_format_optional",
+    actions = ["//cc/toolchains/actions:compile_actions"],
+    env = {"FOO": "{dependency_file}"},
+    format = {"dependency_file": "//cc/toolchains/variables:dependency_file"},
+    requires_not_none = "//cc/toolchains/variables:dependency_file",
+)
+
+util.helper_target(
+    cc_args,
+    name = "bad_env_format_list",
+    actions = ["//cc/toolchains/actions:compile_actions"],
+    env = {"FOO": "{preprocessor_defines}"},
+    format = {"preprocessor_defines": "//cc/toolchains/variables:preprocessor_defines"},
+    tags = ["manual"],
+)
+
+expect_failure_test(
+    name = "bad_env_format_list_test",
+    failure_message = "format only works on string, file, or directory type variables, but preprocessor_defines has type List[string]",
+    target = ":bad_env_format_list",
+)
+
+util.helper_target(
+    cc_args,
+    name = "bad_env_format_optional",
+    actions = ["//cc/toolchains/actions:compile_actions"],
+    env = {"FOO": "{user_compile_flags}"},
+    format = {"user_compile_flags": "//cc/toolchains/variables:user_compile_flags"},
+    tags = ["manual"],
+)
+
+expect_failure_test(
+    name = "bad_env_format_optional_test",
+    failure_message = "format only works on string, file, or directory type variables, but user_compile_flags has type List[string]",
+    target = ":bad_env_format_optional",
+)
+
+util.helper_target(
+    cc_args,
+    name = "bad_env_format_wrong_action",
+    actions = ["//cc/toolchains/actions:compile_actions"],
+    env = {"FOO": "{runtime_solib_name}"},
+    format = {"runtime_solib_name": "//cc/toolchains/variables:runtime_solib_name"},
+    tags = ["manual"],
+)
+
+expect_failure_test(
+    name = "bad_env_format_wrong_action_test",
+    failure_message = "runtime_solib_name is inaccessible from the action",
+    target = ":bad_env_format_wrong_action",
+)
+
 analysis_test_suite(
     name = "test_suite",
     targets = TARGETS,