From aa55c4df164e9a7be0ee592a10b35f505256b213 Mon Sep 17 00:00:00 2001
From: Googler <hvd@google.com>
Date: Wed, 30 Aug 2023 00:48:57 -0700
Subject: [PATCH] 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
---
 .../BazelRuleAnalysisThreadContext.java       |  4 +
 .../starlark/StarlarkRuleClassFunctions.java  | 22 +++++-
 .../starlark/StarlarkRuleContext.java         | 10 +++
 .../analysis/starlark/StarlarkSubrule.java    | 23 +++++-
 .../build/lib/packages/AspectDefinition.java  | 21 ++++-
 .../build/lib/packages/RuleClass.java         | 28 ++++++-
 .../lib/packages/StarlarkDefinedAspect.java   |  7 +-
 .../lib/rules/test/StarlarkTestingModule.java |  3 +-
 .../StarlarkRuleFunctionsApi.java             | 22 +++++-
 .../FakeStarlarkRuleFunctionsApi.java         |  2 +
 .../starlark/StarlarkSubruleTest.java         | 76 ++++++++++++++++++-
 .../build/lib/packages/RuleClassTest.java     |  3 +-
 12 files changed, 204 insertions(+), 17 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/BazelRuleAnalysisThreadContext.java b/src/main/java/com/google/devtools/build/lib/analysis/BazelRuleAnalysisThreadContext.java
index 4e0bcf9cdc0fb0..abf90641c9e00c 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/BazelRuleAnalysisThreadContext.java
+++ b/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.
    *
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 f43430a95356bb..2931ff22c5aa49 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
@@ -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")));
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleContext.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleContext.java
index 50f93b34fdf303..cb02362cce6417 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleContext.java
+++ b/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`.
    *
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 58f056c84b179e..3055e9f15858d2 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
@@ -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 {}
 }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java b/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
index 47a91190212ab4..516db17ddfd156 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/AspectDefinition.java
+++ b/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);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java b/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
index bf775ff13d7cbd..ff2aaf6349b40e 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
+++ b/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;
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java b/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java
index efc498880957fa..b6168dcc9c8cce 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/StarlarkDefinedAspect.java
+++ b/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();
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java b/src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java
index 56ef2a10239b57..d8b7980b25a263 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java
+++ b/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.
     //
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 acbb6938aee855..eca20e08f4fa5f 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
@@ -441,6 +441,15 @@ Object provider(Object doc, Object fields, Object init, StarlarkThread thread)
                     + " allows rules to run actions on multiple execution platforms within a"
                     + " single target. See <a href='${link exec-groups}'>execution groups"
                     + " documentation</a> for more info."),
+        @Param(
+            name = "subrules",
+            allowedTypes = {
+              @ParamType(type = Sequence.class, generic1 = StarlarkSubruleApi.class),
+            },
+            named = true,
+            defaultValue = "[]",
+            positional = false,
+            doc = "Experimental, DO NOT USE!"),
       },
       useStarlarkThread = true)
   StarlarkCallable rule(
@@ -462,6 +471,7 @@ StarlarkCallable rule(
       Object buildSetting,
       Object cfg,
       Object execGroups,
+      Sequence<?> subrules,
       StarlarkThread thread)
       throws EvalException;
 
@@ -648,7 +658,16 @@ StarlarkCallable rule(
                     + " href='../globals/bzl.html#exec_group'><code>exec_group</code>s</a>. If set,"
                     + " allows aspects to run actions on multiple execution platforms within a"
                     + " single instance. See <a href='${link exec-groups}'>execution groups"
-                    + " documentation</a> for more info.")
+                    + " documentation</a> for more info."),
+        @Param(
+            name = "subrules",
+            allowedTypes = {
+              @ParamType(type = Sequence.class, generic1 = StarlarkSubruleApi.class),
+            },
+            named = true,
+            defaultValue = "[]",
+            positional = false,
+            doc = "Experimental, DO NOT USE!")
       },
       useStarlarkThread = true)
   StarlarkAspectApi aspect(
@@ -667,6 +686,7 @@ StarlarkAspectApi aspect(
       Boolean applyToGeneratingRules,
       Sequence<?> execCompatibleWith,
       Object execGroups,
+      Sequence<?> subrules,
       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 f21385e9c51f3b..22e47d3bed82b6 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
@@ -148,6 +148,7 @@ public StarlarkCallable rule(
       Object buildSetting,
       Object cfg,
       Object execGroups,
+      Sequence<?> subrules,
       StarlarkThread thread)
       throws EvalException {
     ImmutableMap.Builder<String, FakeDescriptor> attrsMapBuilder = ImmutableMap.builder();
@@ -204,6 +205,7 @@ public StarlarkAspectApi aspect(
       Boolean applyToFiles,
       Sequence<?> execCompatibleWith,
       Object execGroups,
+      Sequence<?> subrules,
       StarlarkThread thread)
       throws EvalException {
     FakeStarlarkAspect fakeAspect = new FakeStarlarkAspect();
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 911aeee1561905..cf7b239bd69318 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
@@ -72,7 +72,7 @@ public void testSubrule_isCallableOnlyFromRuleOrAspectImplementation() throws Ex
   }
 
   @Test
-  public void testSubrule_implementationMustAcceptSubruleContext() throws Exception {
+  public void testSubrule_ruleMustDeclareSubrule() throws Exception {
     scratch.file(
         "subrule_testing/myrule.bzl",
         "_my_subrule = subrule(implementation = lambda : '')",
@@ -90,6 +90,64 @@ public void testSubrule_implementationMustAcceptSubruleContext() throws Exceptio
     AssertionError error =
         assertThrows(AssertionError.class, () -> getConfiguredTarget("//subrule_testing:foo"));
 
+    assertThat(error)
+        .hasMessageThat()
+        .contains(
+            "Error in subrule(lambda): rule 'my_rule' must declare 'subrule(lambda)' in"
+                + " 'subrules'");
+  }
+
+  @Test
+  public void testSubrule_aspectMustDeclareSubrule() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "_my_subrule = subrule(implementation = lambda ctx: 'dummy aspect result')",
+        "",
+        "def _aspect_impl(ctx,target):",
+        "  res = _my_subrule()",
+        "",
+        "_my_aspect = aspect(implementation = _aspect_impl)",
+        "",
+        "my_rule = rule(",
+        "  implementation = lambda ctx: [],",
+        "  attrs = {'dep' : attr.label(mandatory = True, aspects = [_my_aspect])},",
+        ")");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "java_library(name = 'bar')",
+        "my_rule(name = 'foo', dep = 'bar')");
+
+    AssertionError error =
+        assertThrows(AssertionError.class, () -> getConfiguredTarget("//subrule_testing:foo"));
+
+    assertThat(error)
+        .hasMessageThat()
+        .contains(
+            "Error in subrule(lambda): aspect '//subrule_testing:myrule.bzl%_my_aspect' must"
+                + " declare 'subrule(lambda)' in 'subrules'");
+  }
+
+  @Test
+  public void testSubrule_implementationMustAcceptSubruleContext() throws Exception {
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        "_my_subrule = subrule(implementation = lambda : '')",
+        "",
+        "def _rule_impl(ctx):",
+        "  _my_subrule()",
+        "",
+        "my_rule = rule(implementation = _rule_impl, subrules = [_my_subrule])");
+    scratch.file(
+        "subrule_testing/BUILD",
+        //
+        "load('myrule.bzl', 'my_rule')",
+        "my_rule(name = 'foo')");
+
+    AssertionError error =
+        assertThrows(AssertionError.class, () -> getConfiguredTarget("//subrule_testing:foo"));
+
     assertThat(error)
         .hasMessageThat()
         .contains("Error: lambda() does not accept positional arguments, but got 1");
@@ -106,7 +164,7 @@ public void testSubrule_isCallableFromRule() throws Exception {
         "  res = _my_subrule()",
         "  return MyInfo(result = res)",
         "",
-        "my_rule = rule(implementation = _rule_impl)");
+        "my_rule = rule(implementation = _rule_impl, subrules = [_my_subrule])");
     scratch.file(
         "subrule_testing/BUILD",
         //
@@ -131,7 +189,7 @@ public void testSubrule_isCallableFromAspect() throws Exception {
         "  res = _my_subrule()",
         "  return MyInfo(result = res)",
         "",
-        "_my_aspect = aspect(implementation = _aspect_impl)",
+        "_my_aspect = aspect(implementation = _aspect_impl, subrules = [_my_subrule])",
         "",
         "my_rule = rule(",
         "  implementation = lambda ctx: [ctx.attr.dep[MyInfo]],",
@@ -158,6 +216,18 @@ public void testSubruleInstantiation_outsideAllowlist_failsWithPrivateAPIError()
         "'//foo:bar' cannot use private API", "subrule(implementation = lambda: 0 )");
   }
 
+  @Test
+  public void testSubrulesParamForRule_isPrivateAPI() throws Exception {
+    evOutsideAllowlist.checkEvalErrorContains(
+        "'//foo:bar' cannot use private API", "rule(implementation = lambda: 0, subrules = [1])");
+  }
+
+  @Test
+  public void testSubrulesParamForAspect_isPrivateAPI() throws Exception {
+    evOutsideAllowlist.checkEvalErrorContains(
+        "'//foo:bar' cannot use private API", "aspect(implementation = lambda: 0, subrules = [1])");
+  }
+
   private StructImpl getProvider(String targetLabel, String providerLabel, String providerName)
       throws LabelSyntaxException {
     ConfiguredTarget target = getConfiguredTarget(targetLabel);
diff --git a/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java b/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java
index 666701fa30bac9..be0948ef4761a9 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/RuleClassTest.java
@@ -1058,7 +1058,8 @@ private static RuleClass newRuleClass(
                 .add(RuleClass.NAME_ATTRIBUTE)
                 .add(attributes)
                 .build(),
-        /* buildSetting= */ null);
+        /* buildSetting= */ null,
+        /* subrules= */ ImmutableList.of());
   }
 
   private static RuleClass createParentRuleClass() {