Don't call Starlark implicit outputs functions twice

Previously, a Starlark rule's implicit outputs function would be called twice: once during loading to help populate the complete set of outputs, and then again at the start of analysis while constructing the `ctx.outputs` object. This double Starlark evaluation was not only wasteful, but caused StarlarkRuleContext's constructor to potentially throw EvalException and InterruptedException, complicating refactoring of that code. This change makes it so we remember the result of the first call in Rule and access it later.

Initial versions of this CL attempted to store the implicit output function results in a new field on Rule. However, this added ~0.5% RAM overhead to a large `build --nobuild` invocation, and ~1.5% on a `query 'deps(...)'` of the same target. (This increase was not attributable simply to the addition of the field itself, but to the map entries it contained.)

The new strategy is to take an existing field that held all the outputs (implicit and explicit), and expand it to also include the string keys of Starlark implicit outputs. Not only does this avoid adding new redundant fields, it allows us to contract an existing one. The tradeoff is that the accessors must do some computational work to copy only the relevant data from the field. The stored map is compressed as a flat list for further space savings.

Along the way, I refactored the code that populates this collection of outputs. The diff is larger than I would've liked, but essentially the bodies of populateExplicitOutputFiles() and populateImplicitOutputFiles() got lifted into lambdas in the caller (with a little inlining of helpers), to better update state during the computation. There may be further refactoring that could be done, e.g. changing the mode of one failure to do a rule error instead of a LabelSyntaxException, or combining validation logic into one place, but I'm trying to avoid any more semantic changes than are absolutely necessary.

On the benchmark alluded to above, we save 0.93% RAM on a query invocation, and 0.40% on a build --nobuild.

Cleanup work toward #11437.

PiperOrigin-RevId: 354631199

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

@@ -1606,6 +1606,8 @@ public Artifact getImplicitOutputArtifact(String path) {
    * content-based path for this artifact (see {@link
    * BuildConfiguration#useContentBasedOutputPaths()}).
    */
+  // TODO(bazel-team): Consider removing contentBasedPath stuff, which is unused as of 18 months
+  // after its introduction in cl/252148134.
   private Artifact getImplicitOutputArtifact(String path, boolean contentBasedPath) {
     return getPackageRelativeArtifact(path, getBinOrGenfilesDirectory(), contentBasedPath);
   }

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

@@ -334,6 +334,10 @@ public StarlarkCallable rule(
 
     if (implicitOutputs != Starlark.NONE) {
       if (implicitOutputs instanceof StarlarkFunction) {
+        // TODO(brandjon): Embedding bazelContext in a callback is not thread safe! Instead
+        // construct a new BazelStarlarkContext with copies of the relevant fields that are safe to
+        // share. Maybe create a method on BazelStarlarkContext for safely constructing a child
+        // context.
         StarlarkCallbackHelper callback =
             new StarlarkCallbackHelper(
                 (StarlarkFunction) implicitOutputs, thread.getSemantics(), bazelContext);

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

@@ -84,14 +84,8 @@ public static ConfiguredTarget buildRule(
       String toolsRepository)
       throws InterruptedException, RuleErrorException, ActionConflictException {
     String expectFailure = ruleContext.attributes().get("expect_failure", Type.STRING);
-    StarlarkRuleContext starlarkRuleContext = null;
+    StarlarkRuleContext starlarkRuleContext = new StarlarkRuleContext(ruleContext, null);
     try {
-      // The StarlarkRuleContext constructor may throw EvalException due
-      // to StarlarkImplicitOutputsFunction, which shouldn't be called at
-      // analysis time anyway since all such targets should be created during loading.
-      // TODO(adonovan): clean it up.
-      starlarkRuleContext = new StarlarkRuleContext(ruleContext, null);
-
       RuleClass ruleClass = ruleContext.getRule().getRuleClassObject();
       if (ruleClass.getRuleClassType().equals(RuleClass.Builder.RuleClassType.WORKSPACE)) {
         ruleContext.ruleError(

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

@@ -21,6 +21,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
@@ -61,11 +62,10 @@
 import com.google.devtools.build.lib.packages.BuildSetting;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.ImplicitOutputsFunction;
-import com.google.devtools.build.lib.packages.ImplicitOutputsFunction.StarlarkImplicitOutputsFunction;
 import com.google.devtools.build.lib.packages.OutputFile;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.Provider;
-import com.google.devtools.build.lib.packages.RawAttributeMapper;
+import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.RuleClass.ConfiguredTargetFactory.RuleErrorException;
 import com.google.devtools.build.lib.packages.StructImpl;
 import com.google.devtools.build.lib.packages.StructProvider;
@@ -137,10 +137,9 @@ public final class StarlarkRuleContext implements StarlarkRuleContextApi<Constra
    *
    * @param aspectDescriptor aspect for which the context is created, or <code>null</code> if it is
    *     for a rule.
-   * @throws InterruptedException
    */
   public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor aspectDescriptor)
-      throws EvalException, InterruptedException, RuleErrorException {
+      throws RuleErrorException {
     // Init ruleContext first, we need it to obtain the StarlarkSemantics used by
     // StarlarkActionFactory (and possibly others).
     this.ruleContext = Preconditions.checkNotNull(ruleContext);
@@ -150,40 +149,37 @@ public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor a
     this.hostFragments = new FragmentCollection(ruleContext, HostTransition.INSTANCE);
     this.aspectDescriptor = aspectDescriptor;
 
+    Rule rule = ruleContext.getRule();
+
     if (aspectDescriptor == null) {
       this.isForAspect = false;
-      Collection<Attribute> attributes = ruleContext.getRule().getAttributes();
-      Outputs outputs = new Outputs(this);
-
-      ImplicitOutputsFunction implicitOutputsFunction =
-          ruleContext.getRule().getImplicitOutputsFunction();
-
-      if (implicitOutputsFunction instanceof StarlarkImplicitOutputsFunction) {
-        StarlarkImplicitOutputsFunction func =
-            (StarlarkImplicitOutputsFunction) implicitOutputsFunction;
-        for (Map.Entry<String, String> entry :
-            func.calculateOutputs(
-                    ruleContext.getAnalysisEnvironment().getEventHandler(),
-                    RawAttributeMapper.of(ruleContext.getRule()))
-                .entrySet()) {
-          outputs.addOutput(
-              entry.getKey(),
-              ruleContext.getImplicitOutputArtifact(entry.getValue()));
-        }
-      }
+      Collection<Attribute> attributes = rule.getAttributes();
 
+      // Populate ctx.outputs.
+      Outputs outputs = new Outputs(this);
+      // These getters do some computational work to return a view, so ensure we only do it once.
+      ImmutableListMultimap<String, OutputFile> explicitOutMap = rule.getExplicitOutputFileMap();
+      ImmutableMap<String, OutputFile> implicitOutMap = rule.getStarlarkImplicitOutputFileMap();
+      // Add the explicit outputs -- values of attributes of type OUTPUT or OUTPUT_LIST.
+      // We must iterate over the attribute definitions, and not just the entries in the
+      // explicitOutMap, because the latter omits empty output attributes, which must still
+      // generate None or [] fields in the struct.
       for (Attribute a : attributes) {
+        // Skip non-output attrs.
         String attrName = a.getName();
         Type<?> type = a.getType();
         if (type.getLabelClass() != LabelClass.OUTPUT) {
           continue;
         }
+
+        // Grab all associated outputs.
         ImmutableList.Builder<Artifact> artifactsBuilder = ImmutableList.builder();
-        for (OutputFile outputFile : ruleContext.getRule().getOutputFileMap().get(attrName)) {
+        for (OutputFile outputFile : explicitOutMap.get(attrName)) {
           artifactsBuilder.add(ruleContext.createOutputArtifact(outputFile));
         }
         StarlarkList<Artifact> artifacts = StarlarkList.immutableCopyOf(artifactsBuilder.build());
 
+        // For singular output attributes, unwrap sole element or else use None for arity mismatch.
         if (type == BuildType.OUTPUT) {
           if (artifacts.size() == 1) {
             outputs.addOutput(attrName, Iterables.getOnlyElement(artifacts));
@@ -193,15 +189,23 @@ public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor a
         } else if (type == BuildType.OUTPUT_LIST) {
           outputs.addOutput(attrName, artifacts);
         } else {
-          throw new IllegalArgumentException(
-              "Type of " + attrName + "(" + type + ") is not output type ");
+          throw new AssertionError(
+              String.format("Attribute %s has unexpected output type %s", attrName, type));
         }
       }
+      // Add the implicit outputs. In the case where the rule has a native-defined implicit outputs
+      // function, nothing is added. Note that Rule ensures that Starlark-defined implicit output
+      // keys don't conflict with output attribute names.
+      // TODO(bazel-team): Also see about requiring the key to be a valid Starlark identifier.
+      for (Map.Entry<String, OutputFile> e : implicitOutMap.entrySet()) {
+        outputs.addOutput(e.getKey(), ruleContext.createOutputArtifact(e.getValue()));
+      }
 
       this.outputsObject = outputs;
 
+      // Populate ctx.attr.
       StarlarkAttributesCollection.Builder builder = StarlarkAttributesCollection.builder(this);
-      for (Attribute attribute : ruleContext.getRule().getAttributes()) {
+      for (Attribute attribute : attributes) {
         Object value = ruleContext.attributes().get(attribute.getName(), attribute.getType());
         builder.addAttribute(attribute, value);
       }
@@ -212,13 +216,13 @@ public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor a
     } else { // ASPECT
       this.isForAspect = true;
       this.outputsObject = null;
-
       ImmutableCollection<Attribute> attributes =
           ruleContext.getMainAspect().getDefinition().getAttributes().values();
+
       StarlarkAttributesCollection.Builder aspectBuilder =
           StarlarkAttributesCollection.builder(this);
       for (Attribute attribute : attributes) {
-        Object defaultValue = attribute.getDefaultValue(ruleContext.getRule());
+        Object defaultValue = attribute.getDefaultValue(rule);
         if (defaultValue instanceof ComputedDefault) {
           defaultValue = ((ComputedDefault) defaultValue).getDefault(ruleContext.attributes());
         }
@@ -229,7 +233,7 @@ public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor a
       this.splitAttributes = null;
       StarlarkAttributesCollection.Builder ruleBuilder = StarlarkAttributesCollection.builder(this);
 
-      for (Attribute attribute : ruleContext.getRule().getAttributes()) {
+      for (Attribute attribute : rule.getAttributes()) {
         Object value = ruleContext.attributes().get(attribute.getName(), attribute.getType());
         ruleBuilder.addAttribute(attribute, value);
       }
@@ -239,7 +243,7 @@ public StarlarkRuleContext(RuleContext ruleContext, @Nullable AspectDescriptor a
           continue;
         }
         for (Attribute attribute : aspect.getDefinition().getAttributes().values()) {
-          Object defaultValue = attribute.getDefaultValue(ruleContext.getRule());
+          Object defaultValue = attribute.getDefaultValue(rule);
           if (defaultValue instanceof ComputedDefault) {
             defaultValue = ((ComputedDefault) defaultValue).getDefault(ruleContext.attributes());
           }
@@ -276,13 +280,14 @@ public Outputs(StarlarkRuleContext context) {
       this.context = context;
     }
 
-    private void addOutput(String key, Object value)
-        throws EvalException {
+    private void addOutput(String key, Object value) throws RuleErrorException {
       Preconditions.checkState(!context.isImmutable(),
           "Cannot add outputs to immutable Outputs object");
+      // TODO(bazel-team): We should reject outputs whose key is not an identifier. Today this is
+      // allowed, and the resulting ctx.outputs value can be retrieved using getattr().
       if (outputs.containsKey(key)
           || (context.isExecutable() && EXECUTABLE_OUTPUT_NAME.equals(key))) {
-        throw Starlark.errorf("Multiple outputs with the same key: %s", key);
+        context.getRuleContext().throwWithRuleError("Multiple outputs with the same key: " + key);
       }
       outputs.put(key, value);
     }