Refactor AddArgsInfo into ExpandArgsInfo

BEGIN_PUBLIC
Refactor AddArgsInfo into ExpandArgsInfo

This allows us to create a similar mechanism to the current toolchain, while maintaining type safety.
END_PUBLIC

PiperOrigin-RevId: 615939056
Change-Id: I9b6763150194f8a76dfd8da730a3e2d45accbe20

6 files changed, 135 insertions, 34 deletions

diff --git a/cc/toolchains/args.bzl b/cc/toolchains/args.bzl
index 3bedfa8..e15b039 100644
--- a/cc/toolchains/args.bzl
+++ b/cc/toolchains/args.bzl
@@ -13,6 +13,7 @@
 # limitations under the License.
 """All providers for rule-based bazel toolchain config."""
 
+load("//cc:cc_toolchain_config_lib.bzl", "flag_group")
 load(
     "//cc/toolchains/impl:collect.bzl",
     "collect_action_types",
@@ -22,32 +23,41 @@ load(
 load(
     ":cc_toolchain_info.bzl",
     "ActionTypeSetInfo",
-    "AddArgsInfo",
     "ArgsInfo",
     "ArgsListInfo",
+    "ExpandArgsInfo",
     "FeatureConstraintInfo",
 )
 
 visibility("public")
 
 def _cc_args_impl(ctx):
-    add_args = [AddArgsInfo(
-        label = ctx.label,
-        args = tuple(ctx.attr.args),
-        files = depset([]),
-    )]
+    if not ctx.attr.args and not ctx.attr.env:
+        fail("cc_args requires at least one of args and env")
 
     actions = collect_action_types(ctx.attr.actions)
     files = collect_files(ctx.attr.data)
     requires = collect_provider(ctx.attr.requires_any_of, FeatureConstraintInfo)
 
+    expand = None
+    if ctx.attr.args:
+        # TODO: This is temporary until cc_expand_args is implemented.
+        expand = ExpandArgsInfo(
+            label = ctx.label,
+            expand = tuple(),
+            iterate_over = None,
+            files = files,
+            requires_types = {},
+            legacy_flag_group = flag_group(flags = ctx.attr.args),
+        )
+
     args = ArgsInfo(
         label = ctx.label,
         actions = actions,
         requires_any_of = tuple(requires),
-        files = files,
-        args = add_args,
+        expand = expand,
         env = ctx.attr.env,
+        files = files,
     )
     return [
         args,
@@ -74,7 +84,6 @@ See @rules_cc//cc/toolchains/actions:all for valid options.
 """,
         ),
         "args": attr.string_list(
-            mandatory = True,
             doc = """Arguments that should be added to the command-line.
 
 These are evaluated in order, with earlier args appearing earlier in the
diff --git a/cc/toolchains/cc_toolchain_info.bzl b/cc/toolchains/cc_toolchain_info.bzl
index 054eed7..feb1697 100644
--- a/cc/toolchains/cc_toolchain_info.bzl
+++ b/cc/toolchains/cc_toolchain_info.bzl
@@ -45,13 +45,16 @@ ActionTypeSetInfo = provider(
     },
 )
 
-AddArgsInfo = provider(
+ExpandArgsInfo = provider(
     doc = "A provider representation of Args.add/add_all/add_joined parameters",
     # @unsorted-dict-items
     fields = {
         "label": "(Label) The label defining this provider. Place in error messages to simplify debugging",
-        "args": "(Sequence[str]) The command-line arguments to add",
+        "expand": "(Sequence[ExpandArgsInfo]) The nested arg expansion. Mutually exclusive with args",
+        "iterate_over": "(Optional[str]) The variable to iterate over",
         "files": "(depset[File]) The files required to use this variable",
+        "requires_types": "(dict[str, str]) A mapping from variables to their expected type name (not type). This means that we can require the generic type Option, rather than an Option[T]",
+        "legacy_flag_group": "(flag_group) The flag_group this corresponds to",
     },
 )
 
@@ -62,7 +65,7 @@ ArgsInfo = provider(
         "label": "(Label) The label defining this provider. Place in error messages to simplify debugging",
         "actions": "(depset[ActionTypeInfo]) The set of actions this is associated with",
         "requires_any_of": "(Sequence[FeatureConstraintInfo]) This will be enabled if any of the listed predicates are met. Equivalent to with_features",
-        "args": "(Sequence[AddArgsInfo]) The command-line arguments to add.",
+        "expand": "(Optional[ExpandArgsInfo]) The args to expand. Equivalent to a flag group.",
         "files": "(depset[File]) Files required for the args",
         "env": "(dict[str, str]) Environment variables to apply",
     },
diff --git a/cc/toolchains/impl/legacy_converter.bzl b/cc/toolchains/impl/legacy_converter.bzl
index a216f66..e4e54bb 100644
--- a/cc/toolchains/impl/legacy_converter.bzl
+++ b/cc/toolchains/impl/legacy_converter.bzl
@@ -20,7 +20,6 @@ load(
     legacy_env_set = "env_set",
     legacy_feature = "feature",
     legacy_feature_set = "feature_set",
-    legacy_flag_group = "flag_group",
     legacy_flag_set = "flag_set",
     legacy_tool = "tool",
     legacy_with_feature_set = "with_feature_set",
@@ -50,10 +49,14 @@ def convert_feature_constraint(constraint):
         not_features = sorted([ft.name for ft in constraint.none_of.to_list()]),
     )
 
-def _convert_add_arg(add_arg):
-    return [legacy_flag_group(flags = list(add_arg.args))]
+def convert_args(args):
+    """Converts an ArgsInfo to flag_sets and env_sets.
 
-def _convert_args(args):
+    Args:
+        args: (ArgsInfo) The args to convert
+    Returns:
+        struct(flag_sets = List[flag_set], env_sets = List[env_sets])
+    """
     actions = _convert_actions(args.actions)
     with_features = [
         convert_feature_constraint(fc)
@@ -61,14 +64,11 @@ def _convert_args(args):
     ]
 
     flag_sets = []
-    if args.args:
-        flag_groups = []
-        for add_args in args.args:
-            flag_groups.extend(_convert_add_arg(add_args))
+    if args.expand != None:
         flag_sets.append(legacy_flag_set(
             actions = actions,
             with_features = with_features,
-            flag_groups = flag_groups,
+            flag_groups = [args.expand.legacy_flag_group],
         ))
 
     env_sets = []
@@ -93,7 +93,7 @@ def _convert_args_sequence(args_sequence):
     flag_sets = []
     env_sets = []
     for args in args_sequence:
-        legacy_args = _convert_args(args)
+        legacy_args = convert_args(args)
         flag_sets.extend(legacy_args.flag_sets)
         env_sets.extend(legacy_args.env_sets)
 
diff --git a/tests/rule_based_toolchain/args/BUILD b/tests/rule_based_toolchain/args/BUILD
index b6003a1..585ec91 100644
--- a/tests/rule_based_toolchain/args/BUILD
+++ b/tests/rule_based_toolchain/args/BUILD
@@ -18,6 +18,17 @@ util.helper_target(
     env = {"BAR": "bar"},
 )
 
+util.helper_target(
+    cc_args,
+    name = "env_only",
+    actions = ["//tests/rule_based_toolchain/actions:all_compile"],
+    data = [
+        "//tests/rule_based_toolchain/testdata:file1",
+        "//tests/rule_based_toolchain/testdata:multiple",
+    ],
+    env = {"BAR": "bar"},
+)
+
 analysis_test_suite(
     name = "test_suite",
     targets = TARGETS,
diff --git a/tests/rule_based_toolchain/args/args_test.bzl b/tests/rule_based_toolchain/args/args_test.bzl
index d4937b7..fbd4ce9 100644
--- a/tests/rule_based_toolchain/args/args_test.bzl
+++ b/tests/rule_based_toolchain/args/args_test.bzl
@@ -14,11 +14,23 @@
 """Tests for the cc_args rule."""
 
 load(
+    "//cc:cc_toolchain_config_lib.bzl",
+    "env_entry",
+    "env_set",
+    "flag_group",
+    "flag_set",
+)
+load(
     "//cc/toolchains:cc_toolchain_info.bzl",
     "ActionTypeInfo",
     "ArgsInfo",
     "ArgsListInfo",
 )
+load(
+    "//cc/toolchains/impl:legacy_converter.bzl",
+    "convert_args",
+)
+load("//tests/rule_based_toolchain:subjects.bzl", "subjects")
 
 visibility("private")
 
@@ -28,13 +40,17 @@ _SIMPLE_FILES = [
     "tests/rule_based_toolchain/testdata/multiple2",
 ]
 
-def _test_simple_args_impl(env, targets):
+_CONVERTED_ARGS = subjects.struct(
+    flag_sets = subjects.collection,
+    env_sets = subjects.collection,
+)
+
+def _simple_test(env, targets):
     simple = env.expect.that_target(targets.simple).provider(ArgsInfo)
     simple.actions().contains_exactly([
         targets.c_compile.label,
         targets.cpp_compile.label,
     ])
-    simple.args().contains_exactly([targets.simple.label])
     simple.env().contains_exactly({"BAR": "bar"})
     simple.files().contains_exactly(_SIMPLE_FILES)
 
@@ -44,12 +60,54 @@ def _test_simple_args_impl(env, targets):
     c_compile.args().contains_exactly([targets.simple[ArgsInfo]])
     c_compile.files().contains_exactly(_SIMPLE_FILES)
 
+    converted = env.expect.that_value(
+        convert_args(targets.simple[ArgsInfo]),
+        factory = _CONVERTED_ARGS,
+    )
+    converted.env_sets().contains_exactly([env_set(
+        actions = ["c_compile", "cpp_compile"],
+        env_entries = [env_entry(key = "BAR", value = "bar")],
+    )])
+
+    converted.flag_sets().contains_exactly([flag_set(
+        actions = ["c_compile", "cpp_compile"],
+        flag_groups = [flag_group(flags = ["--foo", "foo"])],
+    )])
+
+def _env_only_test(env, targets):
+    env_only = env.expect.that_target(targets.env_only).provider(ArgsInfo)
+    env_only.actions().contains_exactly([
+        targets.c_compile.label,
+        targets.cpp_compile.label,
+    ])
+    env_only.env().contains_exactly({"BAR": "bar"})
+    env_only.files().contains_exactly(_SIMPLE_FILES)
+
+    c_compile = env.expect.that_target(targets.simple).provider(ArgsListInfo).by_action().get(
+        targets.c_compile[ActionTypeInfo],
+    )
+    c_compile.files().contains_exactly(_SIMPLE_FILES)
+
+    converted = env.expect.that_value(
+        convert_args(targets.env_only[ArgsInfo]),
+        factory = _CONVERTED_ARGS,
+    )
+    converted.env_sets().contains_exactly([env_set(
+        actions = ["c_compile", "cpp_compile"],
+        env_entries = [env_entry(key = "BAR", value = "bar")],
+    )])
+
+    converted.flag_sets().contains_exactly([])
+
 TARGETS = [
     ":simple",
+    ":env_only",
     "//tests/rule_based_toolchain/actions:c_compile",
     "//tests/rule_based_toolchain/actions:cpp_compile",
 ]
 
+# @unsorted-dict-items
 TESTS = {
-    "simple_test": _test_simple_args_impl,
+    "simple_test": _simple_test,
+    "env_only_test_test": _env_only_test,
 }
diff --git a/tests/rule_based_toolchain/subjects.bzl b/tests/rule_based_toolchain/subjects.bzl
index fd1a953..f7eab87 100644
--- a/tests/rule_based_toolchain/subjects.bzl
+++ b/tests/rule_based_toolchain/subjects.bzl
@@ -21,9 +21,9 @@ load(
     "ActionTypeConfigSetInfo",
     "ActionTypeInfo",
     "ActionTypeSetInfo",
-    "AddArgsInfo",
     "ArgsInfo",
     "ArgsListInfo",
+    "ExpandArgsInfo",
     "FeatureConstraintInfo",
     "FeatureInfo",
     "FeatureSetInfo",
@@ -40,6 +40,10 @@ visibility("//tests/rule_based_toolchain/...")
 # This makes it rather awkward for copybara.
 runfiles_subject = lambda value, meta: _subjects.depset_file(value.files, meta = meta)
 
+# The string type has .equals(), which is all we can really do for an unknown
+# type.
+unknown_subject = _subjects.str
+
 # buildifier: disable=name-conventions
 _ActionTypeFactory = generate_factory(
     ActionTypeInfo,
@@ -102,13 +106,27 @@ _FeatureConstraintFactory = generate_factory(
     ),
 )
 
+_EXPAND_ARGS_FLAGS = dict(
+    expand = None,
+    files = _subjects.depset_file,
+    iterate_over = optional_subject(_subjects.str),
+    legacy_flag_group = unknown_subject,
+    requires_types = _subjects.dict,
+)
+
 # buildifier: disable=name-conventions
-_AddArgsFactory = generate_factory(
-    AddArgsInfo,
-    "AddArgsInfo",
-    dict(
-        args = _subjects.collection,
-        files = _subjects.depset_file,
+_FakeExpandArgsFactory = generate_factory(
+    ExpandArgsInfo,
+    "ExpandArgsInfo",
+    _EXPAND_ARGS_FLAGS,
+)
+
+# buildifier: disable=name-conventions
+_ExpandArgsFactory = generate_factory(
+    ExpandArgsInfo,
+    "ExpandArgsInfo",
+    _EXPAND_ARGS_FLAGS | dict(
+        expand = ProviderSequence(_FakeExpandArgsFactory),
     ),
 )
 
@@ -118,9 +136,10 @@ _ArgsFactory = generate_factory(
     "ArgsInfo",
     dict(
         actions = ProviderDepset(_ActionTypeFactory),
-        args = ProviderSequence(_AddArgsFactory),
         env = _subjects.dict,
         files = _subjects.depset_file,
+        # Use .factory so it's not inlined.
+        expand = optional_subject(_ExpandArgsFactory.factory),
         requires_any_of = ProviderSequence(_FeatureConstraintFactory),
     ),
 )
@@ -201,7 +220,7 @@ _ToolchainConfigFactory = generate_factory(
 FACTORIES = [
     _ActionTypeFactory,
     _ActionTypeSetFactory,
-    _AddArgsFactory,
+    _ExpandArgsFactory,
     _ArgsFactory,
     _ArgsListFactory,
     _MutuallyExclusiveCategoryFactory,
@@ -217,6 +236,7 @@ result_fn_wrapper = _result_fn_wrapper
 
 subjects = struct(
     **(structs.to_dict(_subjects) | dict(
+        unknown = unknown_subject,
         result = result_subject,
         optional = optional_subject,
         struct = struct_subject,