diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/Args.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/Args.java index 647df9f201418e..42a7f049d765a6 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/Args.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/Args.java @@ -380,7 +380,7 @@ private static StarlarkCallable validateMapEach(Object fn, boolean allowClosure) // This unfortunately disallows such trivially safe non-global // functions as "lambda x: x". // See https://github.com/bazelbuild/bazel/issues/12701. - if (sfn.getModule().getGlobal(sfn.getName()) != sfn && !allowClosure) { + if (!(sfn.isGlobal() || allowClosure)) { throw Starlark.errorf( "to avoid unintended retention of analysis data structures, " + "the map_each function (declared at %s) must be declared " diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkCustomCommandLine.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkCustomCommandLine.java index c318d95a8e246c..c5e305c5a9f2e5 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkCustomCommandLine.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkCustomCommandLine.java @@ -190,22 +190,29 @@ static final class VectorArg { private final int features; private final StringificationType stringificationType; - @Nullable private final StarlarkSemantics starlarkSemantics; // Null unless map_each is present. + // Null unless map_each is present. + @Nullable private final StarlarkSemantics starlarkSemantics; + // Null unless map_each is a global function. + @Nullable private final StarlarkFunction mapEachGlobalFunction; private VectorArg( int features, StringificationType stringificationType, - @Nullable StarlarkSemantics starlarkSemantics) { + @Nullable StarlarkSemantics starlarkSemantics, + @Nullable StarlarkFunction mapEachGlobalFunction) { this.features = features; this.stringificationType = stringificationType; this.starlarkSemantics = starlarkSemantics; + this.mapEachGlobalFunction = mapEachGlobalFunction; } private static VectorArg create( int features, StringificationType stringificationType, - @Nullable StarlarkSemantics starlarkSemantics) { - return intern(new VectorArg(features, stringificationType, starlarkSemantics)); + @Nullable StarlarkSemantics starlarkSemantics, + @Nullable StarlarkFunction mapEachGlobalFunction) { + return intern( + new VectorArg(features, stringificationType, starlarkSemantics, mapEachGlobalFunction)); } @VisibleForSerialization @@ -234,13 +241,20 @@ private static void push( features |= arg.formatJoined != null ? HAS_FORMAT_JOINED : 0; features |= arg.terminateWith != null ? HAS_TERMINATE_WITH : 0; features |= arg.nestedSet == null && arg.list.size() == 1 ? HAS_SINGLE_ARG : 0; + // Intern global Starlark functions in the VectorArg as they can be reused for all rule + // instances (and possibly even across multiple Args.add_all calls). + StarlarkFunction mapEachGlobalFunction = + arg.mapEach instanceof StarlarkFunction sfn && sfn.isGlobal() ? sfn : null; arguments.add( VectorArg.create( features, arg.nestedSetStringificationType, - arg.mapEach != null ? starlarkSemantics : null)); + arg.mapEach != null ? starlarkSemantics : null, + mapEachGlobalFunction)); if (arg.mapEach != null) { - arguments.add(arg.mapEach); + if (mapEachGlobalFunction == null) { + arguments.add(arg.mapEach); + } arguments.add(arg.location); } if (arg.nestedSet != null) { @@ -303,7 +317,11 @@ private int preprocess( StarlarkCallable mapEach = null; Location location = null; if ((features & HAS_MAP_EACH) != 0) { - mapEach = (StarlarkCallable) arguments.get(argi++); + if (mapEachGlobalFunction != null) { + mapEach = mapEachGlobalFunction; + } else { + mapEach = (StarlarkCallable) arguments.get(argi++); + } location = (Location) arguments.get(argi++); } @@ -500,7 +518,11 @@ private int addToFingerprint( StarlarkCallable mapEach = null; Location location = null; if ((features & HAS_MAP_EACH) != 0) { - mapEach = (StarlarkCallable) arguments.get(argi++); + if (mapEachGlobalFunction != null) { + mapEach = mapEachGlobalFunction; + } else { + mapEach = (StarlarkCallable) arguments.get(argi++); + } location = (Location) arguments.get(argi++); } @@ -729,13 +751,18 @@ public boolean equals(Object o) { } return features == that.features && stringificationType.equals(that.stringificationType) - && Objects.equals(starlarkSemantics, that.starlarkSemantics); + && Objects.equals(starlarkSemantics, that.starlarkSemantics) + // Do not use equals for StarlarkFunction as it is implemented in terms of + // SymbolGenerator.Symbol, which may not change if the function itself does. + && mapEachGlobalFunction == that.mapEachGlobalFunction; } @Override public int hashCode() { - return 31 * HashCodes.hashObjects(stringificationType, starlarkSemantics) - + Integer.hashCode(features); + return 31 + * (31 * HashCodes.hashObjects(stringificationType, starlarkSemantics) + + Integer.hashCode(features)) + + System.identityHashCode(mapEachGlobalFunction); } } diff --git a/src/main/java/net/starlark/java/eval/StarlarkFunction.java b/src/main/java/net/starlark/java/eval/StarlarkFunction.java index 0d9bf733f0c8e5..9828f8e264dfff 100644 --- a/src/main/java/net/starlark/java/eval/StarlarkFunction.java +++ b/src/main/java/net/starlark/java/eval/StarlarkFunction.java @@ -85,6 +85,13 @@ boolean isToplevel() { return rfn.isToplevel(); } + /** + * Whether this function is defined at the top level of a file. + */ + public boolean isGlobal() { + return getModule().getGlobal(getName()) == this; + } + // TODO(adonovan): many functions would be simpler if // parameterNames excluded the *args and **kwargs parameters, // (whose names are immaterial to the callee anyway). Do that.