Add a string-to-label-dict attribute type.

This really just exposes the existing LABEL_DICT_UNARY type that already exists in Java.

Other than writing some boring conversion logic, all pieces fell together pretty pleasantly.

Work towards #7989. That one calls for a string-to-label-list-dict type, but I'm planning to resolve that as WAI, since string-to-label-dict is close enough.

RELNOTES: None.
PiperOrigin-RevId: 686415025
Change-Id: Ib07ada7ab2ede95220ed1cc2d3569996fc4afb88

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkAttrModule.java

@@ -43,6 +43,7 @@
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.starlarkbuildapi.NativeComputedDefaultApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkAttrModuleApi;
+import com.google.devtools.build.lib.starlarkbuildapi.StarlarkAttrModuleApi.Descriptor;
 import com.google.devtools.build.lib.util.FileType;
 import com.google.devtools.build.lib.util.FileTypeSet;
 import java.util.List;
@@ -620,6 +621,51 @@ public Descriptor labelListAttribute(
     return new Descriptor("label_list", attribute);
   }
 
+  @Override
+  public Descriptor stringKeyedLabelDictAttribute(
+      Boolean allowEmpty,
+      Object defaultValue,
+      Object doc,
+      Object allowFiles,
+      Object allowRules,
+      Sequence<?> providers,
+      Sequence<?> flags,
+      Boolean mandatory,
+      Object cfg,
+      Sequence<?> aspects,
+      StarlarkThread thread)
+      throws EvalException {
+    checkContext(thread, "attr.label_keyed_string_dict()");
+    Map<String, Object> kwargs =
+        optionMap(
+            DEFAULT_ARG,
+            defaultValue,
+            ALLOW_FILES_ARG,
+            allowFiles,
+            ALLOW_RULES_ARG,
+            allowRules,
+            PROVIDERS_ARG,
+            providers,
+            FLAGS_ARG,
+            flags,
+            MANDATORY_ARG,
+            mandatory,
+            ALLOW_EMPTY_ARG,
+            allowEmpty,
+            CONFIGURATION_ARG,
+            cfg,
+            ASPECTS_ARG,
+            aspects);
+    ImmutableAttributeFactory attribute =
+        createAttributeFactory(
+            BuildType.LABEL_DICT_UNARY,
+            Starlark.toJavaOptional(doc, String.class),
+            kwargs,
+            thread,
+            "string_keyed_label_dict");
+    return new Descriptor("string_keyed_label_dict", attribute);
+  }
+
   @Override
   public Descriptor labelKeyedStringDictAttribute(
       Boolean allowEmpty,

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkAttributesCollection.java

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.analysis.starlark;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.AliasProvider;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
@@ -147,6 +148,22 @@ private Builder(StarlarkRuleContext ruleContext) {
       this.context = ruleContext;
     }
 
+    static Dict<String, TransitiveInfoCollection> convertStringToLabelMap(
+        Map<String, Label> unconfiguredValue,
+        List<? extends TransitiveInfoCollection> prerequisites) {
+      Map<Label, TransitiveInfoCollection> prerequisiteMap = Maps.newHashMap();
+      for (TransitiveInfoCollection prereq : prerequisites) {
+        prerequisiteMap.put(AliasProvider.getDependencyLabel(prereq), prereq);
+      }
+
+      Dict.Builder<String, TransitiveInfoCollection> builder = Dict.builder();
+      for (var entry : unconfiguredValue.entrySet()) {
+        builder.put(entry.getKey(), prerequisiteMap.get(entry.getValue()));
+      }
+
+      return builder.buildImmutable();
+    }
+
     @SuppressWarnings("unchecked")
     public void addAttribute(Attribute a, Object val) {
       Type<?> type = a.getType();
@@ -170,10 +187,7 @@ public void addAttribute(Attribute a, Object val) {
         return;
       }
 
-      // TODO(b/140636597): Remove the LABEL_DICT_UNARY special case of this conditional
-      // LABEL_DICT_UNARY was previously not treated as a dependency-bearing type, and was put into
-      // Starlark as a Map<String, Label>; this special case preserves that behavior temporarily.
-      if (type.getLabelClass() != LabelClass.DEPENDENCY || type == BuildType.LABEL_DICT_UNARY) {
+      if (type.getLabelClass() != LabelClass.DEPENDENCY) {
         // Attribute values should be type safe
         attrBuilder.put(skyname, Attribute.valueToStarlark(val));
         return;
@@ -229,6 +243,11 @@ public void addAttribute(Attribute a, Object val) {
           || (type == BuildType.LABEL && a.getTransitionFactory().isSplit())) {
         List<?> allPrereq = context.getRuleContext().getPrerequisites(a.getName());
         attrBuilder.put(skyname, StarlarkList.immutableCopyOf(allPrereq));
+      } else if (type == BuildType.LABEL_DICT_UNARY) {
+        List<? extends TransitiveInfoCollection> allPrereq =
+            context.getRuleContext().getPrerequisites(a.getName());
+        attrBuilder.put(skyname, convertStringToLabelMap(
+            BuildType.LABEL_DICT_UNARY.cast(val), allPrereq));
       } else if (type == BuildType.LABEL_KEYED_STRING_DICT) {
         Dict.Builder<TransitiveInfoCollection, String> builder = Dict.builder();
         Map<Label, String> original = BuildType.LABEL_KEYED_STRING_DICT.cast(val);
@@ -238,17 +257,6 @@ public void addAttribute(Attribute a, Object val) {
           builder.put(prereq, original.get(AliasProvider.getDependencyLabel(prereq)));
         }
         attrBuilder.put(skyname, builder.buildImmutable());
-      } else if (type == BuildType.LABEL_DICT_UNARY) {
-        Map<Label, TransitiveInfoCollection> prereqsByLabel = new LinkedHashMap<>();
-        for (TransitiveInfoCollection target :
-            context.getRuleContext().getPrerequisites(a.getName())) {
-          prereqsByLabel.put(target.getLabel(), target);
-        }
-        ImmutableMap.Builder<String, TransitiveInfoCollection> attrValue = ImmutableMap.builder();
-        for (Map.Entry<String, Label> entry : ((Map<String, Label>) val).entrySet()) {
-          attrValue.put(entry.getKey(), prereqsByLabel.get(entry.getValue()));
-        }
-        attrBuilder.put(skyname, attrValue.buildOrThrow());
       } else {
         throw new IllegalArgumentException(
             "Can't transform attribute "

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleContext.java

@@ -36,6 +36,7 @@
 import com.google.devtools.build.lib.analysis.BashCommandConstructor;
 import com.google.devtools.build.lib.analysis.CommandHelper;
 import com.google.devtools.build.lib.analysis.ConfigurationMakeVariableContext;
+import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.FileProvider;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.LocationExpander;
@@ -524,6 +525,16 @@ private static StructImpl buildSplitAttributeInfo(
         if (attr.getType() == BuildType.LABEL) {
           Preconditions.checkState(splitPrereq.getValue().size() == 1);
           value = splitPrereq.getValue().get(0).getConfiguredTarget();
+        } else if (attr.getType() == BuildType.LABEL_DICT_UNARY) {
+          ImmutableList<ConfiguredTarget> prerequisites =
+              splitPrereq.getValue().stream()
+                  .map(ConfiguredTargetAndData::getConfiguredTarget)
+                  .collect(ImmutableList.toImmutableList());
+
+          value =
+              StarlarkAttributesCollection.Builder.convertStringToLabelMap(
+                  ruleContext.attributes().get(attr.getName(), BuildType.LABEL_DICT_UNARY),
+                  prerequisites);
         } else {
           // BuildType.LABEL_LIST
           value =
@@ -862,7 +873,7 @@ public ToolchainContextApi toolchains() throws EvalException {
     } else {
       return StarlarkToolchainContext.create(
           /* targetDescription= */ ruleContext.getToolchainContext().targetDescription(),
-          /* resolveToolchainInfoFunc= */ ruleContext.getToolchainContext()::forToolchainType,
+          /* resolveToolchainDataFunc= */ ruleContext.getToolchainContext()::forToolchainType,
           /* resolvedToolchainTypeLabels= */ ruleContext
               .getToolchainContext()
               .requestedToolchainTypeLabels()

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkAttrModuleApi.java

@@ -530,6 +530,106 @@ Descriptor labelListAttribute(
       StarlarkThread thread)
       throws EvalException;
 
+  @StarlarkMethod(
+      name = "string_keyed_label_dict",
+      doc =
+          "<p>Creates a schema for an attribute whose value is a dictionary where the keys are "
+              + "strings and the values are labels. This is a dependency attribute.</p>"
+              + DEPENDENCY_ATTR_TEXT,
+      parameters = {
+        @Param(name = ALLOW_EMPTY_ARG, defaultValue = "True", doc = ALLOW_EMPTY_DOC, named = true),
+        @Param(
+            name = DEFAULT_ARG,
+            allowedTypes = {
+              @ParamType(type = Dict.class),
+              @ParamType(type = StarlarkFunction.class)
+            },
+            defaultValue = "{}",
+            named = true,
+            positional = false,
+            doc =
+                DEFAULT_DOC
+                    + "Use strings or the <a"
+                    + " href=\"../builtins/Label.html#Label\"><code>Label</code></a> function to"
+                    + " specify default values, for example,"
+                    + " <code>attr.string_keyed_label_dict(default = {\"foo\": \"//a:b\","
+                    + " \"bar\": \"//a:c\"})</code>."),
+        @Param(
+            name = DOC_ARG,
+            allowedTypes = {@ParamType(type = String.class), @ParamType(type = NoneType.class)},
+            defaultValue = "None",
+            doc = DOC_DOC,
+            named = true,
+            positional = false),
+        @Param(
+            name = ALLOW_FILES_ARG,
+            allowedTypes = {
+              @ParamType(type = Boolean.class),
+              @ParamType(type = Sequence.class, generic1 = String.class),
+              @ParamType(type = NoneType.class),
+            },
+            defaultValue = "None",
+            named = true,
+            positional = false,
+            doc = ALLOW_FILES_DOC),
+        @Param(
+            name = ALLOW_RULES_ARG,
+            allowedTypes = {
+              @ParamType(type = Sequence.class, generic1 = String.class),
+              @ParamType(type = NoneType.class),
+            },
+            defaultValue = "None",
+            named = true,
+            positional = false,
+            doc = ALLOW_RULES_DOC),
+        @Param(
+            name = PROVIDERS_ARG,
+            defaultValue = "[]",
+            named = true,
+            positional = false,
+            doc = PROVIDERS_DOC),
+        @Param(
+            name = FLAGS_ARG,
+            allowedTypes = {@ParamType(type = Sequence.class, generic1 = String.class)},
+            defaultValue = "[]",
+            named = true,
+            positional = false,
+            doc = FLAGS_DOC),
+        @Param(
+            name = MANDATORY_ARG,
+            defaultValue = "False",
+            named = true,
+            positional = false,
+            doc = MANDATORY_DOC),
+        @Param(
+            name = CONFIGURATION_ARG,
+            defaultValue = "None",
+            named = true,
+            positional = false,
+            doc = CONFIGURATION_DOC),
+        @Param(
+            name = ASPECTS_ARG,
+            allowedTypes = {@ParamType(type = Sequence.class, generic1 = StarlarkAspectApi.class)},
+            defaultValue = "[]",
+            named = true,
+            positional = false,
+            doc = ASPECTS_ARG_DOC)
+      },
+      useStarlarkThread = true)
+  Descriptor stringKeyedLabelDictAttribute(
+      Boolean allowEmpty,
+      Object defaultValue,
+      Object doc,
+      Object allowFiles,
+      Object allowRules,
+      Sequence<?> providers,
+      Sequence<?> flags,
+      Boolean mandatory,
+      Object cfg,
+      Sequence<?> aspects,
+      StarlarkThread thread)
+      throws EvalException;
+
   @StarlarkMethod(
       name = "label_keyed_string_dict",
       doc =

src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java

@@ -2896,6 +2896,103 @@ public void initializer_stringListDict() throws Exception {
     assertThat(((Map<String, List<String>>) info.getValue("dict")).get("k")).containsExactly("val");
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void stringKeyedLabelDictWithSplitConfiguration() throws Exception {
+    scratch.file(
+        "a/BUILD",
+        """
+        load(":a.bzl", "a")
+        a(name="a", dict={"foo_key": ":foo", "gen_key": ":gen"})
+
+        filegroup(name="foo", srcs=["foo.txt"])
+        genrule(name="gen", srcs=[], outs=["gen.txt"], cmd="exit 1")
+        """);
+
+    scratch.file(
+        "a/a.bzl",
+        """
+        DictInfo = provider(fields=["dict"])
+
+        def _a_impl(ctx):
+            return [DictInfo(dict = ctx.split_attr.dict)]
+
+        def _trans_impl(settings, attr):
+          return {
+            "fastbuild_key": {"//command_line_option:compilation_mode": "fastbuild"},
+            "dbg_key": {"//command_line_option:compilation_mode": "dbg"},
+          }
+
+        trans = transition(
+            implementation = _trans_impl,
+            inputs = [],
+            outputs = ["//command_line_option:compilation_mode"])
+
+        a = rule(
+            implementation=_a_impl,
+            attrs={"dict": attr.string_keyed_label_dict(cfg=trans)})
+        """);
+
+    ConfiguredTarget a = getConfiguredTarget("//a:a");
+    StructImpl info =
+        (StructImpl)
+            a.get(
+                new StarlarkProvider.Key(
+                    keyForBuild(Label.parseCanonical("//a:a.bzl")), "DictInfo"));
+    Map<String, Map<String, ConfiguredTarget>> dict =
+        (Map<String, Map<String, ConfiguredTarget>>) info.getValue("dict");
+    assertThat(dict.keySet()).containsExactly("fastbuild_key", "dbg_key");
+    Map<String, ConfiguredTarget> fastbuild = dict.get("fastbuild_key");
+    Map<String, ConfiguredTarget> dbg = dict.get("dbg_key");
+    assertThat(fastbuild.keySet()).containsExactly("foo_key", "gen_key");
+    assertThat(dbg.keySet()).containsExactly("foo_key", "gen_key");
+
+    assertThat(getFilesToBuild(fastbuild.get("foo_key")).getSingleton().getExecPathString())
+        .isEqualTo("a/foo.txt");
+    assertThat(getFilesToBuild(dbg.get("foo_key")).getSingleton().getExecPathString())
+        .isEqualTo("a/foo.txt");
+    assertThat(getFilesToBuild(fastbuild.get("gen_key")).getSingleton().getExecPathString())
+        .endsWith("-fastbuild/bin/a/gen.txt");
+    assertThat(getFilesToBuild(dbg.get("gen_key")).getSingleton().getExecPathString())
+        .endsWith("-dbg/bin/a/gen.txt");
+  }
+
+  @Test
+  @SuppressWarnings("unchecked")
+  public void stringKeyedLabelDict() throws Exception {
+    scratch.file(
+        "a/BUILD",
+        """
+        load(":a.bzl", "a")
+        a(name="a", dict={"foo_key": ":foo", "gen_key": ":gen"})
+
+        filegroup(name="foo", srcs=["foo.txt"])
+        genrule(name="gen", srcs=[], outs=["gen.txt"], cmd="exit 1")
+        """);
+
+    scratch.file(
+        "a/a.bzl",
+        """
+        DictInfo = provider(fields=["dict"])
+
+        def _a_impl(ctx):
+            return [DictInfo(dict = ctx.attr.dict)]
+
+        a = rule(implementation=_a_impl, attrs={"dict": attr.string_keyed_label_dict()})
+        """);
+
+    ConfiguredTarget a = getConfiguredTarget("//a:a");
+    StructImpl info =
+        (StructImpl)
+            a.get(
+                new StarlarkProvider.Key(
+                    keyForBuild(Label.parseCanonical("//a:a.bzl")), "DictInfo"));
+    Map<String, ConfiguredTarget> dict = (Map<String, ConfiguredTarget>) info.getValue("dict");
+    assertThat(dict.keySet()).containsExactly("foo_key", "gen_key");
+    assertThat(dict.get("foo_key").getLabel()).isEqualTo(Label.parseCanonicalUnchecked("//a:foo"));
+    assertThat(dict.get("gen_key").getLabel()).isEqualTo(Label.parseCanonicalUnchecked("//a:gen"));
+  }
+
   @Test
   @SuppressWarnings("unchecked")
   public void initializer_labelKeyedStringDict() throws Exception {