Intern global map_each functions in VectorArgs

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).

Closes #23685.

PiperOrigin-RevId: 681565448
Change-Id: I80d9048eb6f55a5ad0ef7a76674612d1f7f7832d

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,21 @@ 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), saving a slot in
+      // arguments.
+      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 +318,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 +519,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++);
       }
 
@@ -720,6 +743,7 @@ Builder setTerminateWith(String terminateWith) {
     }
 
     @Override
+    @SuppressWarnings("ReferenceEquality")
     public boolean equals(Object o) {
       if (this == o) {
         return true;
@@ -729,13 +753,19 @@ public boolean equals(Object o) {
       }
       return features == that.features
           && stringificationType.equals(that.stringificationType)
-          && Objects.equals(starlarkSemantics, that.starlarkSemantics);
+          && Objects.equals(starlarkSemantics, that.starlarkSemantics)
+          // Use reference equality to avoid resurrecting a weakly-reachable but value-equal
+          // StarlarkFunction instance. Value-equal instances may be created when a .bzl file is
+          // re-evaluated.
+          && mapEachGlobalFunction == that.mapEachGlobalFunction;
     }
 
     @Override
     public int hashCode() {
-      return 31 * HashCodes.hashObjects(stringificationType, starlarkSemantics)
-          + Integer.hashCode(features);
+      int result = HashCodes.hashObjects(stringificationType, starlarkSemantics);
+      result = 31 * result + Integer.hashCode(features);
+      result = 31 * result + System.identityHashCode(mapEachGlobalFunction);
+      return result;
     }
   }
 

src/main/java/net/starlark/java/eval/StarlarkFunction.java

@@ -85,6 +85,11 @@ boolean isToplevel() {
     return rfn.isToplevel();
   }
 
+  /** Whether this function is defined at the top level of a file. */
+  public boolean isGlobal() {
+    return module.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.