Add uniquify parameter to TemplateDict.add_joined

fmeum · copybara-github · commit 9d250edb63f2 · 2022-11-02T05:47:34.000-07:00
The behavior is analogous to that of Args.add_joined. Closes #16213. PiperOrigin-RevId: 485571903 Change-Id: Id69de92d703d5202bfc7b50abfbbb4326441f242
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/TemplateDict.java
@@ -17,12 +17,14 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.analysis.actions.Substitution;
 import com.google.devtools.build.lib.analysis.actions.Substitution.ComputedSubstitution;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.starlarkbuildapi.TemplateDictApi;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Mutability;
@@ -58,6 +60,7 @@ public TemplateDictApi addJoined(
       Depset valuesSet,
       String joinWith,
       StarlarkCallable mapEach,
+      Boolean uniquify,
       StarlarkThread thread)
       throws EvalException {
     if (mapEach instanceof StarlarkFunction) {
@@ -71,7 +74,7 @@ public TemplateDictApi addJoined(
       }
     }
     substitutions.add(
-        new LazySubstitution(key, thread.getSemantics(), valuesSet, mapEach, joinWith));
+        new LazySubstitution(key, thread.getSemantics(), valuesSet, mapEach, uniquify, joinWith));
     return this;
   }
 
@@ -84,18 +87,21 @@ private static class LazySubstitution extends ComputedSubstitution {
     private final StarlarkSemantics semantics;
     private final Depset valuesSet;
     private final StarlarkCallable mapEach;
+    private final boolean uniquify;
     private final String joinWith;
 
     public LazySubstitution(
         String key,
         StarlarkSemantics semantics,
         Depset valuesSet,
         StarlarkCallable mapEach,
+        boolean uniquify,
         String joinWith) {
       super(key);
       this.semantics = semantics;
       this.valuesSet = valuesSet;
       this.mapEach = mapEach;
+      this.uniquify = uniquify;
       this.joinWith = joinWith;
     }
 
@@ -127,6 +133,19 @@ public String getValue() throws EvalException {
                 "Could not evaluate substitution for %s: %s", val, e.getMessage());
           }
         }
+        if (uniquify) {
+          // Stably deduplicate parts in-place.
+          int count = parts.size();
+          HashSet<String> seen = Sets.newHashSetWithExpectedSize(count);
+          int addIndex = 0;
+          for (int i = 0; i < count; ++i) {
+            String val = parts.get(i);
+            if (seen.add(val)) {
+              parts.set(addIndex++, val);
+            }
+          }
+          parts = parts.subList(0, addIndex);
+        }
         return Joiner.on(joinWith).join(parts);
       }
     }
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/TemplateDictApi.java
@@ -73,9 +73,24 @@ public interface TemplateDictApi extends StarlarkValue {
                 "A Starlark function accepting a single argument and returning a String. This"
                     + " function is applied to each item of the depset specified in the"
                     + " <code>values</code> parameter"),
+        @Param(
+            name = "uniquify",
+            named = true,
+            positional = false,
+            defaultValue = "False",
+            doc =
+                "If true, duplicate strings derived from <code>values</code> will be omitted. Only "
+                    + "the first occurrence of each string will remain. Usually this feature is "
+                    + "not needed because depsets already omit duplicates, but it can be useful "
+                    + "if <code>map_each</code> emits the same string for multiple items."),
       },
       useStarlarkThread = true)
   TemplateDictApi addJoined(
-      String key, Depset values, String joinWith, StarlarkCallable mapEach, StarlarkThread thread)
+      String key,
+      Depset values,
+      String joinWith,
+      StarlarkCallable mapEach,
+      Boolean uniquify,
+      StarlarkThread thread)
       throws EvalException;
 }
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleContextTest.java
@@ -3683,6 +3683,63 @@ public void testTemplateExpansionComputedSubstitution() throws Exception {
                 "td_files_key", "foo.txt%%bar.txt%%baz.txt"));
   }
 
+  @Test
+  public void testTemplateExpansionComputedSubstitutionWithUniquify() throws Exception {
+    setBuildLanguageOptions("--experimental_lazy_template_expansion");
+    scratch.file(
+        "test/rules.bzl",
+        "def _artifact_to_extension(file):",
+        "  return file.extension",
+        "",
+        "def _undertest_impl(ctx):",
+        "  template_dict = ctx.actions.template_dict()",
+        "  template_dict.add_joined('exts', depset(ctx.files.srcs),",
+        "                           map_each = _artifact_to_extension,",
+        "                           uniquify = True,",
+        "                           join_with = '%%',",
+        "                          )",
+        "  ctx.actions.expand_template(output=ctx.outputs.out,",
+        "                              template=ctx.file.template,",
+        "                              computed_substitutions=template_dict,",
+        "                              )",
+        "undertest_rule = rule(",
+        "  implementation = _undertest_impl,",
+        "  outputs = {'out': '%{name}.txt'},",
+        "  attrs = {'template': attr.label(allow_single_file=True),",
+        "           'srcs':attr.label_list(allow_files=True)",
+        "           },",
+        "  _skylark_testable = True,",
+        ")",
+        testingRuleDefinition);
+    scratch.file("test/template.txt", "exts", "exts");
+    scratch.file(
+        "test/BUILD",
+        "load(':rules.bzl', 'undertest_rule', 'testing_rule')",
+        "undertest_rule(",
+        "    name = 'undertest',",
+        "    template = ':template.txt',",
+        "    srcs = ['foo.txt', 'bar.log', 'baz.txt', 'bak.exe', 'far.sh', 'boo.sh'],",
+        ")",
+        "testing_rule(",
+        "    name = 'testing',",
+        "    dep = ':undertest',",
+        ")");
+    StarlarkRuleContext ruleContext = createRuleContext("//test:testing");
+    setRuleContext(ruleContext);
+    ev.update("file", ev.eval("ruleContext.attr.dep.files.to_list()[0]"));
+    ev.update("action", ev.eval("ruleContext.attr.dep[Actions].by_file[file]"));
+
+    assertThat(ev.eval("type(action)")).isEqualTo("Action");
+
+    Object contentUnchecked = ev.eval("action.content");
+    assertThat(contentUnchecked).isInstanceOf(String.class);
+    assertThat(contentUnchecked).isEqualTo("txt%%log%%exe%%sh\ntxt%%log%%exe%%sh\n");
+
+    Object substitutionsUnchecked = ev.eval("action.substitutions");
+    assertThat(substitutionsUnchecked).isInstanceOf(Dict.class);
+    assertThat(substitutionsUnchecked).isEqualTo(ImmutableMap.of("exts", "txt%%log%%exe%%sh"));
+  }
+
   @Test
   public void testTemplateExpansionComputedSubstitutionDuplicateKeys() throws Exception {
     setBuildLanguageOptions("--experimental_lazy_template_expansion");
