1 files changed, 60 insertions, 58 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
index a558e4581d..b423803744 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
@@ -16,7 +16,10 @@ package com.google.devtools.build.lib.query2.engine;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.MapMaker;
+import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.collect.CompactHashSet;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskCallable;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskFuture;
 import java.util.Collections;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
@@ -29,17 +32,18 @@ public final class QueryUtil {
 
   /** A {@link Callback} that can aggregate all the partial results into one set. */
   public interface AggregateAllCallback<T> extends Callback<T> {
+    /** Returns a (mutable) set of all the results. */
     Set<T> getResult();
   }
 
-  /** A {@link OutputFormatterCallback} that can aggregate all the partial results into one set. */
+  /** A {@link OutputFormatterCallback} that is also a {@link AggregateAllCallback}. */
   public abstract static class AggregateAllOutputFormatterCallback<T>
-      extends OutputFormatterCallback<T> implements AggregateAllCallback<T>  {
+      extends ThreadSafeOutputFormatterCallback<T> implements AggregateAllCallback<T>  {
   }
 
   private static class AggregateAllOutputFormatterCallbackImpl<T>
       extends AggregateAllOutputFormatterCallback<T> {
-    private final Set<T> result = CompactHashSet.create();
+    private final Set<T> result = Sets.newConcurrentHashSet();
 
     @Override
     public final void processOutput(Iterable<T> partialResult) {
@@ -52,95 +56,93 @@ public final class QueryUtil {
     }
   }
 
+  private static class OrderedAggregateAllOutputFormatterCallbackImpl<T>
+      extends AggregateAllOutputFormatterCallback<T> {
+    private final Set<T> result = CompactHashSet.create();
+
+    @Override
+    public final synchronized void processOutput(Iterable<T> partialResult) {
+      Iterables.addAll(result, partialResult);
+    }
+
+    @Override
+    public synchronized Set<T> getResult() {
+      return result;
+    }
+  }
+
   /**
-   * Returns a fresh {@link AggregateAllOutputFormatterCallback} that can aggregate all the partial
-   * results into one set.
-   *
-   * <p>Intended to be used by top-level evaluation of {@link QueryExpression}s; contrast with
-   * {@link #newAggregateAllCallback}.
+   * Returns a fresh {@link AggregateAllOutputFormatterCallback} instance whose
+   * {@link AggregateAllCallback#getResult} returns all the elements of the result in the order they
+   * were processed.
    */
   public static <T> AggregateAllOutputFormatterCallback<T>
-      newAggregateAllOutputFormatterCallback() {
-    return new AggregateAllOutputFormatterCallbackImpl<>();
+      newOrderedAggregateAllOutputFormatterCallback() {
+    return new OrderedAggregateAllOutputFormatterCallbackImpl<>();
   }
 
-  /**
-   * Returns a fresh {@link AggregateAllCallback}.
-   *
-   * <p>Intended to be used by {@link QueryExpression} implementations; contrast with
-   * {@link #newAggregateAllOutputFormatterCallback}.
-   */
+  /** Returns a fresh {@link AggregateAllCallback} instance. */
   public static <T> AggregateAllCallback<T> newAggregateAllCallback() {
     return new AggregateAllOutputFormatterCallbackImpl<>();
   }
 
   /**
-   * Fully evaluate a {@code QueryExpression} and return a set with all the results.
+   * Returns a {@link QueryTaskFuture} representing the evaluation of {@code expr} as a (mutable)
+   * {@link Set} comprised of all the results.
    *
    * <p>Should only be used by QueryExpressions when it is the only way of achieving correctness.
    */
-  public static <T> Set<T> evalAll(
-      QueryEnvironment<T> env, VariableContext<T> context, QueryExpression expr)
-          throws QueryException, InterruptedException {
-    AggregateAllCallback<T> callback = newAggregateAllCallback();
-    env.eval(expr, context, callback);
-    return callback.getResult();
-  }
-
-  private abstract static class AbstractUniquifierBase<T, K> implements Uniquifier<T> {
-    @Override
-    public final ImmutableList<T> unique(Iterable<T> newElements) {
-      ImmutableList.Builder<T> result = ImmutableList.builder();
-      for (T element : newElements) {
-        if (unique(element)) {
-          result.add(element);
-        }
-      }
-      return result.build();
-    }
+  public static <T> QueryTaskFuture<Set<T>> evalAll(
+      QueryEnvironment<T> env, VariableContext<T> context, QueryExpression expr) {
+    final AggregateAllCallback<T> callback = newAggregateAllCallback();
+    return env.whenSucceedsCall(
+        env.eval(expr, context, callback),
+        new QueryTaskCallable<Set<T>>() {
+          @Override
+          public Set<T> call() {
+            return callback.getResult();
+          }
+        });
   }
 
   /** A trivial {@link Uniquifier} implementation. */
-  public static class UniquifierImpl<T, K> extends AbstractUniquifierBase<T, K> {
+  public static class UniquifierImpl<T, K> implements Uniquifier<T> {
     private final KeyExtractor<T, K> extractor;
-    private final CompactHashSet<K> alreadySeen = CompactHashSet.create();
+    private final Set<K> alreadySeen;
 
     public UniquifierImpl(KeyExtractor<T, K> extractor) {
-      this.extractor = extractor;
+      this(extractor, /*concurrencyLevel=*/ 1);
     }
 
-    @Override
-    public final boolean unique(T element) {
-      return alreadySeen.add(extractor.extractKey(element));
-    }
-  }
-
-  /** A trvial {@link ThreadSafeUniquifier} implementation. */
-  public static class ThreadSafeUniquifierImpl<T, K>
-      extends AbstractUniquifierBase<T, K> implements ThreadSafeUniquifier<T> {
-    private final KeyExtractor<T, K> extractor;
-    private final Set<K> alreadySeen;
-
-    public ThreadSafeUniquifierImpl(KeyExtractor<T, K> extractor, int concurrencyLevel) {
+    public UniquifierImpl(KeyExtractor<T, K> extractor, int concurrencyLevel) {
       this.extractor = extractor;
       this.alreadySeen = Collections.newSetFromMap(
           new MapMaker().concurrencyLevel(concurrencyLevel).<K, Boolean>makeMap());
     }
 
     @Override
-    public final boolean unique(T element) {
+    public boolean unique(T element) {
       return alreadySeen.add(extractor.extractKey(element));
     }
+
+    @Override
+    public ImmutableList<T> unique(Iterable<T> newElements) {
+      ImmutableList.Builder<T> result = ImmutableList.builder();
+      for (T element : newElements) {
+        if (unique(element)) {
+          result.add(element);
+        }
+      }
+      return result.build();
+    }
   }
 
-  /** A trivial {@link ThreadSafeMinDepthUniquifier} implementation. */
-  public static class ThreadSafeMinDepthUniquifierImpl<T, K>
-      implements ThreadSafeMinDepthUniquifier<T> {
+  /** A trivial {@link MinDepthUniquifier} implementation. */
+  public static class MinDepthUniquifierImpl<T, K> implements MinDepthUniquifier<T> {
     private final KeyExtractor<T, K> extractor;
     private final ConcurrentMap<K, AtomicInteger> alreadySeenAtDepth;
 
-    public ThreadSafeMinDepthUniquifierImpl(
-        KeyExtractor<T, K> extractor, int concurrencyLevel) {
+    public MinDepthUniquifierImpl(KeyExtractor<T, K> extractor, int concurrencyLevel) {
       this.extractor = extractor;
       this.alreadySeenAtDepth = new MapMaker().concurrencyLevel(concurrencyLevel).makeMap();
     }