From 8764ebb95e1c27336a72cdc16a2dc68c0fc50c84 Mon Sep 17 00:00:00 2001 From: Fabian Meumertzheim Date: Wed, 2 Nov 2022 05:45:59 -0700 Subject: [PATCH] Add uniquify parameter to TemplateDict.add_joined The behavior is analogous to that of Args.add_joined. Closes #16213. PiperOrigin-RevId: 485571903 Change-Id: Id69de92d703d5202bfc7b50abfbbb4326441f242 --- .../lib/analysis/starlark/TemplateDict.java | 21 ++++++- .../lib/starlarkbuildapi/TemplateDictApi.java | 17 +++++- .../lib/starlark/StarlarkRuleContextTest.java | 57 +++++++++++++++++++ 3 files changed, 93 insertions(+), 2 deletions(-) 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 index 1b956607f9ee58..58587ae5ccc973 100644 --- 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,6 +87,7 @@ 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( @@ -91,11 +95,13 @@ public LazySubstitution( 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 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 index e2175609a9bbe1..39c648d3b84613 100644 --- 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" + " values parameter"), + @Param( + name = "uniquify", + named = true, + positional = false, + defaultValue = "False", + doc = + "If true, duplicate strings derived from values 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 map_each 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 index 0b2f77a7346a70..88a8d8bc6407df 100644 --- 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");