Added label_injection to crate_universe annotations

crate_universe/extensions.bzl

@@ -377,7 +377,9 @@ load("@bazel_tools//tools/build_defs/repo:git.bzl", "new_git_repository")
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load(
     "//crate_universe/private:common_utils.bzl",
+    "apply_label_injections",
     "new_cargo_bazel_fn",
+    "sanitize_label_injections",
 )
 load("//crate_universe/private:crates_repository.bzl", "SUPPORTED_PLATFORM_TRIPLES")
 load(
@@ -407,13 +409,6 @@ load("//rust/platform:triple.bzl", "get_host_triple")
 load("//rust/platform:triple_mappings.bzl", "system_to_binary_ext")
 load(":defs.bzl", _crate_universe_crate = "crate")
 
-# A list of labels which may be relative (and if so, is within the repo the rule is generated in).
-#
-# If I were to write ":foo", with attr.label_list, it would evaluate to
-# "@@//:foo". However, for a tag such as deps, ":foo" should refer to
-# "@@rules_rust~crates~<crate>//:foo".
-_relative_label_list = attr.string_list
-
 _OPT_BOOL_VALUES = {
     "auto": None,
     "off": False,
@@ -1010,6 +1005,7 @@ def _crate_impl(module_ctx):
             repositories = annotation_dict.pop("repositories")
             crate = annotation_dict.pop("crate")
             version = annotation_dict.pop("version")
+            label_injections = sanitize_label_injections(annotation_dict.pop("label_injections", {}))
 
             # The crate.annotation function can take in either a list or a bool.
             # For the tag-based method, because it has type safety, we have to
@@ -1042,6 +1038,16 @@ def _crate_impl(module_ctx):
             if replacement:
                 annotation_dict["override_targets"]["bin"] = str(replacement)
 
+            # Replace any apparent labels with canonical ones. This is important for cases where an apparent label
+            # (e.g. `@foo` vs canonical `@@rules_rust+0.1.2+toolchains~foo`) is used in an annotation. To match the
+            # user intent of the annotation, these labels must be resolved at the same module level the annotation
+            # is defined.
+            for attr_name in _ANNOTATION_SELECT_ATTRS:
+                annotation_dict[attr_name] = apply_label_injections(
+                    label_mapping = label_injections,
+                    attribute = annotation_dict[attr_name],
+                )
+
             if not repositories:
                 _insert_annotation(module_annotations, crate, version, annotation_dict)
             else:
@@ -1305,6 +1311,13 @@ _ANNOTATION_NORMAL_ATTRS = {
         values = _OPT_BOOL_VALUES.keys(),
         default = "auto",
     ),
+    "label_injections": attr.label_keyed_string_dict(
+        doc = (
+            "A mapping of labels to stable names used in the annotation. This is necessary for cases where a `build_script_data` annotation is given and the new label is used in location " +
+            "expansion via `build_script_env`. E.g. `build_script_data = [\"@xz//:lzma\"]` and `build_script_env = {\"LZMA_BIN\": \"$(execpath @xz//:lzma)\"}`. This example would require " +
+            "`label_injections = {\"@xz\": \"@xz\"}` where the key will be resolved to a canonical label and the value is the apparent label used in annotations."
+        ),
+    ),
     "override_target_bin": attr.label(
         doc = "An optional alternate target to use when something depends on this crate to allow the parent repo to provide its own version of this dependency.",
     ),
@@ -1331,6 +1344,13 @@ _ANNOTATION_NORMAL_ATTRS = {
     ),
 }
 
+# A list of labels which may be relative (and if so, is within the repo the rule is generated in).
+#
+# If I were to write ":foo", with attr.label_list, it would evaluate to
+# "@@//:foo". However, for a tag such as deps, ":foo" should refer to
+# "@@rules_rust~crates~<crate>//:foo".
+_relative_label_list = attr.string_list
+
 _ANNOTATION_SELECT_ATTRS = {
     "build_script_compile_data": _relative_label_list(
         doc = "A list of labels to add to a crate's `cargo_build_script::compile_data` attribute.",

crate_universe/private/common_utils.bzl

@@ -211,3 +211,96 @@ def parse_alias_rule(value):
         bzl = str(bzl),
         rule = rule,
     )
+
+def sanitize_label_injections(label_injections):
+    """Load the `label_injections` attribute and sanitize it for use in replacements.
+
+    Args:
+        label_injections (dict[Label, str]): A mapping of labels (canonical) to string labels (apparent).
+
+    Returns:
+        dict[str, str]: A mapping of string
+    """
+    updated = {}
+    for canonical, apparent in label_injections.items():
+        canon_repo, _, _ = str(canonical).partition("//")
+        _, _, target = apparent.partition("//")
+        if target:
+            target = "//" + target
+
+        updated["{}{}".format(canon_repo, target)] = apparent
+
+    return updated
+
+def apply_label_injections(*, label_mapping, attribute):
+    """Apply label_injections to convert canonical labels to apparent labels.
+
+    Args:
+        label_mapping: The mapping of canonical label to apparent label.
+        attribute: The attribute to update.
+
+    Returns:
+        [dict | list | str]: The updated attribute based on `label_mapping`.
+    """
+    if not label_mapping:
+        return attribute
+
+    if not attribute:
+        return attribute
+
+    attr_type = type(attribute)
+
+    if attr_type == "dict":
+        updated = dict(attribute.items())
+        value_type = type(updated.values()[0])
+
+        if value_type == "string":
+            updated = dict(attribute.items())
+            for canonical, apparent in label_mapping.items():
+                updated = {
+                    key.replace(apparent, canonical): value.replace(apparent, canonical)
+                    for key, value in updated.items()
+                }
+            return updated
+
+        if value_type == "list":
+            for canonical, apparent in label_mapping.items():
+                updated = {
+                    key.replace(apparent, canonical): [v.replace(apparent, canonical) for v in value]
+                    for key, value in updated.items()
+                }
+            return updated
+
+        if value_type == "dict":
+            for canonical, apparent in label_mapping.items():
+                updated = {
+                    key.replace(apparent, canonical): {
+                        k.replace(apparent, canonical): v.replace(apparent, canonical)
+                        for k, v in value.items()
+                    }
+                    for key, value in updated.items()
+                }
+            return updated
+
+        fail("Unknown dict value type `{}` for: {}".format(
+            value_type,
+            attribute,
+        ))
+
+    if attr_type == "list":
+        updated = [i for i in attribute]
+        for canonical, apparent in label_mapping.items():
+            updated = [entry.replace(apparent, canonical) for entry in updated]
+        return updated
+
+    if attr_type == "string":
+        updated = str(attribute)
+        for canonical, apparent in label_mapping.items():
+            updated = updated.replace(apparent, canonical)
+
+        return updated
+
+    fail("Unexpected attribute type `{}` for: {}".format(
+        attr_type,
+        attribute,
+    ))

crate_universe/private/tests/label_injections/BUILD.bazel

@@ -0,0 +1,4 @@
+load(":label_injections_test.bzl", "label_injections_test_suite")
+
+############################ UNIT TESTS #############################
+label_injections_test_suite(name = "label_injections_test_suite")