Pull out a class to hold the data that is transferred from a compilation action

to the include scanner and slightly reshuffle code.

RELNOTES: None.
PiperOrigin-RevId: 204906167

4 files changed, 94 insertions, 71 deletions

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 492131bfdd..a9f8b406f3 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
@@ -28,6 +28,7 @@ import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.rules.cpp.CppHelper.PregreppedHeader;
+import com.google.devtools.build.lib.rules.cpp.IncludeScanner.IncludeScanningHeaderData;
 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.skylarkbuildapi.cpp.CcCompilationContextApi;
@@ -197,45 +198,46 @@ public final class CcCompilationContext implements CcCompilationContextApi {
     return headerInfo.textualHeaders;
   }
 
-  public Map<Artifact, Artifact> createLegalGeneratedScannerFileMap() {
-    Map<Artifact, Artifact> legalOuts = new HashMap<>();
+  public IncludeScanningHeaderData createIncludeScanningHeaderData(
+      boolean usePic, boolean createModularHeaders) {
+    // We'd prefer for these types to use ImmutableSet/ImmutableMap. However, constructing these is
+    // substantially more costly in a way that shows up in profiles.
+    Map<PathFragment, Artifact> pathToLegalOutputArtifact = new HashMap<>();
+    Map<Artifact, Artifact> pregreppedHeaders = new HashMap<>();
+    Set<Artifact> modularHeaders = new HashSet<>();
     for (HeaderInfo transitiveHeaderInfo : transitiveHeaderInfos) {
+      boolean isModule = createModularHeaders && transitiveHeaderInfo.getModule(usePic) != null;
       for (Artifact a : transitiveHeaderInfo.modularHeaders) {
         if (!a.isSourceArtifact()) {
-          legalOuts.put(a, null);
+          pathToLegalOutputArtifact.put(a.getExecPath(), a);
+        }
+        if (isModule) {
+          modularHeaders.add(a);
         }
       }
       for (Artifact a : transitiveHeaderInfo.textualHeaders) {
         if (!a.isSourceArtifact()) {
-          legalOuts.put(a, null);
+          pathToLegalOutputArtifact.put(a.getExecPath(), a);
         }
       }
       for (PregreppedHeader pregreppedHeader : transitiveHeaderInfo.pregreppedHeaders) {
         Artifact hdr = pregreppedHeader.originalHeader();
         Preconditions.checkState(!hdr.isSourceArtifact(), hdr);
-        legalOuts.put(hdr, pregreppedHeader.greppedHeader());
+        pregreppedHeaders.put(hdr, pregreppedHeader.greppedHeader());
       }
     }
-    return Collections.unmodifiableMap(legalOuts);
+    modularHeaders.removeAll(headerInfo.modularHeaders);
+    modularHeaders.removeAll(headerInfo.textualHeaders);
+    return new IncludeScanningHeaderData(
+        Collections.unmodifiableMap(pathToLegalOutputArtifact),
+        Collections.unmodifiableSet(modularHeaders),
+        Collections.unmodifiableMap(pregreppedHeaders));
   }
 
   public NestedSet<Artifact> getTransitiveModules(boolean usePic) {
     return usePic ? transitivePicModules : transitiveModules;
   }
 
-  public Set<Artifact> getModularHeaders(boolean usePic) {
-    Set<Artifact> result = new HashSet<>();
-    for (HeaderInfo transitiveHeaderInfo : transitiveHeaderInfos) {
-      if (transitiveHeaderInfo.getModule(usePic) != null) {
-        result.addAll(transitiveHeaderInfo.modularHeaders);
-      }
-    }
-    // Remove headers belonging to this module.
-    result.removeAll(headerInfo.modularHeaders);
-    result.removeAll(headerInfo.textualHeaders);
-    return Collections.unmodifiableSet(result);
-  }
-
   public ImmutableSet<Artifact> getUsedModules(boolean usePic, Set<Artifact> usedHeaders) {
     ImmutableSet.Builder<Artifact> result = ImmutableSet.builder();
     for (HeaderInfo transitiveHeaderInfo : transitiveHeaderInfos) {
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 474459b379..cfb597ebcd 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
@@ -58,6 +58,7 @@ 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.CcToolchainFeatures.FeatureConfiguration;
 import com.google.devtools.build.lib.rules.cpp.CppCompileActionContext.Reply;
+import com.google.devtools.build.lib.rules.cpp.IncludeScanner.IncludeScanningHeaderData;
 import com.google.devtools.build.lib.skyframe.ActionExecutionValue;
 import com.google.devtools.build.lib.util.DependencySet;
 import com.google.devtools.build.lib.util.Fingerprint;
@@ -458,17 +459,9 @@ public class CppCompileAction extends AbstractAction
     return outputFile;
   }
 
-  @Override
-  public Map<Artifact, Artifact> getLegalGeneratedScannerFileMap() {
-    return ccCompilationContext.createLegalGeneratedScannerFileMap();
-  }
-
-  @Override
-  @Nullable
-  public Set<Artifact> getModularHeaders() {
-    return useHeaderModules && cppConfiguration.getPruneCppInputDiscovery()
-        ? ccCompilationContext.getModularHeaders(usePic)
-        : null;
+  public IncludeScanningHeaderData getIncludeScanningHeaderData() {
+    return ccCompilationContext.createIncludeScanningHeaderData(
+        usePic, useHeaderModules && cppConfiguration.getPruneCppInputDiscovery());
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScannable.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScannable.java
index 5ac4690891..9eb3806989 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScannable.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScannable.java
@@ -20,8 +20,6 @@ import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Collection;
 import java.util.List;
-import java.util.Map;
-import java.util.Set;
 import javax.annotation.Nullable;
 
 /**
@@ -94,23 +92,6 @@ public interface IncludeScannable {
   NestedSet<Artifact> getDeclaredIncludeSrcs();
 
   /**
-   * Returns a map of (generated header:.includes file listing the header's includes) which may be
-   * reached during include scanning. Generated files which are reached, but not in the key set,
-   * must be ignored.
-   *
-   * <p>If grepping of output files is not enabled via --extract_generated_inclusions, keys
-   * should just map to null.
-   */
-  Map<Artifact, Artifact> getLegalGeneratedScannerFileMap();
-
-  /**
-   * Returns a set of headers that are modular, i.e. are going to be read as a serialized AST rather
-   * than from the textual source file. Depending on the implementation, it is likely that further
-   * input discovery through such headers is unnecessary as the serialized AST is self-contained.
-   */
-  Set<Artifact> getModularHeaders();
-
-  /**
    * Returns an artifact that is the executable for {@link ExtractInclusionAction}.
    */
   Artifact getGrepIncludes();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScanner.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScanner.java
index 0512adcc4f..5ceb0a580d 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScanner.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/IncludeScanner.java
@@ -64,11 +64,14 @@ public interface IncludeScanner {
    * <p>{@code mainSource} is the source file relative to which the {@code cmdlineIncludes} are
    * interpreted.</p>
    */
-  void process(Artifact mainSource, Collection<Artifact> sources,
-      Map<Artifact, Artifact> legalOutputPaths, List<PathFragment> includeDirs,
-      List<PathFragment> quoteIncludeDirs, List<String> cmdlineIncludes,
-      Set<Artifact> includes, ActionExecutionContext actionExecutionContext, Artifact grepIncludes,
-      Set<Artifact> modularHeaders)
+  void process(
+      Artifact mainSource,
+      Collection<Artifact> sources,
+      IncludeScanningHeaderData includeScanningHeaderData,
+      List<String> cmdlineIncludes,
+      Set<Artifact> includes,
+      ActionExecutionContext actionExecutionContext,
+      Artifact grepIncludes)
       throws IOException, ExecException, InterruptedException;
 
   /** Supplies IncludeScanners upon request. */
@@ -97,8 +100,10 @@ public interface IncludeScanner {
      * @param actionExecutionContext the context for {@code action}.
      * @param profilerTaskName what the {@link Profiler} should record this call for.
      */
-    public static Collection<Artifact> scanForIncludedInputs(IncludeScannable action,
+    public static Collection<Artifact> scanForIncludedInputs(
+        IncludeScannable action,
         IncludeScannerSupplier includeScannerSupplier,
+        IncludeScanningHeaderData includeScanningHeaderData,
         ActionExecutionContext actionExecutionContext,
         String profilerTaskName)
         throws ExecException, InterruptedException {
@@ -110,8 +115,6 @@ public interface IncludeScanner {
 
       Profiler profiler = Profiler.instance();
       try (SilentCloseable c = profiler.profile(ProfilerTask.SCANNER, profilerTaskName)) {
-
-        Map<Artifact, Artifact> legalOutputPaths = action.getLegalGeneratedScannerFileMap();
         // Deduplicate include directories. This can occur especially with "built-in" and "system"
         // include directories because of the way we retrieve them. Duplicate include directories
         // really mess up #include_next directives.
@@ -121,31 +124,28 @@ public interface IncludeScanner {
 
         includeDirs.addAll(action.getSystemIncludeDirs());
 
-          // Add the system include paths to the list of include paths.
-          for (PathFragment pathFragment : action.getBuiltInIncludeDirectories()) {
-            if (pathFragment.isAbsolute()) {
-              absoluteBuiltInIncludeDirs.add(pathFragment);
-            }
-            includeDirs.add(pathFragment);
+        // Add the system include paths to the list of include paths.
+        for (PathFragment pathFragment : action.getBuiltInIncludeDirectories()) {
+          if (pathFragment.isAbsolute()) {
+            absoluteBuiltInIncludeDirs.add(pathFragment);
           }
+          includeDirs.add(pathFragment);
+        }
 
-          List<PathFragment> includeDirList = ImmutableList.copyOf(includeDirs);
-          IncludeScanner scanner = includeScannerSupplier.scannerFor(quoteIncludeDirs,
-              includeDirList);
+        List<PathFragment> includeDirList = ImmutableList.copyOf(includeDirs);
+        IncludeScanner scanner =
+            includeScannerSupplier.scannerFor(quoteIncludeDirs, includeDirList);
 
         Artifact mainSource = action.getMainIncludeScannerSource();
         Collection<Artifact> sources = action.getIncludeScannerSources();
         scanner.process(
             mainSource,
             sources,
-            legalOutputPaths,
-            quoteIncludeDirs,
-            includeDirList,
+            includeScanningHeaderData,
             cmdlineIncludes,
             includes,
             actionExecutionContext,
-            action.getGrepIncludes(),
-            action.getModularHeaders());
+            action.getGrepIncludes());
 
       } catch (IOException e) {
         throw new EnvironmentalExecException(e.getMessage());
@@ -172,4 +172,51 @@ public interface IncludeScanner {
       return inputs;
     }
   }
+
+  /**
+   * Holds pre-aggregated information that the {@link IncludeScanner} needs from the compilation
+   * action.
+   */
+  class IncludeScanningHeaderData {
+    /**
+     * Lookup table to find the {@link Artifact}s of generated files based on their {@link
+     * Artifact#execPath}.
+     */
+    private final Map<PathFragment, Artifact> pathToLegalOutputArtifact;
+
+    /**
+     * The set of headers that are modular, i.e. are going to be read as a serialized AST rather
+     * than from the textual source file. Depending on the implementation, it is likely that further
+     * input discovery through such headers is unnecessary as the serialized AST is self-contained.
+     */
+    private final Set<Artifact> modularHeaders;
+
+    /**
+     * Lookup table to find pregrepped .includes files that contain a single line per include. These
+     * are essentially a cache to make it unnecessary for the {@link IncludeScanner} to read large
+     * generated files repeatedly.
+     */
+    private final Map<Artifact, Artifact> pregreppedHeaders;
+
+    public IncludeScanningHeaderData(
+        Map<PathFragment, Artifact> pathToLegalOutputArtifact,
+        Set<Artifact> modularHeaders,
+        Map<Artifact, Artifact> pregreppedHeaders) {
+      this.pathToLegalOutputArtifact = pathToLegalOutputArtifact;
+      this.modularHeaders = modularHeaders;
+      this.pregreppedHeaders = pregreppedHeaders;
+    }
+
+    public Set<Artifact> getModularHeaders() {
+      return modularHeaders;
+    }
+
+    public Map<Artifact, Artifact> getPregreppedHeaders() {
+      return pregreppedHeaders;
+    }
+
+    public Map<PathFragment, Artifact> getPathToLegalOutputArtifact() {
+      return pathToLegalOutputArtifact;
+    }
+  }
 }