Split SkyQueryEnvironment#getReverseDeps into a few helper methods and use them in ParallelSkyQueryUtils.

In AllRdepsUnboundedVisitor#getVisitResult, only get rdeps of the nodes in question. As long as we're careful to only visit nodes have corresponding Targets (which we already are doing), there's no need to also get the Targets corresponding to the nodes in question.

--
MOS_MIGRATED_REVID=135680956

1 files changed, 36 insertions, 14 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
index b920787f2b..000261cc1c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
@@ -345,15 +345,15 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     return super.evaluateQuery(expr, batchCallback);
   }
 
-  private Map<Target, Collection<Target>> makeTargetsMap(Map<SkyKey, Iterable<SkyKey>> input)
+  private Map<SkyKey, Collection<Target>> targetifyValues(Map<SkyKey, Iterable<SkyKey>> input)
       throws InterruptedException {
-    ImmutableMap.Builder<Target, Collection<Target>> result = ImmutableMap.builder();
+    ImmutableMap.Builder<SkyKey, Collection<Target>> result = ImmutableMap.builder();
 
     Map<SkyKey, Target> allTargets =
         makeTargetsFromSkyKeys(Sets.newHashSet(Iterables.concat(input.values())));
 
-    for (Map.Entry<SkyKey, Target> entry : makeTargetsFromSkyKeys(input.keySet()).entrySet()) {
-      Iterable<SkyKey> skyKeys = input.get(entry.getKey());
+    for (Map.Entry<SkyKey, Iterable<SkyKey>> entry : input.entrySet()) {
+      Iterable<SkyKey> skyKeys = entry.getValue();
       Set<Target> targets = CompactHashSet.createWithExpectedSize(Iterables.size(skyKeys));
       for (SkyKey key : skyKeys) {
         Target target = allTargets.get(key);
@@ -361,19 +361,38 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
           targets.add(target);
         }
       }
-      result.put(entry.getValue(), targets);
+      result.put(entry.getKey(), targets);
     }
     return result.build();
   }
 
-  private Map<Target, Collection<Target>> getRawFwdDeps(Iterable<Target> targets)
+  private Map<Target, Collection<Target>> targetifyKeys(Map<SkyKey, Collection<Target>> input)
       throws InterruptedException {
-    return makeTargetsMap(graph.getDirectDeps(makeTransitiveTraversalKeys(targets)));
+    Map<SkyKey, Target> targets = makeTargetsFromSkyKeys(input.keySet());
+    ImmutableMap.Builder<Target, Collection<Target>> resultBuilder = ImmutableMap.builder();
+    for (Map.Entry<SkyKey, Collection<Target>> entry : input.entrySet()) {
+      SkyKey key = entry.getKey();
+      Target target = targets.get(key);
+      if (target != null) {
+        resultBuilder.put(target, entry.getValue());
+      }
+    }
+    return resultBuilder.build();
+  }
+
+  private Map<Target, Collection<Target>> targetifyKeysAndValues(
+      Map<SkyKey, Iterable<SkyKey>> input) throws InterruptedException {
+    return targetifyKeys(targetifyValues(input));
   }
 
-  private Map<Target, Collection<Target>> getRawReverseDeps(Iterable<Target> targets)
+  private Map<Target, Collection<Target>> getRawFwdDeps(Iterable<Target> targets)
       throws InterruptedException {
-    return makeTargetsMap(graph.getReverseDeps(makeTransitiveTraversalKeys(targets)));
+    return targetifyKeysAndValues(graph.getDirectDeps(makeTransitiveTraversalKeys(targets)));
+  }
+
+  private Map<SkyKey, Collection<Target>> getRawReverseDeps(
+      Iterable<SkyKey> transitiveTraversalKeys) throws InterruptedException {
+    return targetifyValues(graph.getReverseDeps(transitiveTraversalKeys));
   }
 
   private Set<Label> getAllowedDeps(Rule rule) throws InterruptedException {
@@ -421,20 +440,23 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
 
   @Override
   public Collection<Target> getReverseDeps(Iterable<Target> targets) throws InterruptedException {
-    Map<Target, Collection<Target>> rawReverseDeps = getRawReverseDeps(targets);
-    warnIfMissingTargets(targets, rawReverseDeps.keySet());
+    return getReverseDepsOfTransitiveTraversalKeys(Iterables.transform(targets, TARGET_TO_SKY_KEY));
+  }
 
+  Collection<Target> getReverseDepsOfTransitiveTraversalKeys(
+      Iterable<SkyKey> transitiveTraversalKeys) throws InterruptedException {
+    Map<SkyKey, Collection<Target>> rawReverseDeps = getRawReverseDeps(transitiveTraversalKeys);
     return processRawReverseDeps(rawReverseDeps);
   }
 
-  private Collection<Target> processRawReverseDeps(Map<Target, Collection<Target>> rawReverseDeps)
+  private Collection<Target> processRawReverseDeps(Map<SkyKey, Collection<Target>> rawReverseDeps)
       throws InterruptedException {
     Set<Target> result = CompactHashSet.create();
     CompactHashSet<Target> visited =
         CompactHashSet.createWithExpectedSize(totalSizeOfCollections(rawReverseDeps.values()));
 
     Set<Label> keys = CompactHashSet.create(Collections2.transform(rawReverseDeps.keySet(),
-        TARGET_LABEL_FUNCTION));
+        SKYKEY_TO_LABEL));
     for (Collection<Target> parentCollection : rawReverseDeps.values()) {
       for (Target parent : parentCollection) {
         if (visited.add(parent)) {
@@ -1211,7 +1233,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
         Callback<Target> callback,
         Queue<Map.Entry<SkyKey, Iterable<SkyKey>>> reverseDepsQueue)
         throws QueryException, InterruptedException {
-      Collection<Target> children = processRawReverseDeps(makeTargetsMap(reverseDepsMap));
+      Collection<Target> children = processRawReverseDeps(targetifyValues(reverseDepsMap));
       Iterable<Target> currentInUniverse = Iterables.filter(children, universe);
       ImmutableList<Target> uniqueChildren = uniquifier.unique(currentInUniverse);
       reverseDepsMap.clear();