Intern global map_each functions in VectorArgs

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 "

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);
  }
  }
 

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.