Add support for subrules to depend on other subrules

PiperOrigin-RevId: 581245635
Change-Id: I258c0e874af78451a9cd3fc6a4228750a60485a9

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java

@@ -1270,6 +1270,7 @@ public StarlarkSubruleApi subrule(
       Dict<?, ?> attrsUnchecked,
       Sequence<?> toolchainsUnchecked,
       Sequence<?> fragmentsUnchecked,
+      Sequence<?> subrulesUnchecked,
       StarlarkThread thread)
       throws EvalException {
     if (!thread.getSemantics().getBool(BuildLanguageOptions.EXPERIMENTAL_RULE_EXTENSION_API)) {
@@ -1312,7 +1313,12 @@ 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, ImmutableSet.copyOf(fragments));
+    return new StarlarkSubrule(
+        implementation,
+        attrs,
+        toolchains,
+        ImmutableSet.copyOf(fragments),
+        ImmutableSet.copyOf(Sequence.cast(subrulesUnchecked, StarlarkSubrule.class, "subrules")));
   }
 
   private static ImmutableSet<ToolchainTypeRequirement> parseToolchainTypes(

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleContext.java

@@ -52,6 +52,7 @@
 import com.google.devtools.build.lib.analysis.config.FragmentCollection;
 import com.google.devtools.build.lib.analysis.platform.ConstraintValueInfo;
 import com.google.devtools.build.lib.analysis.starlark.StarlarkActionFactory.StarlarkActionContext;
+import com.google.devtools.build.lib.analysis.starlark.StarlarkSubrule.SubruleContext;
 import com.google.devtools.build.lib.analysis.stringtemplate.ExpansionException;
 import com.google.devtools.build.lib.analysis.test.InstrumentedFilesCollector;
 import com.google.devtools.build.lib.analysis.test.InstrumentedFilesInfo;
@@ -192,8 +193,9 @@ public final class StarlarkRuleContext
    */
   private int resolveCommandScriptCounter = 0;
 
-  // for temporarily freezing mutability, such as while evaluating a subrule
-  private boolean lockedForSubruleEvaluation = false;
+  // for temporarily freezing mutability, while evaluating a subrule this is set to the
+  // corresponding subrule context, or is null otherwise
+  @Nullable private SubruleContext lockedForSubruleEvaluation = null;
 
   /**
    * Creates a new StarlarkRuleContext wrapping ruleContext.
@@ -332,11 +334,11 @@ ImmutableSet<? extends StarlarkSubruleApi> getSubrules() {
     }
   }
 
-  void setLockedForSubrule(boolean isLocked) {
-    this.lockedForSubruleEvaluation = isLocked;
+  void setLockedForSubrule(@Nullable SubruleContext lockedBy) {
+    this.lockedForSubruleEvaluation = lockedBy;
   }
 
-  boolean getLockedForSubrule() {
+  SubruleContext getLockedForSubrule() {
     return lockedForSubruleEvaluation;
   }
 
@@ -564,7 +566,7 @@ private static StructImpl buildSplitAttributeInfo(
 
   @Override
   public boolean isImmutable() {
-    return ruleContext == null || lockedForSubruleEvaluation;
+    return ruleContext == null || lockedForSubruleEvaluation != null;
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java

@@ -75,6 +75,7 @@ public class StarlarkSubrule implements StarlarkExportable, StarlarkCallable, St
   private final StarlarkFunction implementation;
   private final ImmutableSet<ToolchainTypeRequirement> toolchains;
   private final ImmutableSet<String> fragments;
+  private final ImmutableSet<StarlarkSubrule> subrules;
 
   // following fields are set on export
   @Nullable private String exportedName = null;
@@ -84,11 +85,13 @@ public StarlarkSubrule(
       StarlarkFunction implementation,
       ImmutableMap<String, Descriptor> attributes,
       ImmutableSet<ToolchainTypeRequirement> toolchains,
-      ImmutableSet<String> fragments) {
+      ImmutableSet<String> fragments,
+      ImmutableSet<StarlarkSubrule> subrules) {
     this.implementation = implementation;
     this.attributes = SubruleAttribute.from(attributes);
     this.toolchains = toolchains;
     this.fragments = fragments;
+    this.subrules = subrules;
   }
 
   @Override
@@ -113,11 +116,14 @@ public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwarg
         BazelRuleAnalysisThreadContext.fromOrFail(thread, getName())
             .getRuleContext()
             .getStarlarkRuleContext();
-    if (ruleContext.getLockedForSubrule()) {
-      throw Starlark.errorf("subrules cannot call other subrules");
-    }
-    ImmutableSet<? extends StarlarkSubruleApi> declaredSubrules = ruleContext.getSubrules();
-    if (!declaredSubrules.contains(this)) {
+    SubruleContext callerSubruleContext = ruleContext.getLockedForSubrule();
+    if (callerSubruleContext != null) {
+      if (!callerSubruleContext.subrule.getDeclaredSubrules().contains(this)) {
+        throw Starlark.errorf(
+            "subrule %s must declare %s in 'subrules'",
+            callerSubruleContext.subrule.getName(), getName());
+      }
+    } else if (!ruleContext.getSubrules().contains(this)) {
       throw getUndeclaredSubruleError(ruleContext);
     }
     ImmutableSet.Builder<FilesToRunProvider> runfilesFromDeps = ImmutableSet.builder();
@@ -164,15 +170,21 @@ public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwarg
     ImmutableList<Object> positionals =
         ImmutableList.builder().add(subruleContext).addAll(args).build();
     try {
-      ruleContext.setLockedForSubrule(true);
+      ruleContext.setLockedForSubrule(subruleContext);
       return Starlark.call(
           thread, implementation, positionals, Dict.immutableCopyOf(namedArgs.buildOrThrow()));
     } finally {
       subruleContext.nullify();
-      ruleContext.setLockedForSubrule(false);
+      // callerSubruleContext may be null if this subrule was called from the rule itself, but in
+      // that case null is exactly what we want to set here
+      ruleContext.setLockedForSubrule(callerSubruleContext);
     }
   }
 
+  private ImmutableSet<StarlarkSubrule> getDeclaredSubrules() {
+    return subrules;
+  }
+
   private EvalException getUndeclaredSubruleError(StarlarkRuleContext starlarkRuleContext) {
     if (starlarkRuleContext.isForAspect()) {
       return Starlark.errorf(
@@ -229,37 +241,35 @@ public void export(EventHandler handler, Label extensionLabel, String exportedNa
    */
   static ImmutableList<Pair<String, Descriptor>> discoverAttributes(
       ImmutableList<? extends StarlarkSubruleApi> subrules) throws EvalException {
-    ImmutableSet.Builder<StarlarkSubruleApi> uniqueSubrules = ImmutableSet.builder();
-    for (StarlarkSubruleApi subrule : subrules) {
-      // TODO: b/293304174 - use all transitive subrules once subrules can depend on other subrules
-      uniqueSubrules.add(subrule);
-    }
     ImmutableList.Builder<Pair<String, Descriptor>> attributes = ImmutableList.builder();
-    for (StarlarkSubruleApi subrule : uniqueSubrules.build()) {
-      if (subrule instanceof StarlarkSubrule) {
-        attributes.addAll(((StarlarkSubrule) subrule).attributesForRule());
-      }
+    for (StarlarkSubrule subrule : getTransitiveSubrules(subrules)) {
+      attributes.addAll(subrule.attributesForRule());
     }
     return attributes.build();
   }
 
   /** Returns all toolchain types to be lifted from the given subrules to a rule/aspect */
   static ImmutableSet<ToolchainTypeRequirement> discoverToolchains(
       ImmutableList<? extends StarlarkSubruleApi> subrules) {
-    ImmutableSet.Builder<StarlarkSubruleApi> uniqueSubrules = ImmutableSet.builder();
-    for (StarlarkSubruleApi subrule : subrules) {
-      // TODO: b/293304174 - use all transitive subrules once subrules can depend on other subrules
-      uniqueSubrules.add(subrule);
-    }
     ImmutableSet.Builder<ToolchainTypeRequirement> toolchains = ImmutableSet.builder();
-    for (StarlarkSubruleApi subrule : uniqueSubrules.build()) {
-      if (subrule instanceof StarlarkSubrule) {
-        toolchains.addAll(((StarlarkSubrule) subrule).toolchains);
-      }
+    for (StarlarkSubrule subrule : getTransitiveSubrules(subrules)) {
+      toolchains.addAll(subrule.toolchains);
     }
     return toolchains.build();
   }
 
+  private static ImmutableSet<StarlarkSubrule> getTransitiveSubrules(
+      ImmutableCollection<? extends StarlarkSubruleApi> subrules) {
+    ImmutableSet.Builder<StarlarkSubrule> uniqueSubrules = ImmutableSet.builder();
+    for (StarlarkSubruleApi subruleApi : subrules) {
+      if (subruleApi instanceof StarlarkSubrule) {
+        StarlarkSubrule subrule = (StarlarkSubrule) subruleApi;
+        uniqueSubrules.add(subrule).addAll(getTransitiveSubrules(subrule.getDeclaredSubrules()));
+      }
+    }
+    return uniqueSubrules.build();
+  }
+
   /**
    * The context object passed to the implementation function of a subrule.
    *
@@ -272,7 +282,7 @@ static ImmutableSet<ToolchainTypeRequirement> discoverToolchains(
       name = "subrule_ctx",
       category = DocCategory.BUILTIN,
       doc = "A context object passed to the implementation function of a subrule.")
-  private static class SubruleContext implements StarlarkActionContext {
+  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;
@@ -362,7 +372,6 @@ public ArtifactRoot newFileRoot() {
 
     @Override
     public void checkMutable(String attrName) throws EvalException {
-      // TODO: b/293304174 - check if subrule is locked once subrules can call other subrules
       if (isImmutable()) {
         throw Starlark.errorf(
             "cannot access field or method '%s' of subrule context outside of its own"
@@ -373,7 +382,7 @@ public void checkMutable(String attrName) throws EvalException {
 
     @Override
     public boolean isImmutable() {
-      return starlarkRuleContext == null;
+      return starlarkRuleContext == null || starlarkRuleContext.getLockedForSubrule() != this;
     }
 
     @Override
@@ -425,6 +434,7 @@ public Object maybeOverrideToolchain(Object toolchainUnchecked) throws EvalExcep
           : Iterables.getOnlyElement(requestedToolchains);
     }
 
+    // TODO: b/293304174 - maybe simplify all this by just relying on starlarkRuleContext
     private void nullify() {
       this.subrule = null;
       this.starlarkRuleContext = null;

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/StarlarkRuleFunctionsApi.java

@@ -849,14 +849,22 @@ ExecGroupApi execGroup(
             defaultValue = "[]",
             doc =
                 "List of names of configuration fragments that the subrule requires in target"
-                    + " configuration.")
+                    + " configuration."),
+        @Param(
+            name = "subrules",
+            allowedTypes = {@ParamType(type = Sequence.class, generic1 = StarlarkSubruleApi.class)},
+            named = true,
+            positional = false,
+            defaultValue = "[]",
+            doc = "List of other subrules needed by this subrule.")
       },
       useStarlarkThread = true)
   StarlarkSubruleApi subrule(
       StarlarkFunction implementation,
       Dict<?, ?> attrs,
       Sequence<?> toolchains,
       Sequence<?> fragments,
+      Sequence<?> subrules,
       StarlarkThread thread)
       throws EvalException;
 }

src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeStarlarkRuleFunctionsApi.java

@@ -252,6 +252,7 @@ public StarlarkSubruleApi subrule(
       Dict<?, ?> attrs,
       Sequence<?> toolchains,
       Sequence<?> fragments,
+      Sequence<?> subrules,
       StarlarkThread thread) {
     return new FakeStarlarkSubrule();
   }