Use NestedSets to store topLevelModules and discoveredModules. Also improve

computeTransitivelyUsedModules to return early on missing values instead of
starting to create datastructures.

RELNOTES: None.
PiperOrigin-RevId: 208351193

3 files changed, 60 insertions, 34 deletions

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 7a49ba025c..f4b9c688d2 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
@@ -74,6 +74,7 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -145,11 +146,6 @@ public class CppCompileAction extends AbstractAction
    */
   private Set<Artifact> usedModules = null;
 
-  /** Used modules that are not transitively used through other topLevelModules. */
-  private Iterable<Artifact> topLevelModules = null;
-
-  private CcToolchainVariables overwrittenVariables = null;
-
   /**
    * This field is set only for C++ module compiles (compiling .cppmap files into .pcm files). It
    * stores the modules necessary for building this module as they will later also be required for
@@ -159,7 +155,12 @@ public class CppCompileAction extends AbstractAction
    * <p>This field is populated either based on the discovered headers in {@link #discoverInputs} or
    * extracted from the action inputs when restoring it from the action cache.
    */
-  private ImmutableList<Artifact> discoveredModules = null;
+  private NestedSet<Artifact> discoveredModules = null;
+
+  /** Used modules that are not transitively used through other topLevelModules. */
+  private NestedSet<Artifact> topLevelModules = null;
+
+  private CcToolchainVariables overwrittenVariables = null;
 
   /**
    * Creates a new action to compile C/C++ source files.
@@ -453,17 +454,38 @@ public class CppCompileAction extends AbstractAction
       usedModules =
           ccCompilationContext.getUsedModules(usePic, ImmutableSet.copyOf(additionalInputs));
     }
-    Set<Artifact> transitivelyUsedModules =
+    Map<Artifact, NestedSet<Artifact>> transitivelyUsedModules =
         computeTransitivelyUsedModules(
             actionExecutionContext.getEnvironmentForDiscoveringInputs(), usedModules);
     if (transitivelyUsedModules == null) {
       return null;
     }
 
-    ImmutableList<Artifact> discoveredModules =
-        ImmutableList.copyOf(Sets.union(usedModules, transitivelyUsedModules));
-    topLevelModules = ImmutableList.copyOf(Sets.difference(usedModules, transitivelyUsedModules));
+    // Compute top-level modules, i.e. used modules that aren't also dependencies of other
+    // used modules. Combining the NestedSets of transitive deps of the top-level modules also
+    // gives us an effective way to compute and store discoveredModules.
+    Set<Artifact> topLevel = new LinkedHashSet<>(usedModules);
     usedModules = null;
+    for (NestedSet<Artifact> transitive : transitivelyUsedModules.values()) {
+      // It is better to iterate over each nested set here instead of creating a joint one and
+      // iterating over it, as this makes use of NestedSet's memoization (each of them has likely
+      // been iterated over before). Don't use Set.removeAll() here as that iterates over the
+      // smaller set (topLevel, which would support efficient lookup) and looks up in the larger one
+      // (transitive, which is a linear scan).
+      for (Artifact module : transitive) {
+        topLevel.remove(module);
+      }
+    }
+    NestedSetBuilder<Artifact> topLevelModulesBuilder = NestedSetBuilder.stableOrder();
+    NestedSetBuilder<Artifact> discoveredModulesBuilder = NestedSetBuilder.stableOrder();
+    for (Artifact module : topLevel) {
+      topLevelModulesBuilder.add(module);
+      discoveredModulesBuilder.addTransitive(transitivelyUsedModules.get(module));
+    }
+    topLevelModules = topLevelModulesBuilder.build();
+    discoveredModulesBuilder.addTransitive(topLevelModules);
+    NestedSet<Artifact> discoveredModules = discoveredModulesBuilder.build();
+
     additionalInputs = Iterables.concat(additionalInputs, discoveredModules);
     if (outputFile.isFileType(CppFileTypes.CPP_MODULE)) {
       this.discoveredModules = discoveredModules;
@@ -512,7 +534,7 @@ public class CppCompileAction extends AbstractAction
    */
   @Override
   @Nullable
-  public ImmutableList<Artifact> getDiscoveredModules() {
+  public NestedSet<Artifact> getDiscoveredModules() {
     return discoveredModules;
   }
 
@@ -916,7 +938,8 @@ public class CppCompileAction extends AbstractAction
     super.updateInputs(inputs);
     if (outputFile.isFileType(CppFileTypes.CPP_MODULE)) {
       discoveredModules =
-          ImmutableList.copyOf(
+          NestedSetBuilder.wrap(
+              Order.STABLE_ORDER,
               Iterables.filter(inputs, input -> input.isFileType(CppFileTypes.CPP_MODULE)));
     }
   }
@@ -1255,38 +1278,41 @@ public class CppCompileAction extends AbstractAction
   /**
    * For the given {@code usedModules}, looks up modules discovered by their generating actions.
    *
-   * <p>The returned value only contains elements of {@code usedModules} if they happen to be
-   * transitively used from other elements of {@code usedModules}. It can be null when skyframe
-   * lookups return null.
+   * <p>The returned value only contains a map from elements of {@code usedModules} to the
+   * {@link #discoveredModules} required to use them. If dependent actions have not been executed
+   * yet (and thus {@link #discoveredModules} aren't known yet, returns null.
    */
   @Nullable
-  private static Set<Artifact> computeTransitivelyUsedModules(
+  private static Map<Artifact, NestedSet<Artifact>> computeTransitivelyUsedModules(
       SkyFunction.Environment env, Set<Artifact> usedModules) throws InterruptedException {
     // ActionLookupKey → ActionLookupValue
     Map<SkyKey, SkyValue> actionLookupValues =
         env.getValues(
             Iterables.transform(
                 usedModules, module -> (ActionLookupKey) module.getArtifactOwner()));
+    if (env.valuesMissing()) {
+      return null;
+    }
     ArrayList<ActionLookupData> executionValueLookups = new ArrayList<>(usedModules.size());
     for (Artifact module : usedModules) {
       ActionLookupData lookupData = lookupDataFromModule(actionLookupValues, module);
-      if (lookupData == null) {
-        return null;
-      }
-      executionValueLookups.add(lookupData);
+      executionValueLookups.add(Preconditions.checkNotNull(lookupData, module));
     }
 
-    Set<Artifact> additionalModules = Sets.newLinkedHashSet();
     // ActionLookupData → ActionExecutionValue
     Map<SkyKey, SkyValue> actionExecutionValues = env.getValues(executionValueLookups);
-    for (ActionLookupData lookup : executionValueLookups) {
+    if (env.valuesMissing()) {
+      return null;
+    }
+    ImmutableMap.Builder<Artifact, NestedSet<Artifact>> transitivelyUsedModules =
+        ImmutableMap.builderWithExpectedSize(usedModules.size());
+    int pos = 0;
+    for (Artifact module : usedModules) {
+      ActionLookupData lookup = executionValueLookups.get(pos++);
       ActionExecutionValue value = (ActionExecutionValue) actionExecutionValues.get(lookup);
-      if (value == null) {
-        return null;
-      }
-      additionalModules.addAll(value.getDiscoveredModules());
+      transitivelyUsedModules.put(module, value.getDiscoveredModules());
     }
-    return additionalModules;
+    return transitivelyUsedModules.build();
   }
 
   @Nullable
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 75d5e53ff0..1e3b19f87b 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
@@ -14,7 +14,6 @@
 
 package com.google.devtools.build.lib.rules.cpp;
 
-import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -97,5 +96,5 @@ public interface IncludeScannable {
   Artifact getGrepIncludes();
 
   /** Returns modules necessary for building and using the output of this action. */
-  ImmutableList<Artifact> getDiscoveredModules();
+  NestedSet<Artifact> getDiscoveredModules();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
index d6c258f8de..4293dddf6f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
@@ -31,6 +31,7 @@ import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.FileStateValue;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.actions.FilesetOutputSymlink;
+import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.skyframe.serialization.UnshareableValue;
@@ -79,7 +80,7 @@ public class ActionExecutionValue implements SkyValue {
 
   @Nullable private final ImmutableList<FilesetOutputSymlink> outputSymlinks;
 
-  @Nullable private final ImmutableList<Artifact> discoveredModules;
+  @Nullable private final NestedSet<Artifact> discoveredModules;
 
   /**
    * @param artifactData Map from Artifacts to corresponding FileValues.
@@ -97,7 +98,7 @@ public class ActionExecutionValue implements SkyValue {
       Map<Artifact, TreeArtifactValue> treeArtifactData,
       Map<Artifact, FileArtifactValue> additionalOutputData,
       @Nullable ImmutableList<FilesetOutputSymlink> outputSymlinks,
-      @Nullable ImmutableList<Artifact> discoveredModules) {
+      @Nullable NestedSet<Artifact> discoveredModules) {
     this.artifactData = ImmutableMap.<Artifact, FileValue>copyOf(artifactData);
     this.additionalOutputData = ImmutableMap.copyOf(additionalOutputData);
     this.treeArtifactData = ImmutableMap.copyOf(treeArtifactData);
@@ -110,7 +111,7 @@ public class ActionExecutionValue implements SkyValue {
       Map<Artifact, TreeArtifactValue> treeArtifactData,
       Map<Artifact, FileArtifactValue> additionalOutputData,
       @Nullable ImmutableList<FilesetOutputSymlink> outputSymlinks,
-      @Nullable ImmutableList<Artifact> discoveredModules,
+      @Nullable NestedSet<Artifact> discoveredModules,
       boolean notifyOnActionCacheHitAction) {
     return notifyOnActionCacheHitAction
         ? new CrossServerUnshareableActionExecutionValue(
@@ -172,7 +173,7 @@ public class ActionExecutionValue implements SkyValue {
   }
 
   @Nullable
-  public ImmutableList<Artifact> getDiscoveredModules() {
+  public NestedSet<Artifact> getDiscoveredModules() {
     return discoveredModules;
   }
 
@@ -246,7 +247,7 @@ public class ActionExecutionValue implements SkyValue {
         Map<Artifact, TreeArtifactValue> treeArtifactData,
         Map<Artifact, FileArtifactValue> additionalOutputData,
         @Nullable ImmutableList<FilesetOutputSymlink> outputSymlinks,
-        @Nullable ImmutableList<Artifact> discoveredModules) {
+        @Nullable NestedSet<Artifact> discoveredModules) {
       super(
           artifactData, treeArtifactData, additionalOutputData, outputSymlinks, discoveredModules);
     }