Allow multiple trimming transition factories to be added.

If multiple trimming transition factories are added, they are composed in
the order they were added.

RELNOTES: None.
PiperOrigin-RevId: 198934666

3 files changed, 64 insertions, 10 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java b/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java
index c333d0563b..4821f6f961 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ConfiguredRuleClassProvider.java
@@ -32,6 +32,7 @@ import com.google.devtools.build.lib.analysis.buildinfo.BuildInfoFactory;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration.Fragment;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
+import com.google.devtools.build.lib.analysis.config.ComposingRuleTransitionFactory;
 import com.google.devtools.build.lib.analysis.config.ConfigurationFragmentFactory;
 import com.google.devtools.build.lib.analysis.config.DefaultsPackage;
 import com.google.devtools.build.lib.analysis.config.FragmentOptions;
@@ -413,17 +414,22 @@ public class ConfiguredRuleClassProvider implements RuleClassProvider {
     }
 
     /**
-     * Sets the transition factory that produces a trimming transition to be run over all targets
+     * Adds a transition factory that produces a trimming transition to be run over all targets
      * after other transitions.
      *
-     * <p>This is a temporary measure for supporting manual trimming of feature flags, and support
-     * for this transition factory will likely be removed at some point in the future (whenever
-     * automatic trimming is sufficiently workable).
+     * <p>Transitions are run in the order they're added.
+     *
+     * <p>This is a temporary measure for supporting trimming of test rules and manual trimming of
+     * feature flags, and support for this transition factory will likely be removed at some point
+     * in the future (whenever automatic trimming is sufficiently workable).
      */
-    public Builder setTrimmingTransitionFactory(RuleTransitionFactory factory) {
-      Preconditions.checkState(
-          trimmingTransitionFactory == null, "Trimming transition factory already set");
-      trimmingTransitionFactory = Preconditions.checkNotNull(factory);
+    public Builder addTrimmingTransitionFactory(RuleTransitionFactory factory) {
+      if (trimmingTransitionFactory == null) {
+        trimmingTransitionFactory = Preconditions.checkNotNull(factory);
+      } else {
+        trimmingTransitionFactory = new ComposingRuleTransitionFactory(
+            trimmingTransitionFactory, Preconditions.checkNotNull(factory));
+      }
       return this;
     }
 
@@ -435,7 +441,7 @@ public class ConfiguredRuleClassProvider implements RuleClassProvider {
     @VisibleForTesting(/* for testing trimming transition factories without relying on prod use */)
     public Builder overrideTrimmingTransitionFactoryForTesting(RuleTransitionFactory factory) {
       trimmingTransitionFactory = null;
-      return this.setTrimmingTransitionFactory(factory);
+      return this.addTrimmingTransitionFactory(factory);
     }
 
     @Override
diff --git a/src/main/java/com/google/devtools/build/lib/rules/config/ConfigRules.java b/src/main/java/com/google/devtools/build/lib/rules/config/ConfigRules.java
index 6a5fdaec6a..7531493e77 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/config/ConfigRules.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/config/ConfigRules.java
@@ -32,7 +32,7 @@ public final class ConfigRules implements RuleSet {
 
   @Override
   public void init(ConfiguredRuleClassProvider.Builder builder) {
-    builder.setTrimmingTransitionFactory(
+    builder.addTrimmingTransitionFactory(
         new ConfigFeatureFlagTaggedTrimmingTransitionFactory(BaseRuleClasses.TAGGED_TRIMMING_ATTR));
     builder.addRuleDefinition(new ConfigRuleClasses.ConfigBaseRule());
     builder.addRuleDefinition(new ConfigRuleClasses.ConfigSettingRule());
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ConfigurationsForTargetsWithTrimmedConfigurationsTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ConfigurationsForTargetsWithTrimmedConfigurationsTest.java
index 3f5f37c02b..cc438ce9b4 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ConfigurationsForTargetsWithTrimmedConfigurationsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ConfigurationsForTargetsWithTrimmedConfigurationsTest.java
@@ -310,6 +310,54 @@ public class ConfigurationsForTargetsWithTrimmedConfigurationsTest
   }
 
   @Test
+  public void trimmingTransitionsAreComposedInOrderOfAdding() throws Exception {
+    RuleTransitionFactory firstTrimmingTransitionFactory =
+        (rule) ->
+            new AddArgumentToTestArgsTransition(
+                "first trimming transition for " + rule.getLabel().toString());
+    RuleTransitionFactory secondTrimmingTransitionFactory =
+        (rule) ->
+            new AddArgumentToTestArgsTransition(
+                "second trimming transition for " + rule.getLabel().toString());
+    ConfiguredRuleClassProvider.Builder builder = new ConfiguredRuleClassProvider.Builder();
+    TestRuleClassProvider.addStandardRules(builder);
+    builder.addRuleDefinition(TestAspects.BASE_RULE);
+    builder.addRuleDefinition(TEST_BASE_RULE);
+    builder.overrideTrimmingTransitionFactoryForTesting(firstTrimmingTransitionFactory);
+    builder.addTrimmingTransitionFactory(secondTrimmingTransitionFactory);
+    useRuleClassProvider(builder.build());
+    scratch.file(
+        "a/skylark.bzl",
+        "def _impl(ctx):",
+        "  return",
+        "skylark_rule = rule(",
+        "  implementation = _impl,",
+        "  attrs = {",
+        "    'deps': attr.label_list(),",
+        "    '_base': attr.label(default = '//a:base'),",
+        "  }",
+        ")");
+    scratch.file(
+        "a/BUILD",
+        "load(':skylark.bzl', 'skylark_rule')",
+        // ensure that all Skylark rules get the TestConfiguration fragment
+        "test_base(name = 'base')",
+        // skylark rules get trimmed
+        "skylark_rule(name = 'skylark_solo', deps = [':base'])");
+
+    ConfiguredTarget configuredTarget;
+    BuildConfiguration config;
+
+    configuredTarget = Iterables.getOnlyElement(update("//a:skylark_solo").getTargetsToBuild());
+    config = getConfiguration(configuredTarget);
+    assertThat(config.getFragment(TestConfiguration.class).getTestArguments())
+        .containsExactly(
+            "first trimming transition for //a:skylark_solo",
+            "second trimming transition for //a:skylark_solo")
+        .inOrder();
+  }
+
+  @Test
   public void testRuleClassTransition() throws Exception {
     setRulesAvailableInTests(TestAspects.BASE_RULE, TEST_BASE_RULE, ATTRIBUTE_TRANSITION_RULE,
         RULE_CLASS_TRANSITION_RULE);