Add private subrules param to rule() and aspect()

Rules or aspects that invoke a subrule must declare them in the `subrules` parameter

PiperOrigin-RevId: 561258778
Change-Id: Id228c26ba09bf4f0677116ac5b804b927e5c087a

src/main/java/com/google/devtools/build/lib/analysis/BazelRuleAnalysisThreadContext.java

@@ -52,6 +52,10 @@ public String getContextForUncheckedException() {
     return ruleContext.getLabel().toString();
   }
 
+  public RuleContext getRuleContext() {
+    return ruleContext;
+  }
+
   /**
    * Retrieves this context from a Starlark thread.
    *

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

@@ -303,6 +303,7 @@ public StarlarkRuleFunction rule(
       Object buildSetting,
       Object cfg,
       Object execGroups,
+      Sequence<?> subrules,
       StarlarkThread thread)
       throws EvalException {
     // Ensure we're initializing a .bzl file, which also means we have a RuleDefinitionEnvironment.
@@ -314,6 +315,10 @@ public StarlarkRuleFunction rule(
 
     LabelConverter labelConverter = LabelConverter.forBzlEvaluatingThread(thread);
 
+    if (!subrules.isEmpty()) {
+      BuiltinRestriction.failIfCalledOutsideAllowlist(thread, ALLOWLIST_SUBRULES);
+    }
+
     return createRule(
         // Contextual parameters.
         bazelContext,
@@ -339,7 +344,8 @@ public StarlarkRuleFunction rule(
         analysisTest,
         buildSetting,
         cfg,
-        execGroups);
+        execGroups,
+        subrules);
   }
 
   /**
@@ -377,7 +383,8 @@ public static StarlarkRuleFunction createRule(
       Object analysisTest,
       Object buildSetting,
       Object cfg,
-      Object execGroups)
+      Object execGroups,
+      Sequence<?> subrules)
       throws EvalException {
 
     // analysis_test=true implies test=true.
@@ -523,6 +530,9 @@ public static StarlarkRuleFunction createRule(
           parseExecCompatibleWith(execCompatibleWith, labelConverter));
     }
 
+    builder.setSubrules(
+        Sequence.cast(subrules, StarlarkSubruleApi.class, "subrules").getImmutableList());
+
     return new StarlarkRuleFunction(
         builder,
         type,
@@ -586,6 +596,7 @@ public StarlarkAspect aspect(
       Boolean applyToGeneratingRules,
       Sequence<?> rawExecCompatibleWith,
       Object rawExecGroups,
+      Sequence<?> subrules,
       StarlarkThread thread)
       throws EvalException {
     LabelConverter labelConverter = LabelConverter.forBzlEvaluatingThread(thread);
@@ -705,6 +716,10 @@ public StarlarkAspect aspect(
       }
     }
 
+    if (!subrules.isEmpty()) {
+      BuiltinRestriction.failIfCalledOutsideAllowlist(thread, ALLOWLIST_SUBRULES);
+    }
+
     return new StarlarkDefinedAspect(
         implementation,
         Starlark.toJavaOptional(doc, String.class).map(Starlark::trimDocString),
@@ -720,7 +735,8 @@ public StarlarkAspect aspect(
         parseToolchainTypes(toolchains, labelConverter),
         applyToGeneratingRules,
         execCompatibleWith,
-        execGroups);
+        execGroups,
+        ImmutableSet.copyOf(Sequence.cast(subrules, StarlarkSubruleApi.class, "subrules")));
   }
 
   /**

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

@@ -76,6 +76,7 @@
 import com.google.devtools.build.lib.shell.ShellUtils;
 import com.google.devtools.build.lib.shell.ShellUtils.TokenizationException;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkRuleContextApi;
+import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import com.google.devtools.build.lib.starlarkbuildapi.platform.ToolchainContextApi;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.ArrayList;
@@ -298,6 +299,15 @@ public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor a
     }
   }
 
+  /** Returns the subrules declared by the rule or aspect represented by this context. */
+  ImmutableSet<? extends StarlarkSubruleApi> getSubrules() {
+    if (isForAspect()) {
+      return getRuleContext().getMainAspect().getDefinition().getSubrules();
+    } else {
+      return getRuleContext().getRule().getRuleClassObject().getSubrules();
+    }
+  }
+
   /**
    * Represents `ctx.outputs`.
    *

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

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.analysis.starlark;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.analysis.BazelRuleAnalysisThreadContext;
 import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import net.starlark.java.eval.Dict;
@@ -43,12 +44,32 @@ public String getName() {
   @Override
   public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwargs)
       throws EvalException, InterruptedException {
-    BazelRuleAnalysisThreadContext.fromOrFail(thread, getName());
+    StarlarkRuleContext ruleContext =
+        BazelRuleAnalysisThreadContext.fromOrFail(thread, getName())
+            .getRuleContext()
+            .getStarlarkRuleContext();
+    ImmutableSet<? extends StarlarkSubruleApi> declaredSubrules = ruleContext.getSubrules();
+    if (!declaredSubrules.contains(this)) {
+      throw getUndeclaredSubruleError(ruleContext);
+    }
     SubruleContext subruleContext = new SubruleContext();
     ImmutableList<Object> positionals =
         ImmutableList.builder().add(subruleContext).addAll(args).build();
     return Starlark.call(thread, implementation, positionals, kwargs);
   }
 
+  private EvalException getUndeclaredSubruleError(StarlarkRuleContext starlarkRuleContext) {
+    if (starlarkRuleContext.isForAspect()) {
+      return Starlark.errorf(
+          "aspect '%s' must declare '%s' in 'subrules'",
+          starlarkRuleContext.getRuleContext().getMainAspect().getAspectClass().getName(),
+          this.getName());
+    } else {
+      return Starlark.errorf(
+          "rule '%s' must declare '%s' in 'subrules'",
+          starlarkRuleContext.getRuleContext().getRule().getRuleClass(), this.getName());
+    }
+  }
+
   private static class SubruleContext {}
 }

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

@@ -28,6 +28,7 @@
 import com.google.devtools.build.lib.packages.ConfigurationFragmentPolicy.MissingFragmentPolicy;
 import com.google.devtools.build.lib.packages.Type.LabelClass;
 import com.google.devtools.build.lib.packages.Type.LabelVisitor;
+import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Collection;
 import java.util.HashSet;
@@ -83,6 +84,7 @@ public final class AspectDefinition {
 
   private final ImmutableSet<Label> execCompatibleWith;
   private final ImmutableMap<String, ExecGroup> execGroups;
+  private final ImmutableSet<? extends StarlarkSubruleApi> subrules;
 
   public AdvertisedProviderSet getAdvertisedProviders() {
     return advertisedProviders;
@@ -101,7 +103,8 @@ private AspectDefinition(
       boolean applyToGeneratingRules,
       ImmutableSet<AspectClass> requiredAspectClasses,
       ImmutableSet<Label> execCompatibleWith,
-      ImmutableMap<String, ExecGroup> execGroups) {
+      ImmutableMap<String, ExecGroup> execGroups,
+      ImmutableSet<? extends StarlarkSubruleApi> subrules) {
     this.aspectClass = aspectClass;
     this.advertisedProviders = advertisedProviders;
     this.requiredProviders = requiredProviders;
@@ -115,6 +118,7 @@ private AspectDefinition(
     this.requiredAspectClasses = requiredAspectClasses;
     this.execCompatibleWith = execCompatibleWith;
     this.execGroups = execGroups;
+    this.subrules = subrules;
   }
 
   public String getName() {
@@ -147,6 +151,11 @@ public ImmutableMap<String, ExecGroup> execGroups() {
     return execGroups;
   }
 
+  /** Returns the subrules declared by this aspect. */
+  public ImmutableSet<? extends StarlarkSubruleApi> getSubrules() {
+    return subrules;
+  }
+
   /**
    * Returns {@link RequiredProviders} that a configured target must have so that this aspect can be
    * applied to it.
@@ -262,6 +271,7 @@ public static final class Builder {
     private ImmutableSet<AspectClass> requiredAspectClasses = ImmutableSet.of();
     private ImmutableSet<Label> execCompatibleWith = ImmutableSet.of();
     private ImmutableMap<String, ExecGroup> execGroups = ImmutableMap.of();
+    private ImmutableSet<? extends StarlarkSubruleApi> subrules = ImmutableSet.of();
 
     public Builder(AspectClass aspectClass) {
       this.aspectClass = aspectClass;
@@ -561,6 +571,12 @@ public Builder execGroups(ImmutableMap<String, ExecGroup> execGroups) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder subrules(ImmutableSet<? extends StarlarkSubruleApi> subrules) {
+      this.subrules = subrules;
+      return this;
+    }
+
     /**
      * Builds the aspect definition.
      *
@@ -587,7 +603,8 @@ public AspectDefinition build() {
           applyToGeneratingRules,
           requiredAspectClasses,
           execCompatibleWith,
-          execGroups);
+          execGroups,
+          subrules);
     }
   }
 }

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

@@ -54,6 +54,7 @@
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
+import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import com.google.devtools.build.lib.util.HashCodes;
 import com.google.devtools.build.lib.util.StringUtil;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -769,6 +770,8 @@ public String toString() {
         AdvertisedProviderSet.builder();
     private StarlarkCallable configuredTargetFunction = null;
     private BuildSetting buildSetting = null;
+
+    private ImmutableList<? extends StarlarkSubruleApi> subrules = ImmutableList.of();
     private Function<? super Rule, Map<String, Label>> externalBindingsFunction =
         NO_EXTERNAL_BINDINGS;
     private Function<? super Rule, ? extends List<String>> toolchainsToRegisterFunction =
@@ -946,7 +949,8 @@ public RuleClass build(String name, String key) {
           execGroups,
           outputFileKind,
           ImmutableList.copyOf(attributes.values()),
-          buildSetting);
+          buildSetting,
+          subrules);
     }
 
     private static void checkAttributes(
@@ -1269,6 +1273,12 @@ public Builder setBuildSetting(BuildSetting buildSetting) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder setSubrules(ImmutableList<? extends StarlarkSubruleApi> subrules) {
+      this.subrules = subrules;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder setExternalBindingsFunction(Function<? super Rule, Map<String, Label>> func) {
       this.externalBindingsFunction = func;
@@ -1622,8 +1632,12 @@ public Attribute.Builder<?> copy(String name) {
   @Nullable private final BuildSetting buildSetting;
 
   /**
-   * Returns the extra bindings a workspace function adds to the WORKSPACE file.
+   * The subrules associated with this rule. Empty for all rule classes except Starlark-defined
+   * rules that explicitly pass {@code subrules = [...]} to their {@code rule()} declaration
    */
+  private final ImmutableSet<? extends StarlarkSubruleApi> subrules;
+
+  /** Returns the extra bindings a workspace function adds to the WORKSPACE file. */
   private final Function<? super Rule, Map<String, Label>> externalBindingsFunction;
 
   /** Returns the toolchains a workspace function wants to have registered in the WORKSPACE file. */
@@ -1715,7 +1729,8 @@ public Attribute.Builder<?> copy(String name) {
       Map<String, ExecGroup> execGroups,
       OutputFile.Kind outputFileKind,
       ImmutableList<Attribute> attributes,
-      @Nullable BuildSetting buildSetting) {
+      @Nullable BuildSetting buildSetting,
+      ImmutableList<? extends StarlarkSubruleApi> subrules) {
     this.name = name;
     this.callstack = callstack;
     this.key = key;
@@ -1751,7 +1766,7 @@ public Attribute.Builder<?> copy(String name) {
     this.executionPlatformConstraints = ImmutableSet.copyOf(executionPlatformConstraints);
     this.execGroups = ImmutableMap.copyOf(execGroups);
     this.buildSetting = buildSetting;
-
+    this.subrules = ImmutableSet.copyOf(subrules);
     // Create the index and collect non-configurable attributes while doing some validation checks.
     Preconditions.checkState(
         !attributes.isEmpty() && attributes.get(0).equals(NAME_ATTRIBUTE),
@@ -2649,4 +2664,9 @@ public boolean isPackageMetadataRule() {
     // END-INTERNAL
     return false;
   }
+
+  public ImmutableSet<? extends StarlarkSubruleApi> getSubrules() {
+    Preconditions.checkState(isStarlark());
+    return subrules;
+  }
 }

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

@@ -25,6 +25,7 @@
 import com.google.devtools.build.lib.analysis.config.ToolchainTypeRequirement;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.starlarkbuildapi.StarlarkSubruleApi;
 import java.io.Serializable;
 import java.util.Objects;
 import java.util.Optional;
@@ -55,6 +56,7 @@ public final class StarlarkDefinedAspect implements StarlarkExportable, Starlark
   private final boolean applyToGeneratingRules;
   private final ImmutableSet<Label> execCompatibleWith;
   private final ImmutableMap<String, ExecGroup> execGroups;
+  private final ImmutableSet<? extends StarlarkSubruleApi> subrules;
 
   private StarlarkAspectClass aspectClass;
 
@@ -78,7 +80,8 @@ public StarlarkDefinedAspect(
       ImmutableSet<ToolchainTypeRequirement> toolchainTypes,
       boolean applyToGeneratingRules,
       ImmutableSet<Label> execCompatibleWith,
-      ImmutableMap<String, ExecGroup> execGroups) {
+      ImmutableMap<String, ExecGroup> execGroups,
+      ImmutableSet<? extends StarlarkSubruleApi> subrules) {
     this.implementation = implementation;
     this.documentation = documentation.orElse(null);
     this.attributeAspects = attributeAspects;
@@ -93,6 +96,7 @@ public StarlarkDefinedAspect(
     this.applyToGeneratingRules = applyToGeneratingRules;
     this.execCompatibleWith = execCompatibleWith;
     this.execGroups = execGroups;
+    this.subrules = subrules;
   }
 
   public StarlarkCallable getImplementation() {
@@ -216,6 +220,7 @@ private AspectDefinition buildDefinition(AspectParameters aspectParams) {
     builder.requiredAspectClasses(requiredAspectsClasses.build());
     builder.execCompatibleWith(execCompatibleWith);
     builder.execGroups(execGroups);
+    builder.subrules(subrules);
     return builder.build();
   }
 

src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java

@@ -124,7 +124,8 @@ public void analysisTest(
             /* analysisTest= */ Boolean.TRUE,
             /* buildSetting= */ Starlark.NONE,
             /* cfg= */ Starlark.NONE,
-            /* execGroups= */ Starlark.NONE);
+            /* execGroups= */ Starlark.NONE,
+            /* subrules= */ StarlarkList.empty());
 
     // Export the rule.
     //