Batch all DependencyResolver#getTarget calls. This leads to some duplicate iteration, but that should be cheap, while requesting packages sequentially can hurt...

PiperOrigin-RevId: 208126130

2 files changed, 20 insertions, 13 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/LabelVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/LabelVisitor.java
index ab1f1eaaa3..af06cde70e 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/LabelVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/LabelVisitor.java
@@ -26,7 +26,6 @@ import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.packages.AggregatingAttributeMapper;
 import com.google.devtools.build.lib.packages.AspectDefinition;
 import com.google.devtools.build.lib.packages.Attribute;
-import com.google.devtools.build.lib.packages.AttributeMap;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.DependencyFilter;
 import com.google.devtools.build.lib.packages.InputFile;
@@ -353,15 +352,13 @@ final class LabelVisitor {
     private void visitRule(final Rule rule, final int depth, final int count)
         throws InterruptedException {
       // Follow all labels defined by this rule:
-      AggregatingAttributeMapper.of(rule).visitLabels(new AttributeMap.AcceptsLabelAttribute() {
-        @Override
-        public void acceptLabelAttribute(Label label, Attribute attribute) {
-          if (!edgeFilter.apply(rule, attribute)) {
-            return;
-          }
-          enqueueTarget(rule, attribute, label, depth, count);
-        }
-      });
+      AggregatingAttributeMapper.of(rule)
+          .visitLabels()
+          .stream()
+          .filter(depEdge -> edgeFilter.apply(rule, depEdge.getAttribute()))
+          .forEach(
+              depEdge ->
+                  enqueueTarget(rule, depEdge.getAttribute(), depEdge.getLabel(), depth, count));
     }
 
     @ThreadSafe
diff --git a/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java
index 27685267e8..81298a46f7 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java
@@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Streams;
 import com.google.devtools.build.lib.analysis.AspectCollection;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.Dependency;
@@ -58,6 +59,8 @@ import java.io.OutputStream;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
@@ -242,9 +245,16 @@ public class TransitionsOutputFormatterCallback extends CqueryThreadsafeCallback
     }
 
     @Override
-    protected Target getTarget(Target from, Label label, NestedSetBuilder<Cause> rootCauses)
-        throws InterruptedException {
-      return partialResultMap.get(label);
+    protected Map<Label, Target> getTargets(
+        Iterable<Label> labels,
+        Target fromTarget,
+        NestedSetBuilder<Cause> rootCauses,
+        int labelsSizeHint) {
+      return Streams.stream(labels)
+          .distinct()
+          .filter(Objects::nonNull)
+          .filter(partialResultMap::containsKey)
+          .collect(Collectors.toMap(Function.identity(), partialResultMap::get));
     }
 
     @Override