Optimize storage of output files in Rule.

* Store the `outputKey` in `OutputFile` because it costs nothing extra to do so (2 fields vs 3 fields both 24 bytes). This removes storage of the `outputKey` string in `flattenedOutputFiles`, reducing the size of the list.
* Encode whether output files are explicit in the class type so that we don't need to store an `int` counting the number of implicit outputs in `Rule`. Removing the `int` field reduces the shallow heap size of `Rule` by 8 bytes.
* Optimize for the common case of a rule with a single output file by storing it bare. Also forego `ImmutableList` wrappers in favor of `OutputFile[]`.

I thought about making `Rule` immutable with subclasses for the different number of output files, but that would require some major refactoring.

PiperOrigin-RevId: 518961414
Change-Id: I03d8ff4d883f2e2ea4468e150b2081a293eb14af

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

@@ -18,48 +18,67 @@
 import java.util.List;
 import net.starlark.java.syntax.Location;
 
-/**
- * A generated file that is the output of a rule.
- */
-public final class OutputFile extends FileTarget {
+/** A generated file that is the output of a rule. */
+public abstract class OutputFile extends FileTarget {
 
-  private final Rule generatingRule;
+  /**
+   * Constructs an implicit output file with the given label, which must be in the generating rule's
+   * package.
+   *
+   * @param outputKey either the map key returned by {@link
+   *     ImplicitOutputsFunction.StarlarkImplicitOutputsFunction#calculateOutputs} or the empty
+   *     string for natively defined implicit outputs
+   */
+  static OutputFile createImplicit(Label label, Rule generatingRule, String outputKey) {
+    return new Implicit(label, generatingRule, outputKey);
+  }
 
   /**
-   * Constructs an OutputFile with the given label, which must be in the generating rule's package.
+   * Constructs an explicit output file with the given label, which must be in the generating rule's
+   * package.
+   *
+   * @param attrName the output attribute's name; used as the {@linkplain #getOutputKey output key}
    */
-  OutputFile(Label label, Rule generatingRule) {
+  static OutputFile createExplicit(Label label, Rule generatingRule, String attrName) {
+    return new Explicit(label, generatingRule, attrName);
+  }
+
+  private final Rule generatingRule;
+  private final String outputKey;
+
+  private OutputFile(Label label, Rule generatingRule, String outputKey) {
     super(generatingRule.getPackage(), label);
     this.generatingRule = generatingRule;
+    this.outputKey = outputKey;
   }
 
   @Override
-  public RuleVisibility getVisibility() {
+  public final RuleVisibility getVisibility() {
     return generatingRule.getVisibility();
   }
 
   @Override
-  public Iterable<Label> getVisibilityDependencyLabels() {
+  public final Iterable<Label> getVisibilityDependencyLabels() {
     return generatingRule.getVisibilityDependencyLabels();
   }
 
   @Override
-  public List<Label> getVisibilityDeclaredLabels() {
+  public final List<Label> getVisibilityDeclaredLabels() {
     return generatingRule.getVisibilityDeclaredLabels();
   }
 
   @Override
-  public boolean isConfigurable() {
+  public final boolean isConfigurable() {
     return true;
   }
 
   /** Returns the rule which generates this output file. */
-  public Rule getGeneratingRule() {
+  public final Rule getGeneratingRule() {
     return generatingRule;
   }
 
   @Override
-  public Package getPackage() {
+  public final Package getPackage() {
     return generatingRule.getPackage();
   }
 
@@ -73,30 +92,69 @@ public enum Kind {
     FILESET
   }
 
-  /**
-   * Returns the kind of this output file.
-   */
-  public Kind getKind() {
+  /** Returns the kind of this output file. */
+  public final Kind getKind() {
     return generatingRule.getRuleClassObject().getOutputFileKind();
   }
 
   @Override
-  public String getTargetKind() {
+  public final String getTargetKind() {
     return targetKind();
   }
 
   @Override
-  public Rule getAssociatedRule() {
+  public final Rule getAssociatedRule() {
     return generatingRule;
   }
 
   @Override
-  public Location getLocation() {
+  public final Location getLocation() {
     return generatingRule.getLocation();
   }
 
+  /**
+   * Returns this output file's output key.
+   *
+   * <p>An output key is an identifier used to access the output in {@code ctx.outputs}, or the
+   * empty string in the case of an output that's not exposed there. For explicit outputs, the
+   * output key is the name of the attribute under which that output appears. For Starlark-defined
+   * implicit outputs, the output key is determined by the dict returned from the Starlark function.
+   * Native-defined implicit outputs are not named in this manner, and so are invisible to {@code
+   * ctx.outputs} and use the empty string key. (It'd be pathological for the empty string to be
+   * used as a key in the other two cases, but this class makes no attempt to prohibit that.)
+   */
+  final String getOutputKey() {
+    return outputKey;
+  }
+
+  abstract boolean isImplicit();
+
   /** Returns the target kind for all output files. */
   public static String targetKind() {
     return "generated file";
   }
+
+  private static final class Implicit extends OutputFile {
+
+    Implicit(Label label, Rule generatingRule, String outputKey) {
+      super(label, generatingRule, outputKey);
+    }
+
+    @Override
+    boolean isImplicit() {
+      return true;
+    }
+  }
+
+  private static final class Explicit extends OutputFile {
+
+    Explicit(Label label, Rule generatingRule, String attrName) {
+      super(label, generatingRule, attrName);
+    }
+
+    @Override
+    boolean isImplicit() {
+      return false;
+    }
+  }
 }

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

@@ -84,6 +84,8 @@ public class Rule implements Target, DependencyFilter.AttributeInfoProvider {
 
   private static final int ATTR_SIZE_THRESHOLD = 126;
 
+  private static final OutputFile[] NO_OUTPUTS = new OutputFile[0];
+
   private final Package pkg;
   private final Label label;
   private final RuleClass ruleClass;
@@ -126,34 +128,22 @@ public class Rule implements Target, DependencyFilter.AttributeInfoProvider {
   private boolean containsErrors = false;
 
   /**
-   * A compact representation of a multimap from "output keys" to output files.
-   *
-   * <p>An output key is an identifier used to access the output in {@code ctx.outputs}, or the
-   * empty string in the case of an output that's not exposed there. For explicit outputs, the
-   * output key is the name of the attribute under which that output appears. For Starlark-defined
-   * implicit outputs, the output key is determined by the dict returned from the Starlark function.
-   * Native-defined implicit outputs are not named in this manner, and so are invisible to {@code
-   * ctx.outputs} and use the empty string key. (It'd be pathological for the empty string to be
-   * used as a key in the other two cases, but this class makes no attempt to prohibit that.)
+   * Output files generated by this rule.
    *
-   * <p>Rather than naively store an ImmutableListMultimap, we save space by compressing it as an
-   * ImmutableList, where each key is followed by all the values having that key. We distinguish
-   * keys (Strings) from values (OutputFiles) by the fact that they have different types. The
-   * accessor methods traverse the list and create a more user-friendly view.
+   * <p>To save memory, this field is either {@link #NO_OUTPUTS} for zero outputs, an {@link
+   * OutputFile} for a single output, or an {@code OutputFile[]} for multiple outputs.
    *
-   * <p>To distinguish implicit outputs from explicit outputs, we store all the implicit outputs in
-   * the list first, and record how many implicit output keys there are in a separate field.
+   * <p>In the case of multiple outputs, all implicit outputs come before any explicit outputs in
+   * the array.
    *
    * <p>The order of the implicit outputs is the same as returned by the implicit output function.
    * This allows a native rule implementation and native implicit outputs function to agree on the
    * index of a given kind of output. The order of explicit outputs preserves the attribute
    * iteration order and the order of values in a list attribute; the latter is important so that
    * {@code ctx.outputs.some_list} has a well-defined order.
    */
-  // Both of these fields are initialized by populateOutputFiles().
-  private ImmutableList<Object> flattenedOutputFileMap;
-
-  private int numImplicitOutputKeys;
+  // Initialized by populateOutputFilesInternal().
+  private Object outputFiles;
 
   Rule(Package pkg, Label label, RuleClass ruleClass, Location location, CallStack callstack) {
     this.pkg = checkNotNull(pkg);
@@ -262,8 +252,8 @@ public boolean isConfigurableAttribute(String attributeName) {
   }
 
   /**
-   * Returns the attribute definition whose name is {@code attrName}, or null
-   * if not found.  (Use get[X]Attr for the actual value.)
+   * Returns the attribute definition whose name is {@code attrName}, or null if not found. (Use
+   * get[X]Attr for the actual value.)
    *
    * @deprecated use {@link AbstractAttributeMapper#getAttributeDefinition} instead
    */
@@ -285,14 +275,7 @@ public Attribute getAttributeDefinition(String attrName) {
    * outputs will retain the relative order in which they were declared.
    */
   public ImmutableList<OutputFile> getOutputFiles() {
-    // Discard the String keys, taking only the OutputFile values.
-    ImmutableList.Builder<OutputFile> result = ImmutableList.builder();
-    for (Object o : flattenedOutputFileMap) {
-      if (o instanceof OutputFile) {
-        result.add((OutputFile) o);
-      }
-    }
-    return result.build();
+    return ImmutableList.copyOf(outputFilesArray());
   }
 
   /**
@@ -301,15 +284,11 @@ public ImmutableList<OutputFile> getOutputFiles() {
    */
   ImmutableList<OutputFile> getImplicitOutputFiles() {
     ImmutableList.Builder<OutputFile> result = ImmutableList.builder();
-    int seenKeys = 0;
-    for (Object o : flattenedOutputFileMap) {
-      if (o instanceof String) {
-        if (++seenKeys > numImplicitOutputKeys) {
-          break;
-        }
-      } else {
-        result.add((OutputFile) o);
+    for (OutputFile output : outputFilesArray()) {
+      if (!output.isImplicit()) {
+        break;
       }
+      result.add(output);
     }
     return result.build();
   }
@@ -326,14 +305,9 @@ ImmutableList<OutputFile> getImplicitOutputFiles() {
    */
   public ImmutableListMultimap<String, OutputFile> getExplicitOutputFileMap() {
     ImmutableListMultimap.Builder<String, OutputFile> result = ImmutableListMultimap.builder();
-    int seenKeys = 0;
-    String key = null;
-    for (Object o : flattenedOutputFileMap) {
-      if (o instanceof String) {
-        seenKeys++;
-        key = (String) o;
-      } else if (seenKeys > numImplicitOutputKeys) {
-        result.put(key, (OutputFile) o);
+    for (OutputFile output : outputFilesArray()) {
+      if (!output.isImplicit()) {
+        result.put(output.getOutputKey(), output);
       }
     }
     return result.build();
@@ -353,21 +327,21 @@ public ImmutableMap<String, OutputFile> getStarlarkImplicitOutputFileMap() {
       return ImmutableMap.of();
     }
     ImmutableMap.Builder<String, OutputFile> result = ImmutableMap.builder();
-    int seenKeys = 0;
-    String key = null;
-    for (Object o : flattenedOutputFileMap) {
-      if (o instanceof String) {
-        if (++seenKeys > numImplicitOutputKeys) {
-          break;
-        }
-        key = (String) o;
-      } else {
-        result.put(key, (OutputFile) o);
+    for (OutputFile output : outputFilesArray()) {
+      if (!output.isImplicit()) {
+        break;
       }
+      result.put(output.getOutputKey(), output);
     }
     return result.buildOrThrow();
   }
 
+  private OutputFile[] outputFilesArray() {
+    return outputFiles instanceof OutputFile
+        ? new OutputFile[] {(OutputFile) outputFiles}
+        : (OutputFile[]) outputFiles;
+  }
+
   @Override
   public Location getLocation() {
     return location;
@@ -916,14 +890,9 @@ private void populateOutputFilesInternal(
       ImplicitOutputsFunction implicitOutputsFunction,
       boolean checkLabels)
       throws LabelSyntaxException, InterruptedException {
-    Preconditions.checkState(flattenedOutputFileMap == null);
-
-    // We associate each output with its String key (or empty string if there's no key) as we go,
-    // and compress it down to a flat list afterwards. We use ImmutableListMultimap because it's
-    // more efficient than LinkedListMultimap and provides ordering guarantees among keys (whereas
-    // ArrayListMultimap doesn't).
-    ImmutableListMultimap.Builder<String, OutputFile> outputFileMap =
-        ImmutableListMultimap.builder();
+    Preconditions.checkState(outputFiles == null);
+
+    List<OutputFile> outputs = new ArrayList<>();
     // Detects collisions where the same output key is used for both an explicit and implicit entry.
     HashSet<String> implicitOutputKeys = new HashSet<>();
 
@@ -953,8 +922,7 @@ private void populateOutputFilesInternal(
           }
           validateOutputLabel(label, eventHandler);
 
-          OutputFile file = new OutputFile(label, this);
-          outputFileMap.put(outputKey, file);
+          outputs.add(OutputFile.createImplicit(label, this, outputKey));
           implicitOutputKeys.add(outputKey);
         };
 
@@ -1007,8 +975,7 @@ private void populateOutputFilesInternal(
           }
           validateOutputLabel(outputLabel, eventHandler);
 
-          OutputFile outputFile = new OutputFile(outputLabel, this);
-          outputFileMap.put(attrName, outputFile);
+          outputs.add(OutputFile.createExplicit(outputLabel, this, attrName));
         };
 
     // Populate the explicit outputs.
@@ -1029,16 +996,13 @@ private void populateOutputFilesInternal(
       }
     }
 
-    // Flatten the result into the final list.
-    ImmutableList.Builder<Object> builder = ImmutableList.builder();
-    for (Map.Entry<String, Collection<OutputFile>> e : outputFileMap.build().asMap().entrySet()) {
-      builder.add(e.getKey());
-      for (OutputFile out : e.getValue()) {
-        builder.add(out);
-      }
+    if (outputs.isEmpty()) {
+      outputFiles = NO_OUTPUTS;
+    } else if (outputs.size() == 1) {
+      outputFiles = outputs.get(0);
+    } else {
+      outputFiles = outputs.toArray(OutputFile[]::new);
     }
-    flattenedOutputFileMap = builder.build();
-    numImplicitOutputKeys = implicitOutputKeys.size();
   }
 
   private void validateOutputLabel(Label label, EventHandler eventHandler) {