Add BazelModuleContext to BzlCompileFunction

Without it, the guard is evaluated during compilation and the modules removed before they get to load stage.

PiperOrigin-RevId: 476296931
Change-Id: I46d08ff8e6fed5f00acd0a3799da46c48383dd06

src/main/java/com/google/devtools/build/lib/cmdline/BUILD

@@ -12,6 +12,7 @@ filegroup(
 java_library(
     name = "cmdline",
     srcs = [
+        "BazelCompileContext.java",
         "BazelModuleContext.java",
         "Label.java",
         "LabelConstants.java",

src/main/java/com/google/devtools/build/lib/cmdline/BazelCompileContext.java

@@ -0,0 +1,42 @@
+// Copyright 2022 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.cmdline;
+
+import com.google.auto.value.AutoValue;
+import javax.annotation.Nullable;
+import net.starlark.java.eval.Module;
+
+/**
+ * BazelCompileContext records Bazel-specific information associated with a .bzl {@link
+ * net.starlark.java.eval.Module} during bzl compilation.
+ *
+ * <p>Maintainer's note: This object is determined prior to the module's compilation in
+ * BzlCompileFunction. It is saved in the {@code Module} used for compilation as {@link
+ * Module#getClientData client data}. The {@code Module} used during .bzl evaluation is separate and
+ * uses {@link BazelModuleContext} as client data.
+ */
+@AutoValue
+public abstract class BazelCompileContext {
+  /** Label associated with the Starlark {@link net.starlark.java.eval.Module}. */
+  @Nullable
+  public abstract Label label();
+
+  /** Returns the name of the module's .bzl file, as provided to the parser. */
+  public abstract String filename();
+
+  public static BazelCompileContext create(Label label, String filename) {
+    return new AutoValue_BazelCompileContext(label, filename);
+  }
+}

src/main/java/com/google/devtools/build/lib/cmdline/BazelModuleContext.java

@@ -22,10 +22,11 @@
  * BazelModuleContext records Bazel-specific information associated with a .bzl {@link
  * net.starlark.java.eval.Module}.
  *
- * <p>Maintainer's note: This object is determined prior to the module's evaluation. It is saved in
- * the {@code Module} as {@link Module#getClientData client data}. For data that is computed after
- * the module's evaluation and which need not be exposed outside the module-loading machinery,
- * consider {@link BzlLoadValue}.
+ * <p>Maintainer's note: This object is determined prior to the module's evaluation in
+ * BzlLoadFunction. It is saved in the {@code Module} as {@link Module#getClientData client data}.
+ * The {@code Module} used during .bzl compilation is separate and uses {@link BazelCompileContext}
+ * as client data. For data that is computed after the module's evaluation and which need not be
+ * exposed outside the module-loading machinery, consider {@link BzlLoadValue}.
  */
 // Immutability is useful because this object is retrievable from a Module and therefore from a
 // BzlLoadValue.

src/main/java/com/google/devtools/build/lib/skyframe/BzlCompileFunction.java

@@ -16,6 +16,7 @@
 
 import com.google.common.hash.HashFunction;
 import com.google.devtools.build.lib.actions.FileValue;
+import com.google.devtools.build.lib.cmdline.BazelCompileContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.BazelStarlarkEnvironment;
@@ -156,6 +157,7 @@ static BzlCompileValue computeInline(
       // because the "native" object is different. But A) that will be fixed with #11954, and B) we
       // don't care for the same reason as above.
       predeclared = starlarkEnv.getUninjectedBuildBzlEnv();
+
     }
 
     // We have all deps. Parse, resolve, and return.
@@ -171,7 +173,20 @@ static BzlCompileValue computeInline(
     StarlarkFile file = StarlarkFile.parse(input, options);
 
     // compile
-    Module module = Module.withPredeclared(semantics, predeclared);
+    final Module module;
+
+    if (key.kind == BzlCompileValue.Kind.EMPTY_PRELUDE) {
+      // The empty prelude has no label, so we can't use it to filter the predeclareds.
+      // This doesn't matter since the empty prelude doesn't attempt to access any predeclareds
+      // anyway.
+      module = Module.withPredeclared(semantics, predeclared);
+    } else {
+      // The BazelCompileContext holds additional contextual info to be associated with the Module
+      // The information is used to filter predeclareds
+      BazelCompileContext bazelCompileContext =
+          BazelCompileContext.create(key.label, file.getName());
+      module = Module.withPredeclaredAndData(semantics, predeclared, bazelCompileContext);
+    }
     try {
       Program prog = Program.compileFile(file, module);
       return BzlCompileValue.withProgram(prog, digest);

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/core/ContextAndFlagGuardedValue.java

@@ -15,7 +15,7 @@
 package com.google.devtools.build.lib.starlarkbuildapi.core;
 
 import com.google.common.collect.ImmutableSet;
-import com.google.devtools.build.lib.cmdline.BazelModuleContext;
+import com.google.devtools.build.lib.cmdline.BazelCompileContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import java.util.stream.Collectors;
@@ -46,8 +46,10 @@ public static GuardedValue onlyInAllowedReposOrWhenIncompatibleFlagIsFalse(
       public boolean isObjectAccessibleUsingSemantics(
           StarlarkSemantics semantics, @Nullable Object clientData) {
         boolean accessible = flagGuard.isObjectAccessibleUsingSemantics(semantics, clientData);
-        if (!accessible && clientData != null) {
-          BazelModuleContext context = (BazelModuleContext) clientData;
+        // Filtering of predeclareds is only done at compile time, when the client data is
+        // BazelCompileContext and not BazelModuleContext.
+        if (!accessible && clientData != null && clientData instanceof BazelCompileContext) {
+          BazelCompileContext context = (BazelCompileContext) clientData;
           Label label = context.label();
 
           for (PackageIdentifier prefix : allowedPrefixes) {