Stream output of rbuildfiles operator to a callback. This function was missed when we transformed to callbacks.

--
MOS_MIGRATED_REVID=121619124

2 files changed, 41 insertions, 14 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java b/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java
index 624573c132..b3a64c6169 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java
@@ -17,6 +17,7 @@ import com.google.common.base.Function;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
@@ -27,7 +28,6 @@ import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.vfs.PathFragment;
 
 import java.util.List;
-import java.util.Set;
 
 /**
  * An "rbuildfiles" query expression, which computes the set of packages (as represented by their
@@ -76,8 +76,8 @@ public class RBuildFilesFunction implements QueryFunction {
     for (Argument arg : args) {
       fileNames.add(arg.getWord());
     }
-    callback.process((Set<T>)
-        ((SkyQueryEnvironment) env)
-            .getRBuildFiles(Collections2.transform(args, ARGUMENT_TO_PATH_FRAGMENT)));
+    ((SkyQueryEnvironment) env)
+        .getRBuildFiles(
+            Collections2.transform(args, ARGUMENT_TO_PATH_FRAGMENT), (Callback<Target>) callback);
   }
 }
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 349661d8dd..c127f42570 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
@@ -864,12 +864,42 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     return result;
   }
 
+  private static final Function<SkyValue, Package> EXTRACT_PACKAGE =
+      new Function<SkyValue, Package>() {
+        @Override
+        public Package apply(SkyValue skyValue) {
+          return ((PackageValue) skyValue).getPackage();
+        }
+      };
+
+  private static final Predicate<Package> ERROR_FREE_PACKAGE =
+      new Predicate<Package>() {
+        @Override
+        public boolean apply(Package pkg) {
+          return !pkg.containsErrors();
+        }
+      };
+
+  private static final Function<Package, Target> GET_BUILD_FILE =
+      new Function<Package, Target>() {
+        @Override
+        public Target apply(Package pkg) {
+          return pkg.getBuildFile();
+        }
+      };
+
+  private static Iterable<Target> getBuildFilesForPackageValues(Iterable<SkyValue> packageValues) {
+    return Iterables.transform(
+        Iterables.filter(Iterables.transform(packageValues, EXTRACT_PACKAGE), ERROR_FREE_PACKAGE),
+        GET_BUILD_FILE);
+  }
+
   /**
    * Calculates the set of {@link Package} objects, represented as source file targets, that depend
    * on the given list of BUILD files and subincludes (other files are filtered out).
    */
-  @Nullable
-  Set<Target> getRBuildFiles(Collection<PathFragment> fileIdentifiers) {
+  void getRBuildFiles(Collection<PathFragment> fileIdentifiers, Callback<Target> callback)
+      throws QueryException, InterruptedException {
     Collection<SkyKey> files = getSkyKeysForFileFragments(fileIdentifiers);
     Collection<SkyKey> current = graph.getSuccessfulValues(files).keySet();
     Set<SkyKey> resultKeys = CompactHashSet.create();
@@ -889,16 +919,13 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
           current.add(rdep);
         }
       }
-    }
-    Map<SkyKey, SkyValue> packageValues = graph.getSuccessfulValues(resultKeys);
-    ImmutableSet.Builder<Target> result = ImmutableSet.builder();
-    for (SkyValue value : packageValues.values()) {
-      Package pkg = ((PackageValue) value).getPackage();
-      if (!pkg.containsErrors()) {
-        result.add(pkg.getBuildFile());
+      if (resultKeys.size() >= BATCH_CALLBACK_SIZE) {
+        callback.process(
+            getBuildFilesForPackageValues(graph.getSuccessfulValues(resultKeys).values()));
+        resultKeys.clear();
       }
     }
-    return result.build();
+    callback.process(getBuildFilesForPackageValues(graph.getSuccessfulValues(resultKeys).values()));
   }
 
   @Override