Perform builtins injection for WORKSPACE-loaded bzl files.

#17713

Closes #18022.

PiperOrigin-RevId: 525350732
Change-Id: I869653fcd6d4a82ccca49f14e66245c182062731

Author: benjaminp

src/main/java/com/google/devtools/build/lib/packages/BazelStarlarkEnvironment.java

@@ -50,12 +50,14 @@ public final class BazelStarlarkEnvironment {
   private final ImmutableMap<String, Object> uninjectedBuildBzlNativeBindings;
   /** The "native" module fields for a WORKSPACE-loaded bzl module. */
   private final ImmutableMap<String, Object> workspaceBzlNativeBindings;
-  /** The top-level predeclared symbols for a BUILD-loaded bzl module, before builtins injection. */
+  /** The top-level predeclared symbols for a BUILD-loaded bzl module before builtins injection. */
   private final ImmutableMap<String, Object> uninjectedBuildBzlEnv;
   /** The top-level predeclared symbols for BUILD files, before builtins injection and prelude. */
   private final ImmutableMap<String, Object> uninjectedBuildEnv;
-  /** The top-level predeclared symbols for a WORKSPACE-loaded bzl module. */
-  private final ImmutableMap<String, Object> workspaceBzlEnv;
+  /**
+   * The top-level predeclared symbols for a WORKSPACE-loaded bzl module before builtins injection.
+   */
+  private final ImmutableMap<String, Object> uninjectedWorkspaceBzlEnv;
   /** The top-level predeclared symbols for a bzl module in the {@code @_builtins} pseudo-repo. */
   private final ImmutableMap<String, Object> builtinsBzlEnv;
   /** The top-level predeclared symbols for a bzl module in the Bzlmod system. */
@@ -75,7 +77,8 @@ public final class BazelStarlarkEnvironment {
     this.workspaceBzlNativeBindings = createWorkspaceBzlNativeBindings(ruleClassProvider, version);
     this.uninjectedBuildBzlEnv =
         createUninjectedBuildBzlEnv(ruleClassProvider, uninjectedBuildBzlNativeBindings);
-    this.workspaceBzlEnv = createWorkspaceBzlEnv(ruleClassProvider, workspaceBzlNativeBindings);
+    this.uninjectedWorkspaceBzlEnv =
+        createWorkspaceBzlEnv(ruleClassProvider, workspaceBzlNativeBindings);
     // TODO(pcloudy): this should be a bzlmod specific environment, but keep using the workspace
     // envirnment until we implement module rules.
     this.bzlmodBzlEnv = createWorkspaceBzlEnv(ruleClassProvider, workspaceBzlNativeBindings);
@@ -122,9 +125,9 @@ public ImmutableMap<String, Object> getUninjectedBuildEnv() {
     return uninjectedBuildEnv;
   }
 
-  /** Returns the environment for WORKSPACE-loaded bzl files. */
-  public ImmutableMap<String, Object> getWorkspaceBzlEnv() {
-    return workspaceBzlEnv;
+  /** Returns the environment for WORKSPACE-loaded bzl files before builtins injection. */
+  public ImmutableMap<String, Object> getUninjectedWorkspaceBzlEnv() {
+    return uninjectedWorkspaceBzlEnv;
   }
 
   /** Returns the environment for bzl files in the {@code @_builtins} pseudo-repository. */
@@ -338,17 +341,57 @@ private static boolean injectionApplies(String key, Map<String, Boolean> overrid
    * documentation for {@code --experimental_builtins_injection_override}. Non-injected symbols must
    * still obey the above constraints.
    *
-   * @see StarlarkBuiltinsFunction
+   * @see com.google.devtools.build.lib.skyframe.StarlarkBuiltinsFunction
    */
   public ImmutableMap<String, Object> createBuildBzlEnvUsingInjection(
       Map<String, Object> exportedToplevels,
       Map<String, Object> exportedRules,
       List<String> overridesList)
       throws InjectionException {
+    return createBzlEnvUsingInjection(
+        exportedToplevels, exportedRules, overridesList, uninjectedBuildBzlNativeBindings);
+  }
+
+  /**
+   * Constructs an environment for a WORKSPACE-loaded bzl file based on the default environment, the
+   * maps corresponding to the {@code exported_toplevels} and {@code exported_rules} dicts, and the
+   * value of {@code --experimental_builtins_injection_override}.
+   *
+   * @see com.google.devtools.build.lib.skyframe.StarlarkBuiltinsFunction
+   */
+  public ImmutableMap<String, Object> createWorkspaceBzlEnvUsingInjection(
+      Map<String, Object> exportedToplevels,
+      Map<String, Object> exportedRules,
+      List<String> overridesList)
+      throws InjectionException {
+    return createBzlEnvUsingInjection(
+        exportedToplevels, exportedRules, overridesList, workspaceBzlNativeBindings);
+  }
+
+  private ImmutableMap<String, Object> createBzlEnvUsingInjection(
+      Map<String, Object> exportedToplevels,
+      Map<String, Object> exportedRules,
+      List<String> overridesList,
+      Map<String, Object> nativeBase)
+      throws InjectionException {
     Map<String, Boolean> overridesMap = parseInjectionOverridesList(overridesList);
 
+    Map<String, Object> env = new HashMap<>(ruleClassProvider.getEnvironment());
+
+    // Determine "native" bindings.
+    // TODO(#11954): See above comment in createUninjectedBuildBzlEnv.
+    Map<String, Object> nativeBindings = new HashMap<>(nativeBase);
+    for (Map.Entry<String, Object> entry : exportedRules.entrySet()) {
+      String key = entry.getKey();
+      String name = getKeySuffix(key);
+      validateSymbolIsInjectable(name, nativeBindings.keySet(), ruleFunctions.keySet(), "rule");
+      if (injectionApplies(key, overridesMap)) {
+        nativeBindings.put(name, entry.getValue());
+      }
+    }
+    env.put("native", createNativeModule(nativeBindings));
+
     // Determine top-level symbols.
-    Map<String, Object> env = new HashMap<>(uninjectedBuildBzlEnv);
     for (Map.Entry<String, Object> entry : exportedToplevels.entrySet()) {
       String key = entry.getKey();
       String name = getKeySuffix(key);
@@ -362,19 +405,6 @@ public ImmutableMap<String, Object> createBuildBzlEnvUsingInjection(
       }
     }
 
-    // Determine "native" bindings.
-    // TODO(#11954): See above comment in createUninjectedBuildBzlEnv.
-    Map<String, Object> nativeBindings = new HashMap<>(uninjectedBuildBzlNativeBindings);
-    for (Map.Entry<String, Object> entry : exportedRules.entrySet()) {
-      String key = entry.getKey();
-      String name = getKeySuffix(key);
-      validateSymbolIsInjectable(name, nativeBindings.keySet(), ruleFunctions.keySet(), "rule");
-      if (injectionApplies(key, overridesMap)) {
-        nativeBindings.put(name, entry.getValue());
-      }
-    }
-
-    env.put("native", createNativeModule(nativeBindings));
     return ImmutableMap.copyOf(env);
   }
 

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

@@ -572,21 +572,20 @@ private BzlLoadValue computeInternal(
   /**
    * Obtain a suitable StarlarkBuiltinsValue.
    *
-   * <p>For BUILD-loaded .bzl files, this is a real builtins value, obtained using either Skyframe
-   * or inlining of StarlarkBuiltinsFunction (depending on whether {@code inliningState} is
-   * non-null). The returned value includes the StarlarkSemantics.
+   * <p>For BUILD-loaded and WORKSPACE-loaded .bzl files, this is a real builtins value, obtained
+   * using either Skyframe or inlining of StarlarkBuiltinsFunction (depending on whether {@code
+   * inliningState} is non-null). The returned value includes the StarlarkSemantics.
    *
    * <p>For other .bzl files, the builtins computation is not needed and would create a Skyframe
    * cycle if requested, so we instead return an empty builtins value that just wraps the
-   * StarlarkSemantics. (NB: In the case of WORKSPACE-loaded .bzl files, the cycle goes through the
-   * repository remapping value. It's possible this could be avoided if we ever wanted to make this
-   * kind of .bzl file use builtins injection.)
+   * StarlarkSemantics.
    */
   @Nullable
   private StarlarkBuiltinsValue getBuiltins(
       BzlLoadValue.Key key, Environment env, @Nullable InliningState inliningState)
       throws BzlLoadFailedException, InterruptedException {
-    if (!(key instanceof BzlLoadValue.KeyForBuild)) {
+    if (!(key instanceof BzlLoadValue.KeyForBuild)
+        && !(key instanceof BzlLoadValue.KeyForWorkspace)) {
       StarlarkSemantics starlarkSemantics = PrecomputedValue.STARLARK_SEMANTICS.get(env);
       if (starlarkSemantics == null) {
         return null;
@@ -1185,7 +1184,15 @@ private ImmutableMap<String, Object> getAndDigestPredeclaredEnvironment(
       fp.addBytes(builtins.transitiveDigest);
       return builtins.predeclaredForBuildBzl;
     } else if (key instanceof BzlLoadValue.KeyForWorkspace) {
-      return starlarkEnv.getWorkspaceBzlEnv();
+      // TODO(#11437): Remove ability to disable injection by setting flag to empty string.
+      if (builtins
+          .starlarkSemantics
+          .get(BuildLanguageOptions.EXPERIMENTAL_BUILTINS_BZL_PATH)
+          .isEmpty()) {
+        return starlarkEnv.getUninjectedWorkspaceBzlEnv();
+      }
+      fp.addBytes(builtins.transitiveDigest);
+      return builtins.predeclaredForWorkspaceBzl;
     } else if (key instanceof BzlLoadValue.KeyForBzlmod) {
       return starlarkEnv.getBzlmodBzlEnv();
     } else if (key instanceof BzlLoadValue.KeyForBuiltins) {

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

@@ -174,12 +174,18 @@ private static StarlarkBuiltinsValue computeInternal(
               exportedToplevels,
               exportedRules,
               starlarkSemantics.get(BuildLanguageOptions.EXPERIMENTAL_BUILTINS_INJECTION_OVERRIDE));
+      ImmutableMap<String, Object> predeclaredForWorkspaceBzl =
+          starlarkEnv.createWorkspaceBzlEnvUsingInjection(
+              exportedToplevels,
+              exportedRules,
+              starlarkSemantics.get(BuildLanguageOptions.EXPERIMENTAL_BUILTINS_INJECTION_OVERRIDE));
       ImmutableMap<String, Object> predeclaredForBuild =
           starlarkEnv.createBuildEnvUsingInjection(
               exportedRules,
               starlarkSemantics.get(BuildLanguageOptions.EXPERIMENTAL_BUILTINS_INJECTION_OVERRIDE));
       return StarlarkBuiltinsValue.create(
           predeclaredForBuildBzl,
+          predeclaredForWorkspaceBzl,
           predeclaredForBuild,
           exportedToJava,
           transitiveDigest,

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

@@ -64,6 +64,12 @@ static boolean isBuiltinsRepo(RepositoryName repo) {
    */
   public final ImmutableMap<String, Object> predeclaredForBuildBzl;
 
+  /**
+   * Top-level predeclared symbols for a .bzl file loaded on behalf of a WORKSPACE file after
+   * builtins injection has been applied.
+   */
+  public final ImmutableMap<String, Object> predeclaredForWorkspaceBzl;
+
   /**
    * Top-level predeclared symbols for a BUILD file, after builtins injection but before any prelude
    * file has been applied.
@@ -81,11 +87,13 @@ static boolean isBuiltinsRepo(RepositoryName repo) {
 
   private StarlarkBuiltinsValue(
       ImmutableMap<String, Object> predeclaredForBuildBzl,
+      ImmutableMap<String, Object> predeclaredForWorkspaceBzl,
       ImmutableMap<String, Object> predeclaredForBuild,
       ImmutableMap<String, Object> exportedToJava,
       byte[] transitiveDigest,
       StarlarkSemantics starlarkSemantics) {
     this.predeclaredForBuildBzl = predeclaredForBuildBzl;
+    this.predeclaredForWorkspaceBzl = predeclaredForWorkspaceBzl;
     this.predeclaredForBuild = predeclaredForBuild;
     this.exportedToJava = exportedToJava;
     this.transitiveDigest = transitiveDigest;
@@ -94,12 +102,14 @@ private StarlarkBuiltinsValue(
 
   public static StarlarkBuiltinsValue create(
       ImmutableMap<String, Object> predeclaredForBuildBzl,
+      ImmutableMap<String, Object> predeclaredForWorkspaceBzl,
       ImmutableMap<String, Object> predeclaredForBuild,
       ImmutableMap<String, Object> exportedToJava,
       byte[] transitiveDigest,
       StarlarkSemantics starlarkSemantics) {
     return new StarlarkBuiltinsValue(
         predeclaredForBuildBzl,
+        predeclaredForWorkspaceBzl,
         predeclaredForBuild,
         exportedToJava,
         transitiveDigest,
@@ -116,10 +126,11 @@ public static StarlarkBuiltinsValue create(
    */
   public static StarlarkBuiltinsValue createEmpty(StarlarkSemantics starlarkSemantics) {
     return new StarlarkBuiltinsValue(
-        /*predeclaredForBuildBzl=*/ ImmutableMap.of(),
-        /*predeclaredForBuild=*/ ImmutableMap.of(),
-        /*exportedToJava=*/ ImmutableMap.of(),
-        /*transitiveDigest=*/ new byte[] {},
+        /* predeclaredForBuildBzl= */ ImmutableMap.of(),
+        /* predeclaredForWorkspaceBzl= */ ImmutableMap.of(),
+        /* predeclaredForBuild= */ ImmutableMap.of(),
+        /* exportedToJava= */ ImmutableMap.of(),
+        /* transitiveDigest= */ new byte[] {},
         starlarkSemantics);
   }
 

src/test/java/com/google/devtools/build/lib/packages/BazelStarlarkEnvironmentTest.java

@@ -61,7 +61,7 @@ protected ConfiguredRuleClassProvider createRuleClassProvider() {
   public void buildAndWorkspaceBzlEnvsDeclareSameNames() throws Exception {
     BazelStarlarkEnvironment starlarkEnv = pkgFactory.getBazelStarlarkEnvironment();
     Set<String> buildBzlNames = starlarkEnv.getUninjectedBuildBzlEnv().keySet();
-    Set<String> workspaceBzlNames = starlarkEnv.getWorkspaceBzlEnv().keySet();
+    Set<String> workspaceBzlNames = starlarkEnv.getUninjectedWorkspaceBzlEnv().keySet();
     assertThat(buildBzlNames).isEqualTo(workspaceBzlNames);
   }
 
@@ -72,7 +72,7 @@ public void buildAndWorkspaceBzlEnvsAreSameExceptForNative() throws Exception {
     buildBzlEnv.putAll(starlarkEnv.getUninjectedBuildBzlEnv());
     buildBzlEnv.remove("native");
     Map<String, Object> workspaceBzlEnv = new HashMap<>();
-    workspaceBzlEnv.putAll(starlarkEnv.getWorkspaceBzlEnv());
+    workspaceBzlEnv.putAll(starlarkEnv.getUninjectedWorkspaceBzlEnv());
     workspaceBzlEnv.remove("native");
     assertThat(buildBzlEnv).isEqualTo(workspaceBzlEnv);
   }

src/test/java/com/google/devtools/build/lib/rules/repository/RepositoryDelegatorTest.java

@@ -72,6 +72,7 @@
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.build.lib.skyframe.RepositoryMappingFunction;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
+import com.google.devtools.build.lib.skyframe.StarlarkBuiltinsFunction;
 import com.google.devtools.build.lib.skyframe.WorkspaceFileFunction;
 import com.google.devtools.build.lib.starlarkbuildapi.repository.RepositoryBootstrap;
 import com.google.devtools.build.lib.testutil.FoundationTestCase;
@@ -117,7 +118,9 @@ public class RepositoryDelegatorTest extends FoundationTestCase {
   public void setupDelegator() throws Exception {
     rootPath = scratch.dir("/outputbase");
     scratch.file(
-        rootPath.getRelative("MODULE.bazel").getPathString(), "module(name='test',version='0.1')");
+        rootPath.getRelative("MODULE.bazel").getPathString(),
+        "module(name='test',version='0.1')",
+        "bazel_dep(name='bazel_tools',version='1.0')");
     BlazeDirectories directories =
         new BlazeDirectories(
             new ServerDirectories(rootPath, rootPath, rootPath),
@@ -162,6 +165,13 @@ public void setupDelegator() throws Exception {
             .build(ruleClassProvider, fileSystem);
 
     registryFactory = new FakeRegistry.Factory();
+    FakeRegistry registry =
+        registryFactory
+            .newFakeRegistry(scratch.dir("modules").getPathString())
+            .addModule(
+                createModuleKey("bazel_tools", "1.0"),
+                "module(name='bazel_tools', version='1.0');");
+    ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
 
     HashFunction hashFunction = fileSystem.getDigestFunction().getHashFunction();
     evaluator =
@@ -217,6 +227,7 @@ public void setupDelegator() throws Exception {
                     SkyFunctions.BZL_LOAD,
                     BzlLoadFunction.create(
                         pkgFactory, directories, hashFunction, Caffeine.newBuilder().build()))
+                .put(SkyFunctions.STARLARK_BUILTINS, new StarlarkBuiltinsFunction(pkgFactory))
                 .put(SkyFunctions.CONTAINING_PACKAGE_LOOKUP, new ContainingPackageLookupFunction())
                 .put(
                     SkyFunctions.IGNORED_PACKAGE_PREFIXES,
@@ -388,14 +399,6 @@ public void loadRepositoryFromBzlmod() throws Exception {
         rootPath.getRelative("MODULE.bazel").getPathString(),
         "module(name='aaa',version='0.1')",
         "bazel_dep(name='bazel_tools',version='1.0')");
-    FakeRegistry registry =
-        registryFactory
-            .newFakeRegistry(scratch.dir("modules").getPathString())
-            .addModule(
-                createModuleKey("bazel_tools", "1.0"),
-                "module(name='bazel_tools', version='1.0');");
-    ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
-    // Note that bazel_tools is a well-known module, so its repo name will always be "bazel_tools".
     scratch.file(rootPath.getRelative("BUILD").getPathString());
     scratch.file(
         rootPath.getRelative("repo_rule.bzl").getPathString(),

src/test/java/com/google/devtools/build/lib/skyframe/BuiltinsInjectionTest.java

@@ -485,13 +485,8 @@ public void exportsBzlMayBeInErrorWhenInjectionIsDisabled() throws Exception {
     assertContainsEvent("In BUILD: overridable_rule :: <built-in rule overridable_rule>");
   }
 
-  // TODO(#11954): Once WORKSPACE- and BUILD-loaded bzls use the exact same environments, we'll want
-  // to apply injection to both. This is for uniformity, not because we actually care about builtins
-  // injection for WORKSPACE bzls. In the meantime, assert the status quo: WORKSPACE bzls do not use
-  // injection. WORKSPACE and BUILD files themselves probably won't be unified, so WORKSPACE will
-  // likely continue to not use injection.
   @Test
-  public void workspaceAndWorkspaceBzlDoNotUseInjection() throws Exception {
+  public void workspaceLoadedBzlUsesInjectionButNotWORKSPACE() throws Exception {
     writeExportsBzl(
         "exported_toplevels = {'overridable_symbol': 'new_value'}",
         "exported_rules = {'overridable_rule': 'new_rule'}",
@@ -510,9 +505,8 @@ public void workspaceAndWorkspaceBzlDoNotUseInjection() throws Exception {
         "print('In bzl: overridable_symbol :: %s' % overridable_symbol)");
 
     buildAndAssertSuccess();
-    // Builtins for WORKSPACE bzls are populated essentially the same as for BUILD bzls, except that
-    // injection doesn't apply.
-    assertContainsEvent("In bzl: overridable_symbol :: original_value");
+    // Builtins for WORKSPACE bzls are populated the same as for BUILD bzls.
+    assertContainsEvent("In bzl: overridable_symbol :: new_value");
     // We don't assert that the rule isn't injected because the workspace native object doesn't
     // contain our original mock rule. We can test this for WORKSPACE files at the top-level though.
     assertContainsEvent("In WORKSPACE: overridable_rule :: <built-in function overridable_rule>");