Support dict select union

databricks · Mar 17, 2022 · 0ab1610 · 0ab1610
1 parent 542a4e7
commit 0ab1610
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 1 deletion.
diff --git a/src/main/java/com/google/devtools/build/lib/packages/SelectorList.java b/src/main/java/com/google/devtools/build/lib/packages/SelectorList.java
@@ -112,7 +112,12 @@ static SelectorList concat(Object x, Object y) throws EvalException {
 
   @Override
   public SelectorList binaryOp(TokenKind op, Object that, boolean thisLeft) throws EvalException {
-    if (op == TokenKind.PLUS) {
+    if (getNativeType(that).equals(Dict.class)) {
+      if (op == TokenKind.PIPE) {
+        return thisLeft ? concat(this, that) : concat(that, this);
+      }
+    }
+    else if (op == TokenKind.PLUS) {
       return thisLeft ? concat(this, that) : concat(that, this);
     }
     return null;

diff --git a/src/main/java/com/google/devtools/build/lib/packages/Type.java b/src/main/java/com/google/devtools/build/lib/packages/Type.java
@@ -571,6 +571,15 @@ public Map<KeyT, ValueT> convert(Object x, Object what, Object context)
       return ImmutableMap.copyOf(result);
     }
 
+    @Override
+    public Map<KeyT, ValueT> concat(Iterable<Map<KeyT, ValueT>> iterable) {
+      ImmutableMap.Builder<KeyT, ValueT> output = ImmutableMap.builder();
+      for (Map<KeyT, ValueT> map: iterable){
+        output.putAll(map);
+      }
+      return output.build();
+    }
+
     @Override
     public Map<KeyT, ValueT> getDefaultValue() {
       return empty;

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/ConfigurableAttributesTest.java b/src/test/java/com/google/devtools/build/lib/analysis/ConfigurableAttributesTest.java
@@ -777,7 +777,24 @@ public void noMatchCustomErrorMessage() throws Exception {
   @Test
   public void nativeTypeConcatenatedWithSelect() throws Exception {
     writeConfigRules();
+    scratch.file("java/foo/rule.bzl",
+        "def _rule_impl(ctx):",
+        "    return []",
+        "myrule = rule(",
+        "    implementation = _rule_impl,",
+        "    attrs = {",
+        "        'deps': attr.label_keyed_string_dict()",
+        "    },",
+        ")");
     scratch.file("java/foo/BUILD",
+        "load(':rule.bzl', 'myrule')",
+        "myrule(",
+        "    name = 'mytarget',",
+        "    deps = {':always': 'a'} | select({",
+        "        '//conditions:a': {':a': 'a'},",
+        "        '//conditions:b': {':b': 'b'},",
+        "    })",
+        ")",
         "java_binary(",
         "    name = 'binary',",
         "    srcs = ['binary.java'],",
@@ -802,12 +819,34 @@ public void nativeTypeConcatenatedWithSelect() throws Exception {
             "bin java/foo/libb.jar"),
         /*not expected:*/ ImmutableList.of(
             "bin java/foo/liba.jar"));
+    checkRule(
+        "//java/foo:mytarget",
+        "--test_arg=b",
+        /*expected:*/ ImmutableList.of("bin java/foo/libalways.jar", "bin java/foo/libb.jar"),
+        /*not expected:*/ ImmutableList.of("bin java/foo/liba.jar"));
   }
 
   @Test
   public void selectConcatenatedWithNativeType() throws Exception {
     writeConfigRules();
+    scratch.file("java/foo/rule.bzl",
+        "def _rule_impl(ctx):",
+        "    return []",
+        "myrule = rule(",
+        "    implementation = _rule_impl,",
+        "    attrs = {",
+        "        'deps': attr.label_keyed_string_dict()",
+        "    },",
+        ")");
     scratch.file("java/foo/BUILD",
+        "load(':rule.bzl', 'myrule')",
+        "myrule(",
+        "    name = 'mytarget',",
+        "    deps = select({",
+        "        '//conditions:a': {':a': 'a'},",
+        "        '//conditions:b': {':b': 'b'},",
+        "    }) | {':always': 'a'}",
+        ")",
         "java_binary(",
         "    name = 'binary',",
         "    srcs = ['binary.java'],",
@@ -831,12 +870,38 @@ public void selectConcatenatedWithNativeType() throws Exception {
             "bin java/foo/libb.jar"),
         /*not expected:*/ ImmutableList.of(
             "bin java/foo/liba.jar"));
+
+    checkRule(
+        "//java/foo:mytarget",
+        "--test_arg=b",
+        /*expected:*/ ImmutableList.of("bin java/foo/libalways.jar", "bin java/foo/libb.jar"),
+        /*not expected:*/ ImmutableList.of("bin java/foo/liba.jar"));
   }
 
   @Test
   public void selectConcatenatedWithSelect() throws Exception {
     writeConfigRules();
+    scratch.file("java/foo/rule.bzl",
+        "def _rule_impl(ctx):",
+        "    return []",
+        "myrule = rule(",
+        "    implementation = _rule_impl,",
+        "    attrs = {",
+        "        'deps': attr.label_keyed_string_dict()",
+        "    },",
+        ")");
     scratch.file("java/foo/BUILD",
+        "load(':rule.bzl', 'myrule')",
+        "myrule(",
+        "    name = 'mytarget',",
+        "    deps = select({",
+        "        '//conditions:a': {':a': 'a'},",
+        "        '//conditions:b': {':b': 'b'},",
+        "    }) | select({",
+        "        '//conditions:a': {':a2': 'a2'},",
+        "        '//conditions:b': {':b2': 'b2'},",
+        "    })",
+        ")",
         "java_binary(",
         "    name = 'binary',",
         "    srcs = ['binary.java'],",
@@ -868,6 +933,12 @@ public void selectConcatenatedWithSelect() throws Exception {
         /*not expected:*/ ImmutableList.of(
             "bin java/foo/liba.jar",
             "bin java/foo/liba2.jar"));
+
+    checkRule(
+        "//java/foo:mytarget",
+        "--test_arg=b",
+        /*expected:*/ ImmutableList.of("bin java/foo/libb.jar", "bin java/foo/libb2.jar"),
+        /*not expected:*/ ImmutableList.of("bin java/foo/liba.jar", "bin java/foo/liba2.jar"));
   }
 
   @Test

diff --git a/src/test/java/com/google/devtools/build/lib/packages/BuildTypeTest.java b/src/test/java/com/google/devtools/build/lib/packages/BuildTypeTest.java
@@ -432,6 +432,38 @@ public void testSelectorList() throws Exception {
                 .entrySet());
   }
 
+  @Test
+  public void testSelectorDict() throws Exception {
+    Object selector1 = new SelectorValue(ImmutableMap.of("//conditions:a",
+            ImmutableMap.of("//a:a", "a"), "//conditions:b", ImmutableMap.of("//b:b", "b")), "");
+    Object selector2 = new SelectorValue(ImmutableMap.of("//conditions:c",
+            ImmutableMap.of("//c:c", "c"), "//conditions:d", ImmutableMap.of("//d:d", "d")), "");
+    BuildType.SelectorList<Map<Label, String>> selectorList =
+            new BuildType.SelectorList<Map<Label, String>>(
+                    ImmutableList.of(selector1, selector2),
+                    null,
+                    labelConversionContext,
+                    BuildType.LABEL_KEYED_STRING_DICT);
+
+    assertThat(selectorList.getOriginalType()).isEqualTo(BuildType.LABEL_KEYED_STRING_DICT);
+    assertThat(selectorList.getKeyLabels())
+            .containsExactly(
+                    Label.parseAbsolute("//conditions:a", ImmutableMap.of()),
+                    Label.parseAbsolute("//conditions:b", ImmutableMap.of()),
+                    Label.parseAbsolute("//conditions:c", ImmutableMap.of()),
+                    Label.parseAbsolute("//conditions:d", ImmutableMap.of()));
+
+    List<Selector<Map<Label, String>>> selectors = selectorList.getSelectors();
+    assertThat(selectors.get(0).getEntries().entrySet())
+            .containsExactlyElementsIn(
+                    ImmutableMap.of(
+                                    Label.parseAbsolute("//conditions:a", ImmutableMap.of()),
+                                    ImmutableMap.of(Label.create("@//a", "a"), "a"),
+                                    Label.parseAbsolute("//conditions:b", ImmutableMap.of()),
+                                    ImmutableMap.of(Label.create("@//b", "b"), "b"))
+                            .entrySet());
+  }
+
   @Test
   public void testSelectorListMixedTypes() throws Exception {
     Object selector1 =