# Copyright 2024 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementation of cc_tool_map."""

load(
    "//cc/toolchains/impl:collect.bzl",
    "collect_provider",
    "collect_tools",
)
load(
    ":cc_toolchain_info.bzl",
    "ActionTypeSetInfo",
    "ToolConfigInfo",
)

def _cc_tool_map_impl(ctx):
    tools = collect_tools(ctx, ctx.attr.tools)
    action_sets = collect_provider(ctx.attr.actions, ActionTypeSetInfo)

    action_to_tool = {}
    action_to_as = {}
    for i in range(len(action_sets)):
        action_set = action_sets[i]
        tool = tools[i]

        for action in action_set.actions.to_list():
            if action in action_to_as:
                fail("The action %s appears multiple times in your tool_map (as %s and %s)" % (action.label, action_set.label, action_to_as[action].label))
            action_to_as[action] = action_set
            action_to_tool[action] = tool

    return [ToolConfigInfo(label = ctx.label, configs = action_to_tool)]

_cc_tool_map = rule(
    implementation = _cc_tool_map_impl,
    # @unsorted-dict-items
    attrs = {
        "actions": attr.label_list(
            providers = [ActionTypeSetInfo],
            mandatory = True,
            doc = """A list of action names to apply this action to.

See //cc/toolchains/actions:BUILD for valid options.
""",
        ),
        "tools": attr.label_list(
            mandatory = True,
            cfg = "exec",
            allow_files = True,
            doc = """The tool to use for the specified actions.

The tool may be a `cc_tool` or other executable rule.
""",
        ),
    },
    provides = [ToolConfigInfo],
)

def cc_tool_map(name, tools, **kwargs):
    """A toolchain configuration rule that maps toolchain actions to tools.

    A `cc_tool_map` aggregates all the tools that may be used for a given toolchain
    and maps them to their corresponding actions. Conceptually, this is similar to the
    `CXX=/path/to/clang++` environment variables that most build systems use to determine which
    tools to use for a given action. To simplify usage, some actions have been grouped together (for
    example,
    [//third_party/bazel_rules/rules_cc/cc/toolchains/actions:cpp_compile_actions](https://github.com/bazelbuild/rules_cc/tree/main/cc/toolchains/actions/BUILD)) to
    logically express "all the C++ compile actions".

    In Bazel, there is a little more granularity to the mapping, so the mapping doesn't follow the
    traditional `CXX`, `AR`, etc. naming scheme. For a comprehensive list of all the well-known
    actions, see //cc/toolchains/actions:BUILD.

    Example usage:
    ```
    load("//cc/toolchains:tool_map.bzl", "cc_tool_map")

    cc_tool_map(
        name = "all_tools",
        tools = {
            "//cc/toolchains/actions:assembly_actions": ":asm",
            "//cc/toolchains/actions:c_compile": ":clang",
            "//cc/toolchains/actions:cpp_compile_actions": ":clang++",
            "//cc/toolchains/actions:link_actions": ":lld",
            "//cc/toolchains/actions:objcopy_embed_data": ":llvm-objcopy",
            "//cc/toolchains/actions:strip": ":llvm-strip",
            "//cc/toolchains/actions:ar_actions": ":llvm-ar",
        },
    )
    ```

    Note:
       Due to an implementation limitation, if you need to map the same tool to multiple actions,
       you will need to create an intermediate alias for the tool for each set of actions. See
       https://github.com/bazelbuild/rules_cc/issues/235 for more details.

    Args:
        name: (str) The name of the target.
        tools: (Dict[Label, Label]) A mapping between
            `cc_action_type`/`cc_action_type_set` targets
            and the `cc_tool` or executable target that implements that action.
        **kwargs: [common attributes](https://bazel.build/reference/be/common-definitions#common-attributes) that should be applied to this rule.
    """
    _cc_tool_map(
        name = name,
        actions = tools.keys(),
        tools = tools.values(),
        **kwargs
    )