Support Starlark fragments in subrules

Unlike attributes or toolchains, there is fortunately no need to lift these to the rule/aspect class, since these can be looked up directly on the configuration, without any extra Skyframe evaluation. PiperOrigin-RevId: 580786808 Change-Id: Ic3e414b9498b3d9b5b757ade3b5007e4e69260ae
bazel-io · Nov 9, 2023 · bdcf7ce · bdcf7ce
1 parent 9beece3
commit bdcf7ce
Show file tree

Hide file tree

Showing 8 changed files with 247 additions and 18 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/FragmentCollection.java b/src/main/java/com/google/devtools/build/lib/analysis/config/FragmentCollection.java
@@ -13,7 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.analysis.config;
 
-import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableCollection;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
@@ -47,18 +46,6 @@ public ImmutableCollection<String> getFieldNames() {
     return ruleContext.getStarlarkFragmentNames();
   }
 
-  @Override
-  @Nullable
-  public String getErrorMessageForUnknownField(String name) {
-    return String.format(
-        "There is no configuration fragment named '%s'. Available fragments: %s",
-        name, fieldsToString());
-  }
-
-  private String fieldsToString() {
-    return String.format("'%s'", Joiner.on("', '").join(getFieldNames()));
-  }
-
   @Override
   public String toString() {
     return "[ " + fieldsToString() + "]";

diff --git a/...main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/...main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
@@ -1256,13 +1256,16 @@ public StarlarkSubruleApi subrule(
       StarlarkFunction implementation,
       Dict<?, ?> attrsUnchecked,
       Sequence<?> toolchainsUnchecked,
+      Sequence<?> fragmentsUnchecked,
       StarlarkThread thread)
       throws EvalException {
     if (!thread.getSemantics().getBool(BuildLanguageOptions.EXPERIMENTAL_RULE_EXTENSION_API)) {
       BuiltinRestriction.failIfCalledOutsideAllowlist(thread, ALLOWLIST_RULE_EXTENSION_API);
     }
     ImmutableMap<String, Descriptor> attrs =
         ImmutableMap.copyOf(Dict.cast(attrsUnchecked, String.class, Descriptor.class, "attrs"));
+    ImmutableList<String> fragments =
+        Sequence.noneableCast(fragmentsUnchecked, String.class, "fragments").getImmutableList();
     for (Entry<String, Descriptor> attr : attrs.entrySet()) {
       String attrName = attr.getKey();
       Descriptor descriptor = attr.getValue();
@@ -1296,7 +1299,7 @@ public StarlarkSubruleApi subrule(
     if (toolchains.size() > 1) {
       throw Starlark.errorf("subrules may require at most 1 toolchain, got: %s", toolchains);
     }
-    return new StarlarkSubrule(implementation, attrs, toolchains);
+    return new StarlarkSubrule(implementation, attrs, toolchains, ImmutableSet.copyOf(fragments));
   }
 
   private static ImmutableSet<ToolchainTypeRequirement> parseToolchainTypes(

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java
@@ -19,6 +19,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -29,6 +30,7 @@
 import com.google.devtools.build.lib.analysis.BazelRuleAnalysisThreadContext;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.RuleContext;
+import com.google.devtools.build.lib.analysis.config.Fragment;
 import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.analysis.starlark.StarlarkActionFactory.StarlarkActionContext;
 import com.google.devtools.build.lib.analysis.starlark.StarlarkAttrModule.Descriptor;
@@ -39,6 +41,7 @@
 import com.google.devtools.build.lib.packages.AttributeValueSource;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.StarlarkExportable;
+import com.google.devtools.build.lib.starlarkbuildapi.FragmentCollectionApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkActionFactoryApi;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import com.google.devtools.build.lib.starlarkbuildapi.platform.ToolchainContextApi;
@@ -71,6 +74,7 @@ public class StarlarkSubrule implements StarlarkExportable, StarlarkCallable, St
 
   private final StarlarkFunction implementation;
   private final ImmutableSet<ToolchainTypeRequirement> toolchains;
+  private final ImmutableSet<String> fragments;
 
   // following fields are set on export
   @Nullable private String exportedName = null;
@@ -79,10 +83,12 @@ public class StarlarkSubrule implements StarlarkExportable, StarlarkCallable, St
   public StarlarkSubrule(
       StarlarkFunction implementation,
       ImmutableMap<String, Descriptor> attributes,
-      ImmutableSet<ToolchainTypeRequirement> toolchains) {
+      ImmutableSet<ToolchainTypeRequirement> toolchains,
+      ImmutableSet<String> fragments) {
     this.implementation = implementation;
     this.attributes = SubruleAttribute.from(attributes);
     this.toolchains = toolchains;
+    this.fragments = fragments;
   }
 
   @Override
@@ -154,7 +160,7 @@ public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwarg
       namedArgs.put(attr.attrName, value == null ? Starlark.NONE : value);
     }
     SubruleContext subruleContext =
-        new SubruleContext(ruleContext, toolchains, runfilesFromDeps.build());
+        new SubruleContext(this, ruleContext, toolchains, runfilesFromDeps.build());
     ImmutableList<Object> positionals =
         ImmutableList.builder().add(subruleContext).addAll(args).build();
     try {
@@ -269,22 +275,27 @@ static ImmutableSet<ToolchainTypeRequirement> discoverToolchains(
   private static class SubruleContext implements StarlarkActionContext {
     // these fields are effectively final, set to null once this instance is no longer usable by
     // Starlark
+    private StarlarkSubrule subrule;
     private StarlarkRuleContext starlarkRuleContext;
     private ImmutableSet<Label> requestedToolchains;
     private ImmutableSet<FilesToRunProvider> runfilesFromDeps;
     private StarlarkActionFactory actions;
+    private FragmentCollectionApi fragmentCollection;
 
     private SubruleContext(
+        StarlarkSubrule subrule,
         StarlarkRuleContext ruleContext,
         ImmutableSet<ToolchainTypeRequirement> requestedToolchains,
         ImmutableSet<FilesToRunProvider> runfilesFromDeps) {
+      this.subrule = subrule;
       this.starlarkRuleContext = ruleContext;
       this.requestedToolchains =
           requestedToolchains.stream()
               .map(ToolchainTypeRequirement::toolchainType)
               .collect(toImmutableSet());
       this.runfilesFromDeps = runfilesFromDeps;
       this.actions = new StarlarkActionFactory(this);
+      this.fragmentCollection = new SubruleFragmentCollection(this);
     }
 
     @StarlarkMethod(
@@ -335,6 +346,15 @@ private ImmutableSet<Label> getAutomaticExecGroupLabels() {
           .collect(toImmutableSet());
     }
 
+    @StarlarkMethod(
+        name = "fragments",
+        doc = "Allows access to configuration fragments in target configuration.",
+        structField = true)
+    public FragmentCollectionApi getFragmentCollection() throws EvalException {
+      checkMutable("fragments");
+      return fragmentCollection;
+    }
+
     @Override
     public ArtifactRoot newFileRoot() {
       return starlarkRuleContext.getRuleContext().getBinDirectory();
@@ -406,10 +426,12 @@ public Object maybeOverrideToolchain(Object toolchainUnchecked) throws EvalExcep
     }
 
     private void nullify() {
+      this.subrule = null;
       this.starlarkRuleContext = null;
       this.actions = null;
       this.requestedToolchains = null;
       this.runfilesFromDeps = null;
+      this.fragmentCollection = null;
     }
   }
 
@@ -470,4 +492,41 @@ static String getRuleAttrName(
     return valueSource.convertToNativeName(
         "_" + label.getCanonicalForm() + "%" + exportedName + "%" + attrName);
   }
+
+  private static class SubruleFragmentCollection implements FragmentCollectionApi {
+
+    private final SubruleContext subruleContext;
+
+    private SubruleFragmentCollection(SubruleContext subruleContext) {
+      this.subruleContext = subruleContext;
+    }
+
+    @Override
+    @Nullable
+    public Object getValue(String name) throws EvalException {
+      Class<? extends Fragment> fragmentClass =
+          subruleContext.getRuleContext().getConfiguration().getStarlarkFragmentByName(name);
+      if (fragmentClass == null) {
+        return null;
+      }
+      if (!subruleContext.subrule.fragments.contains(name)) {
+        throw Starlark.errorf(
+            "%s has to declare '%s' as a required fragment in order to access it."
+                + " Please update the 'fragments' argument of the subrule definition "
+                + "(for example: fragments = [\"%s\"])",
+            subruleContext.subrule.getName(), name, name);
+      }
+      return subruleContext.getRuleContext().getConfiguration().getFragment(fragmentClass);
+    }
+
+    @Override
+    public ImmutableCollection<String> getFieldNames() {
+      return subruleContext.subrule.fragments;
+    }
+
+    @Override
+    public String toString() {
+      return "[ " + fieldsToString() + "]";
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/FragmentCollectionApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/FragmentCollectionApi.java
@@ -14,7 +14,9 @@
 
 package com.google.devtools.build.lib.starlarkbuildapi;
 
+import com.google.common.base.Joiner;
 import com.google.devtools.build.docgen.annot.DocCategory;
+import javax.annotation.Nullable;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.StarlarkValue;
 import net.starlark.java.eval.Structure;
@@ -32,4 +34,17 @@
             + " fragment reference</a> for a list of available fragments and the <a"
             + " href=\"https://bazel.build/extending/rules#configuration_fragments\">rules"
             + " documentation</a> for how to use them.")
-public interface FragmentCollectionApi extends Structure, StarlarkValue {}
+public interface FragmentCollectionApi extends Structure, StarlarkValue {
+
+  @Override
+  @Nullable
+  default String getErrorMessageForUnknownField(String name) {
+    return String.format(
+        "There is no configuration fragment named '%s'. Available fragments: %s",
+        name, fieldsToString());
+  }
+
+  default String fieldsToString() {
+    return String.format("'%s'", Joiner.on("', '").join(getFieldNames()));
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkRuleFunctionsApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkRuleFunctionsApi.java
@@ -840,13 +840,23 @@ ExecGroupApi execGroup(
                 "If set, the set of toolchains this subrule requires. The list can contain String,"
                     + " Label, or StarlarkToolchainTypeApi objects, in any combination. Toolchains"
                     + " will be found by checking the current platform, and provided to the subrule"
-                    + " implementation via <code>ctx.toolchains</code>.")
+                    + " implementation via <code>ctx.toolchains</code>."),
+        @Param(
+            name = "fragments",
+            allowedTypes = {@ParamType(type = Sequence.class, generic1 = String.class)},
+            named = true,
+            positional = false,
+            defaultValue = "[]",
+            doc =
+                "List of names of configuration fragments that the subrule requires in target"
+                    + " configuration.")
       },
       useStarlarkThread = true)
   StarlarkSubruleApi subrule(
       StarlarkFunction implementation,
       Dict<?, ?> attrs,
       Sequence<?> toolchains,
+      Sequence<?> fragments,
       StarlarkThread thread)
       throws EvalException;
 }
diff --git a/...main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeStarlarkRuleFunctionsApi.java b/...main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeStarlarkRuleFunctionsApi.java
@@ -251,6 +251,7 @@ public StarlarkSubruleApi subrule(
       StarlarkFunction implementation,
       Dict<?, ?> attrs,
       Sequence<?> toolchains,
+      Sequence<?> fragments,
       StarlarkThread thread) {
     return new FakeStarlarkSubrule();
   }

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/BUILD b/src/test/java/com/google/devtools/build/lib/analysis/BUILD
@@ -491,6 +491,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/rules/java:java-compilation",
         "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi",
+        "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/cpp",
         "//src/main/java/net/starlark/java/eval",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//src/test/java/com/google/devtools/build/lib/starlark/util",