From bdcf7ce2ac8e83e2e9569821eb5c9651c45b134a Mon Sep 17 00:00:00 2001
From: Googler <hvd@google.com>
Date: Wed, 8 Nov 2023 22:45:57 -0800
Subject: [PATCH] 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
---
 .../analysis/config/FragmentCollection.java   |  13 --
 .../starlark/StarlarkRuleClassFunctions.java  |   5 +-
 .../analysis/starlark/StarlarkSubrule.java    |  63 +++++++-
 .../FragmentCollectionApi.java                |  17 +-
 .../StarlarkRuleFunctionsApi.java             |  12 +-
 .../FakeStarlarkRuleFunctionsApi.java         |   1 +
 .../google/devtools/build/lib/analysis/BUILD  |   1 +
 .../starlark/StarlarkSubruleTest.java         | 153 ++++++++++++++++++
 8 files changed, 247 insertions(+), 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
index 4c2d13ef4af00d..a99110232c0eb7 100644
--- 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/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
index 60e4acead8d247..62e79f05de39da 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
@@ -1256,6 +1256,7 @@ public StarlarkSubruleApi subrule(
       StarlarkFunction implementation,
       Dict<?, ?> attrsUnchecked,
       Sequence<?> toolchainsUnchecked,
+      Sequence<?> fragmentsUnchecked,
       StarlarkThread thread)
       throws EvalException {
     if (!thread.getSemantics().getBool(BuildLanguageOptions.EXPERIMENTAL_RULE_EXTENSION_API)) {
@@ -1263,6 +1264,8 @@ public StarlarkSubruleApi subrule(
     }
     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
index fda5f4f5f15ba1..3a47a007bdd0ae 100644
--- 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,15 +275,19 @@ 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()
@@ -285,6 +295,7 @@ private SubruleContext(
               .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
index 49feea2c6d37cf..50f217da3056a5 100644
--- 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
index 5435f98651aead..6497f570597fda 100644
--- 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/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeStarlarkRuleFunctionsApi.java b/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeStarlarkRuleFunctionsApi.java
index 8a00349ea1c509..458925c1bf4d59 100644
--- a/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeStarlarkRuleFunctionsApi.java
+++ b/src/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
index cf37909d3fb0a5..c724ff93ed892f 100644
--- 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",
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java b/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java
index f11c53405bdeab..4f405f1a577053 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java
@@ -39,6 +39,7 @@
 import com.google.devtools.build.lib.rules.java.JavaToolchainProvider;
 import com.google.devtools.build.lib.starlark.util.BazelEvaluationTestCase;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
+import com.google.devtools.build.lib.starlarkbuildapi.cpp.CppConfigurationApi;
 import com.google.devtools.build.lib.testutil.TestConstants;
 import net.starlark.java.eval.BuiltinFunction;
 import net.starlark.java.eval.Sequence;
@@ -1020,6 +1021,158 @@ public void testSubruleToolchains_requestedToolchainIsSuppliedToAction() throws
         .isEqualTo(getRuleContext(target).getActionOwner(TestConstants.JAVA_TOOLCHAIN_TYPE));
   }
 
+  @Test
+  public void testSubruleFragments_errorForInvalidFragments() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "def _subrule_impl(ctx):",
+        "  return ctx.fragments.foobar",
+        "_my_subrule = subrule(",
+        "  implementation = _subrule_impl, fragments = ['java', 'cpp']",
+        ")",
+        "MyInfo = provider()",
+        "def _rule_impl(ctx):",
+        "  res = _my_subrule()",
+        "  return MyInfo(result = res)",
+        "",
+        "my_rule = rule(_rule_impl, subrules = [_my_subrule])");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "my_rule(name = 'foo')");
+
+    AssertionError assertionError =
+        assertThrows(AssertionError.class, () -> getConfiguredTarget("//subrule_testing:foo"));
+
+    assertThat(assertionError)
+        .hasMessageThat()
+        .contains(
+            "There is no configuration fragment named 'foobar'. Available fragments: 'java',"
+                + " 'cpp'");
+  }
+
+  @Test
+  public void testSubruleFragments_onlyDeclaredFragmentsAreVisible() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "def _subrule_impl(ctx):",
+        "  return dir(ctx.fragments)",
+        "_my_subrule = subrule(",
+        "  implementation = _subrule_impl,",
+        "  fragments = ['cpp', 'python']",
+        ")",
+        "MyInfo = provider()",
+        "def _rule_impl(ctx):",
+        "  res = _my_subrule()",
+        "  return MyInfo(result = res)",
+        "",
+        "my_rule = rule(_rule_impl, subrules = [_my_subrule], fragments = ['java'])");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "my_rule(name = 'foo')");
+
+    Sequence<String> fragments =
+        Sequence.cast(
+            getProvider("//subrule_testing:foo", "//subrule_testing:myrule.bzl", "MyInfo")
+                .getValue("result"),
+            String.class,
+            "ctx.fragments");
+
+    assertThat(fragments).isNotNull();
+    assertThat(fragments).containsExactly("cpp", "python");
+  }
+
+  @Test
+  public void testSubruleFragments_ruleCannotAccessSubruleFragments() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "def _subrule_impl(ctx):",
+        "  pass",
+        "_my_subrule = subrule(",
+        "  implementation = _subrule_impl,",
+        "  fragments = ['cpp']",
+        ")",
+        "def _rule_impl(ctx):",
+        "  ctx.fragments.cpp",
+        "  return []",
+        "",
+        "my_rule = rule(_rule_impl, subrules = [_my_subrule])");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "my_rule(name = 'foo')");
+
+    AssertionError assertionError =
+        assertThrows(AssertionError.class, () -> getConfiguredTarget("//subrule_testing:foo"));
+
+    assertThat(assertionError)
+        .hasMessageThat()
+        .contains("my_rule has to declare 'cpp' as a required fragment in order to access it");
+  }
+
+  @Test
+  public void testSubruleFragments_canAccessDeclaredFragments() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "def _subrule_impl(ctx):",
+        "  return ctx.fragments.cpp",
+        "_my_subrule = subrule(",
+        "  implementation = _subrule_impl,",
+        "  fragments = ['cpp']",
+        ")",
+        "MyInfo = provider()",
+        "def _rule_impl(ctx):",
+        "  res = _my_subrule()",
+        "  return MyInfo(result = res)",
+        "",
+        "my_rule = rule(_rule_impl, subrules = [_my_subrule], fragments = ['java'])");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "my_rule(name = 'foo')");
+
+    CppConfigurationApi<?> fragment =
+        getProvider("//subrule_testing:foo", "//subrule_testing:myrule.bzl", "MyInfo")
+            .getValue("result", CppConfigurationApi.class);
+
+    assertThat(fragment).isNotNull();
+  }
+
+  @Test
+  public void testSubruleFragments_mustDeclareFragmentsIfAccessed() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "def _subrule_impl(ctx):",
+        "  ctx.fragments.java",
+        "_my_subrule = subrule(",
+        "  implementation = _subrule_impl,",
+        "  fragments = ['cpp', 'python']",
+        ")",
+        "",
+        "def _rule_impl(ctx):",
+        "  res = _my_subrule()",
+        "  return []",
+        "",
+        "my_rule = rule(_rule_impl, subrules = [_my_subrule], fragments = ['java'])");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "my_rule(name = 'foo')");
+
+    AssertionError assertionError =
+        assertThrows(AssertionError.class, () -> getConfiguredTarget("//subrule_testing:foo"));
+
+    assertThat(assertionError)
+        .hasMessageThat()
+        .contains("_my_subrule has to declare 'java' as a required fragment in order to access it");
+  }
+
   @Test
   public void testSubruleInstantiation_outsideAllowlist_failsWithPrivateAPIError()
       throws Exception {