Restructure the NestedSet usage in CcCompilationContext. Goals:
- Don't use NestedSets of NestedSets. Not sure whether this is a performance
  improvement or has serialization benefits, but it does make understanding the
  structure a lot easier.
- Use the same large NestedSet for all header information (whether or not they
  are modular and what pregrepped headers might exist). This speeds up the
  include scanner calls getLegalGeneratedScannerFileMap and getModularHeaders
  because they now iterate over just a single NestedSet and the second
  evaluation of a NestedSet is much cheaper. In a subsequent change, we can
  likely fold all three iterations over this NestedSet into one (in
  getLegalGeneratedScannerFileMap, getModularHeaders and getUsedModules).

Measurements show a small reduction in heap usage.

RELNOTES: None.
PiperOrigin-RevId: 203909308
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationContext.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationContext.java
index 9254e3f..e5bc2f8 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationContext.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCompilationContext.java
@@ -35,10 +35,11 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Objects;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -67,10 +68,10 @@
   /** Non-code mandatory compilation inputs. */
   private final NestedSet<Artifact> nonCodeInputs;
 
-  private final NestedSet<PregreppedHeader> pregreppedHdrs;
-
-  private final ModuleInfo moduleInfo;
-  private final ModuleInfo picModuleInfo;
+  private final HeaderInfo headerInfo;
+  private final NestedSet<HeaderInfo> transitiveHeaderInfos;
+  private final NestedSet<Artifact> transitiveModules;
+  private final NestedSet<Artifact> transitivePicModules;
 
   private final CppModuleMap cppModuleMap;
   private final CppModuleMap verificationModuleMap;
@@ -88,10 +89,11 @@
       NestedSet<PathFragment> declaredIncludeDirs,
       NestedSet<PathFragment> declaredIncludeWarnDirs,
       NestedSet<Artifact> declaredIncludeSrcs,
-      NestedSet<PregreppedHeader> pregreppedHdrs,
       NestedSet<Artifact> nonCodeInputs,
-      ModuleInfo moduleInfo,
-      ModuleInfo picModuleInfo,
+      HeaderInfo headerInfo,
+      NestedSet<HeaderInfo> transitiveHeaderInfos,
+      NestedSet<Artifact> transitiveModules,
+      NestedSet<Artifact> transitivePicModules,
       NestedSet<Artifact> directModuleMaps,
       CppModuleMap cppModuleMap,
       @Nullable CppModuleMap verificationModuleMap,
@@ -102,9 +104,10 @@
     this.declaredIncludeWarnDirs = declaredIncludeWarnDirs;
     this.declaredIncludeSrcs = declaredIncludeSrcs;
     this.directModuleMaps = directModuleMaps;
-    this.pregreppedHdrs = pregreppedHdrs;
-    this.moduleInfo = moduleInfo;
-    this.picModuleInfo = picModuleInfo;
+    this.headerInfo = headerInfo;
+    this.transitiveHeaderInfos = transitiveHeaderInfos;
+    this.transitiveModules = transitiveModules;
+    this.transitivePicModules = transitivePicModules;
     this.cppModuleMap = cppModuleMap;
     this.nonCodeInputs = nonCodeInputs;
     this.verificationModuleMap = verificationModuleMap;
@@ -193,39 +196,67 @@
   }
 
   /** Returns headers given as textual_hdrs in this target. */
-  public ImmutableSet<Artifact> getTextualHdrs() {
-    return moduleInfo.textualHeaders;
+  public Iterable<Artifact> getTextualHdrs() {
+    return headerInfo.textualHeaders;
   }
 
-  /**
-   * Returns the immutable pairs of (header file, pregrepped header file). The value artifacts
-   * (pregrepped header file) are generated by {@link ExtractInclusionAction}.
-   */
-  NestedSet<PregreppedHeader> getPregreppedHeaders() {
-    return pregreppedHdrs;
+  public Map<Artifact, Artifact> createLegalGeneratedScannerFileMap() {
+    Map<Artifact, Artifact> legalOuts = new HashMap<>();
+    for (HeaderInfo transitiveHeaderInfo : transitiveHeaderInfos) {
+      for (Artifact a : transitiveHeaderInfo.modularHeaders) {
+        if (!a.isSourceArtifact()) {
+          legalOuts.put(a, null);
+        }
+      }
+      for (Artifact a : transitiveHeaderInfo.textualHeaders) {
+        if (!a.isSourceArtifact()) {
+          legalOuts.put(a, null);
+        }
+      }
+      for (PregreppedHeader pregreppedHeader : transitiveHeaderInfo.pregreppedHeaders) {
+        Artifact hdr = pregreppedHeader.originalHeader();
+        Preconditions.checkState(!hdr.isSourceArtifact(), hdr);
+        legalOuts.put(hdr, pregreppedHeader.greppedHeader());
+      }
+    }
+    return Collections.unmodifiableMap(legalOuts);
   }
 
   public NestedSet<Artifact> getTransitiveModules(boolean usePic) {
-    return usePic ? picModuleInfo.transitiveModules : moduleInfo.transitiveModules;
+    return usePic ? transitivePicModules : transitiveModules;
   }
 
   public Set<Artifact> getModularHeaders(boolean usePic) {
-    ModuleInfo info = usePic ? picModuleInfo : moduleInfo;
     Set<Artifact> result = new HashSet<>();
-    for (TransitiveModuleHeaders moduleHeaders : info.transitiveModuleHeaders) {
-      result.addAll(moduleHeaders.headers);
+    for (HeaderInfo transitiveHeaderInfo : transitiveHeaderInfos) {
+      if (transitiveHeaderInfo.getModule(usePic) != null) {
+        result.addAll(transitiveHeaderInfo.modularHeaders);
+      }
     }
     // Remove headers belonging to this module.
-    result.removeAll(info.modularHeaders);
-    result.removeAll(info.textualHeaders);
+    result.removeAll(headerInfo.modularHeaders);
+    result.removeAll(headerInfo.textualHeaders);
     return Collections.unmodifiableSet(result);
   }
 
-  public Collection<TransitiveModuleHeaders> getUsedModules(
-      boolean usePic, Set<Artifact> usedHeaders) {
-    return usePic
-        ? picModuleInfo.getUsedModules(usedHeaders)
-        : moduleInfo.getUsedModules(usedHeaders);
+  public Collection<HeaderInfo> getUsedModules(boolean usePic, Set<Artifact> usedHeaders) {
+    List<HeaderInfo> result = new ArrayList<>();
+    for (HeaderInfo transitiveHeaderInfo : transitiveHeaderInfos) {
+      // Do not add the module of the current rule for both:
+      // 1. the module compile itself
+      // 2. compiles of other translation units of the same rule.
+      if (transitiveHeaderInfo.getModule(usePic) == null
+          || transitiveHeaderInfo.getModule(usePic).equals(headerInfo.getModule(usePic))) {
+        continue;
+      }
+      for (Artifact header : transitiveHeaderInfo.modularHeaders) {
+        if (usedHeaders.contains(header)) {
+          result.add(transitiveHeaderInfo);
+          break;
+        }
+      }
+    }
+    return result;
   }
 
   /**
@@ -255,8 +286,8 @@
    */
   protected Set<Artifact> getHeaderModuleSrcs() {
     return new ImmutableSet.Builder<Artifact>()
-        .addAll(moduleInfo.modularHeaders)
-        .addAll(moduleInfo.textualHeaders)
+        .addAll(headerInfo.modularHeaders)
+        .addAll(headerInfo.textualHeaders)
         .build();
   }
 
@@ -282,10 +313,11 @@
         NestedSetBuilder.emptySet(Order.STABLE_ORDER),
         NestedSetBuilder.emptySet(Order.STABLE_ORDER),
         ccCompilationContext.declaredIncludeSrcs,
-        ccCompilationContext.pregreppedHdrs,
         ccCompilationContext.nonCodeInputs,
-        ccCompilationContext.moduleInfo,
-        ccCompilationContext.picModuleInfo,
+        ccCompilationContext.headerInfo,
+        ccCompilationContext.transitiveHeaderInfos,
+        ccCompilationContext.transitiveModules,
+        ccCompilationContext.transitivePicModules,
         ccCompilationContext.directModuleMaps,
         ccCompilationContext.cppModuleMap,
         ccCompilationContext.verificationModuleMap,
@@ -340,11 +372,12 @@
         NestedSetBuilder.stableOrder();
     private final NestedSetBuilder<Artifact> declaredIncludeSrcs =
         NestedSetBuilder.stableOrder();
-    private final NestedSetBuilder<PregreppedHeader> pregreppedHdrs =
-        NestedSetBuilder.stableOrder();
     private final NestedSetBuilder<Artifact> nonCodeInputs = NestedSetBuilder.stableOrder();
-    private final ModuleInfo.Builder moduleInfo = new ModuleInfo.Builder();
-    private final ModuleInfo.Builder picModuleInfo = new ModuleInfo.Builder();
+    private final HeaderInfo.Builder headerInfoBuilder = new HeaderInfo.Builder();
+    private final NestedSetBuilder<HeaderInfo> transitiveHeaderInfo =
+        NestedSetBuilder.stableOrder();
+    private final NestedSetBuilder<Artifact> transitiveModules = NestedSetBuilder.stableOrder();
+    private final NestedSetBuilder<Artifact> transitivePicModules = NestedSetBuilder.stableOrder();
     private final NestedSetBuilder<Artifact> directModuleMaps = NestedSetBuilder.stableOrder();
     private final Set<String> defines = new LinkedHashSet<>();
     private CppModuleMap cppModuleMap;
@@ -392,9 +425,16 @@
       declaredIncludeDirs.addTransitive(otherCcCompilationContext.getDeclaredIncludeDirs());
       declaredIncludeWarnDirs.addTransitive(otherCcCompilationContext.getDeclaredIncludeWarnDirs());
       declaredIncludeSrcs.addTransitive(otherCcCompilationContext.getDeclaredIncludeSrcs());
-      pregreppedHdrs.addTransitive(otherCcCompilationContext.getPregreppedHeaders());
-      moduleInfo.addTransitive(otherCcCompilationContext.moduleInfo);
-      picModuleInfo.addTransitive(otherCcCompilationContext.picModuleInfo);
+      transitiveHeaderInfo.addTransitive(otherCcCompilationContext.transitiveHeaderInfos);
+      transitiveModules.addTransitive(otherCcCompilationContext.transitiveModules);
+      if (otherCcCompilationContext.headerInfo.headerModule != null) {
+        transitiveModules.add(otherCcCompilationContext.headerInfo.headerModule);
+      }
+      transitivePicModules.addTransitive(otherCcCompilationContext.transitivePicModules);
+      if (otherCcCompilationContext.headerInfo.picHeaderModule != null) {
+        transitivePicModules.add(otherCcCompilationContext.headerInfo.picHeaderModule);
+      }
+
       nonCodeInputs.addTransitive(otherCcCompilationContext.nonCodeInputs);
 
       // All module maps of direct dependencies are inputs to the current compile independently of
@@ -528,14 +568,12 @@
     }
 
     public Builder addModularHdrs(Collection<Artifact> headers) {
-      this.moduleInfo.addHeaders(headers);
-      this.picModuleInfo.addHeaders(headers);
+      this.headerInfoBuilder.addHeaders(headers);
       return this;
     }
 
     public Builder addTextualHdrs(Collection<Artifact> headers) {
-      this.moduleInfo.addTextualHeaders(headers);
-      this.picModuleInfo.addTextualHeaders(headers);
+      this.headerInfoBuilder.addTextualHeaders(headers);
       return this;
     }
 
@@ -549,7 +587,7 @@
               .stream()
               .map(pregreppedHeader -> pregreppedHeader.greppedHeader())
               .collect(Collectors.toList()));
-      this.pregreppedHdrs.addAll(pregrepped);
+      this.headerInfoBuilder.addPregreppedHeaders(pregrepped);
       return this;
     }
 
@@ -601,7 +639,7 @@
      * @param headerModule The .pcm file generated for this library.
      */
     public Builder setHeaderModule(Artifact headerModule) {
-      this.moduleInfo.setHeaderModule(headerModule);
+      this.headerInfoBuilder.setHeaderModule(headerModule);
       return this;
     }
 
@@ -610,7 +648,7 @@
      * @param picHeaderModule The .pic.pcm file generated for this library.
      */
     public Builder setPicHeaderModule(Artifact picHeaderModule) {
-      this.picModuleInfo.setHeaderModule(picHeaderModule);
+      this.headerInfoBuilder.setPicHeaderModule(picHeaderModule);
       return this;
     }
 
@@ -623,10 +661,9 @@
 
     @VisibleForTesting // productionVisibility = Visibility.PRIVATE
     public CcCompilationContext build(ActionOwner owner, MiddlemanFactory middlemanFactory) {
-      Preconditions.checkState(
-          Objects.equals(moduleInfo.textualHeaders, picModuleInfo.textualHeaders),
-          "Module and PIC module's textual headers are expected to be identical");
       Artifact prerequisiteStampFile = createMiddleman(owner, middlemanFactory);
+      HeaderInfo headerInfo = headerInfoBuilder.build();
+      transitiveHeaderInfo.add(headerInfo);
 
       return new CcCompilationContext(
           new CommandLineCcCompilationContext(
@@ -643,10 +680,11 @@
           declaredIncludeDirs.build(),
           declaredIncludeWarnDirs.build(),
           declaredIncludeSrcs.build(),
-          pregreppedHdrs.build(),
           nonCodeInputs.build(),
-          moduleInfo.build(),
-          picModuleInfo.build(),
+          headerInfo,
+          transitiveHeaderInfo.build(),
+          transitiveModules.build(),
+          transitivePicModules.build(),
           directModuleMaps.build(),
           cppModuleMap,
           verificationModuleMap,
@@ -692,86 +730,64 @@
   }
 
   /**
-   * Gathers data about the direct and transitive .pcm files belonging to this context. Can be to
-   * either gather data on PIC or on no-PIC .pcm files.
+   * Gathers data about the PIC and no-PIC .pcm files belonging to this context and the associated
+   * information about the headers, e.g. modular vs. textual headers and pre-grepped header files.
    */
   @Immutable
   @AutoCodec
-  public static final class ModuleInfo {
+  public static final class HeaderInfo {
     /**
-     * The module built for this context. If null, then no module is being compiled for this
+     * The modules built for this context. If null, then no module is being compiled for this
      * context.
      */
     private final Artifact headerModule;
+    private final Artifact picHeaderModule;
 
     /** All header files that are compiled into this module. */
-    private final ImmutableSet<Artifact> modularHeaders;
+    private final ImmutableList<Artifact> modularHeaders;
 
     /** All header files that are contained in this module. */
-    private final ImmutableSet<Artifact> textualHeaders;
+    private final ImmutableList<Artifact> textualHeaders;
 
-    /**
-     * All transitive modules that this context depends on, excluding headerModule.
-     */
-    private final NestedSet<Artifact> transitiveModules;
+    private final ImmutableList<PregreppedHeader> pregreppedHeaders;
 
-    /**
-     * All information about mapping transitive headers to transitive modules.
-     */
-    public final NestedSet<TransitiveModuleHeaders> transitiveModuleHeaders;
-
-    public ModuleInfo(
+    public HeaderInfo(
         Artifact headerModule,
-        ImmutableSet<Artifact> modularHeaders,
-        ImmutableSet<Artifact> textualHeaders,
-        NestedSet<Artifact> transitiveModules,
-        NestedSet<TransitiveModuleHeaders> transitiveModuleHeaders) {
+        Artifact picHeaderModule,
+        ImmutableList<Artifact> modularHeaders,
+        ImmutableList<Artifact> textualHeaders,
+        ImmutableList<PregreppedHeader> pregreppedHeaders) {
       this.headerModule = headerModule;
+      this.picHeaderModule = picHeaderModule;
       this.modularHeaders = modularHeaders;
       this.textualHeaders = textualHeaders;
-      this.transitiveModules = transitiveModules;
-      this.transitiveModuleHeaders = transitiveModuleHeaders;
+      this.pregreppedHeaders = pregreppedHeaders;
     }
 
-    public Collection<TransitiveModuleHeaders> getUsedModules(Set<Artifact> usedHeaders) {
-      List<TransitiveModuleHeaders> result = new ArrayList<>();
-      for (TransitiveModuleHeaders transitiveModule : transitiveModuleHeaders) {
-        if (transitiveModule.module.equals(headerModule)) {
-          // Do not add the module of the current rule for both:
-          // 1. the module compile itself
-          // 2. compiles of other translation units of the same rule.
-          continue;
-        }
-        boolean providesUsedHeader = false;
-        for (Artifact header : transitiveModule.headers) {
-          if (usedHeaders.contains(header)) {
-            providesUsedHeader = true;
-            break;
-          }
-        }
-        if (providesUsedHeader) {
-          result.add(transitiveModule);
-        }
-      }
-      return result;
+    public Artifact getModule(boolean pic) {
+      return pic ? picHeaderModule : headerModule;
     }
 
     /**
-     * Builder class for {@link ModuleInfo}.
+     * Builder class for {@link HeaderInfo}.
      */
     public static class Builder {
       private Artifact headerModule = null;
-      private final Set<Artifact> modularHeaders = new LinkedHashSet<>();
-      private final Set<Artifact> textualHeaders = new LinkedHashSet<>();
-      private final NestedSetBuilder<Artifact> transitiveModules = NestedSetBuilder.stableOrder();
-      private final NestedSetBuilder<TransitiveModuleHeaders> transitiveModuleHeaders =
-          NestedSetBuilder.stableOrder();
+      private Artifact picHeaderModule = null;
+      private final Set<Artifact> modularHeaders = new HashSet<>();
+      private final Set<Artifact> textualHeaders = new HashSet<>();
+      private final Set<PregreppedHeader> pregreppedHeaders = new HashSet<>();
 
       public Builder setHeaderModule(Artifact headerModule) {
         this.headerModule = headerModule;
         return this;
       }
 
+      public Builder setPicHeaderModule(Artifact headerModule) {
+        this.picHeaderModule = headerModule;
+        return this;
+      }
+
       public Builder addHeaders(Collection<Artifact> headers) {
         // TODO(djasper): CPP_TEXTUAL_INCLUDEs are currently special cased here and in
         // CppModuleMapAction. These should be moved to a place earlier in the Action construction.
@@ -790,55 +806,19 @@
         return this;
       }
 
-      /**
-       * Adds the {@link ModuleInfo} of a dependency and builds up the transitive data structures.
-       */
-      public Builder addTransitive(ModuleInfo moduleInfo) {
-        if (moduleInfo.headerModule != null) {
-          transitiveModules.add(moduleInfo.headerModule);
-        }
-        transitiveModules.addTransitive(moduleInfo.transitiveModules);
-        transitiveModuleHeaders.addTransitive(moduleInfo.transitiveModuleHeaders);
+      public Builder addPregreppedHeaders(Collection<PregreppedHeader> pregreppedHeaders) {
+        this.pregreppedHeaders.addAll(pregreppedHeaders);
         return this;
       }
 
-      public ModuleInfo build() {
-        ImmutableSet<Artifact> modularHeaders = ImmutableSet.copyOf(this.modularHeaders);
-        NestedSet<Artifact> transitiveModules = this.transitiveModules.build();
-        if (headerModule != null) {
-          transitiveModuleHeaders.add(new TransitiveModuleHeaders(headerModule, modularHeaders));
-        }
-        return new ModuleInfo(
+      public HeaderInfo build() {
+        return new HeaderInfo(
             headerModule,
-            modularHeaders,
-            ImmutableSet.copyOf(this.textualHeaders),
-            transitiveModules,
-            transitiveModuleHeaders.build());
+            picHeaderModule,
+            ImmutableList.copyOf(modularHeaders),
+            ImmutableList.copyOf(textualHeaders),
+            ImmutableList.copyOf(pregreppedHeaders));
       }
     }
   }
-
-  /** Collects data for a specific module in a special format that makes pruning easy. */
-  @Immutable
-  @AutoCodec
-  public static final class TransitiveModuleHeaders {
-    /**
-     * The module that we are calculating information for.
-     */
-    private final Artifact module;
-
-    /**
-     * The headers compiled into this module.
-     */
-    private final ImmutableSet<Artifact> headers;
-
-    public TransitiveModuleHeaders(Artifact module, ImmutableSet<Artifact> headers) {
-      this.module = module;
-      this.headers = headers;
-    }
-
-    public Artifact getModule() {
-      return module;
-    }
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
index 4e81e72..1881fb7 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
@@ -56,9 +56,9 @@
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.rules.cpp.CcCommon.CoptsFilter;
+import com.google.devtools.build.lib.rules.cpp.CcCompilationContext.HeaderInfo;
 import com.google.devtools.build.lib.rules.cpp.CcToolchainFeatures.FeatureConfiguration;
 import com.google.devtools.build.lib.rules.cpp.CppCompileActionContext.Reply;
-import com.google.devtools.build.lib.rules.cpp.CppHelper.PregreppedHeader;
 import com.google.devtools.build.lib.skyframe.ActionExecutionValue;
 import com.google.devtools.build.lib.util.DependencySet;
 import com.google.devtools.build.lib.util.Fingerprint;
@@ -73,8 +73,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
@@ -400,9 +398,9 @@
       // usedHeadersAndModules only contains headers now, so we can pass it to getUsedModules()
       // (and even if it contained other things, it's only used to check for the presence of headers
       // so it would not matter)
-      for (CcCompilationContext.TransitiveModuleHeaders usedModule :
+      for (HeaderInfo usedModule :
           ccCompilationContext.getUsedModules(usePic, usedHeadersAndModules)) {
-        usedModules.add(usedModule.getModule());
+        usedModules.add(usedModule.getModule(usePic));
       }
       usedHeadersAndModules.addAll(usedModules);
     }
@@ -464,19 +462,7 @@
 
   @Override
   public Map<Artifact, Artifact> getLegalGeneratedScannerFileMap() {
-    Map<Artifact, Artifact> legalOuts = new HashMap<>();
-
-    for (Artifact a : ccCompilationContext.getDeclaredIncludeSrcs()) {
-      if (!a.isSourceArtifact()) {
-        legalOuts.put(a, null);
-      }
-    }
-    for (PregreppedHeader pregreppedSrcs : ccCompilationContext.getPregreppedHeaders()) {
-      Artifact hdr = pregreppedSrcs.originalHeader();
-      Preconditions.checkState(!hdr.isSourceArtifact(), hdr);
-      legalOuts.put(hdr, pregreppedSrcs.greppedHeader());
-    }
-    return Collections.unmodifiableMap(legalOuts);
+    return ccCompilationContext.createLegalGeneratedScannerFileMap();
   }
 
   @Override