1 files changed, 132 insertions, 23 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 b423803744..73dd930ac3 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,14 +16,21 @@ 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.MutableMap;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskCallable;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskFuture;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ThreadSafeMutableSet;
+import java.util.AbstractSet;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.Nullable;
 
 /** Several query utilities to make easier to work with query callbacks and uniquifiers. */
 public final class QueryUtil {
@@ -31,19 +38,23 @@ public final class QueryUtil {
   private 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();
+  public interface AggregateAllCallback<T, S extends Set<T>> extends Callback<T> {
+    /** Returns a {@link Set} of all the results. */
+    S getResult();
   }
 
   /** A {@link OutputFormatterCallback} that is also a {@link AggregateAllCallback}. */
-  public abstract static class AggregateAllOutputFormatterCallback<T>
-      extends ThreadSafeOutputFormatterCallback<T> implements AggregateAllCallback<T>  {
+  public abstract static class AggregateAllOutputFormatterCallback<T, S extends Set<T>>
+      extends ThreadSafeOutputFormatterCallback<T> implements AggregateAllCallback<T, S>  {
   }
 
   private static class AggregateAllOutputFormatterCallbackImpl<T>
-      extends AggregateAllOutputFormatterCallback<T> {
-    private final Set<T> result = Sets.newConcurrentHashSet();
+      extends AggregateAllOutputFormatterCallback<T, ThreadSafeMutableSet<T>> {
+    private final ThreadSafeMutableSet<T> result;
+
+    private AggregateAllOutputFormatterCallbackImpl(QueryEnvironment<T> env) {
+      this.result = env.createThreadSafeMutableSet();
+    }
 
     @Override
     public final void processOutput(Iterable<T> partialResult) {
@@ -51,22 +62,35 @@ public final class QueryUtil {
     }
 
     @Override
-    public Set<T> getResult() {
+    public ThreadSafeMutableSet<T> getResult() {
       return result;
     }
   }
 
   private static class OrderedAggregateAllOutputFormatterCallbackImpl<T>
-      extends AggregateAllOutputFormatterCallback<T> {
-    private final Set<T> result = CompactHashSet.create();
+      extends AggregateAllOutputFormatterCallback<T, Set<T>> {
+    private final Set<T> resultSet;
+    private final List<T> resultList;
+
+    private OrderedAggregateAllOutputFormatterCallbackImpl(QueryEnvironment<T> env) {
+      this.resultSet = env.createThreadSafeMutableSet();
+      this.resultList = new ArrayList<>();
+    }
 
     @Override
     public final synchronized void processOutput(Iterable<T> partialResult) {
-      Iterables.addAll(result, partialResult);
+      for (T element : partialResult) {
+        if (resultSet.add(element)) {
+          resultList.add(element);
+        }
+      }
     }
 
     @Override
     public synchronized Set<T> getResult() {
+      // A CompactHashSet's iteration order is the same as its insertion order.
+      CompactHashSet<T> result = CompactHashSet.createWithExpectedSize(resultList.size());
+      result.addAll(resultList);
       return result;
     }
   }
@@ -76,35 +100,120 @@ public final class QueryUtil {
    * {@link AggregateAllCallback#getResult} returns all the elements of the result in the order they
    * were processed.
    */
-  public static <T> AggregateAllOutputFormatterCallback<T>
-      newOrderedAggregateAllOutputFormatterCallback() {
-    return new OrderedAggregateAllOutputFormatterCallbackImpl<>();
+  public static <T> AggregateAllOutputFormatterCallback<T, Set<T>>
+      newOrderedAggregateAllOutputFormatterCallback(QueryEnvironment<T> env) {
+    return new OrderedAggregateAllOutputFormatterCallbackImpl<>(env);
   }
 
   /** Returns a fresh {@link AggregateAllCallback} instance. */
-  public static <T> AggregateAllCallback<T> newAggregateAllCallback() {
-    return new AggregateAllOutputFormatterCallbackImpl<>();
+  public static <T> AggregateAllCallback<T, ThreadSafeMutableSet<T>> newAggregateAllCallback(
+      QueryEnvironment<T> env) {
+    return new AggregateAllOutputFormatterCallbackImpl<>(env);
   }
 
   /**
-   * Returns a {@link QueryTaskFuture} representing the evaluation of {@code expr} as a (mutable)
-   * {@link Set} comprised of all the results.
+   * Returns a {@link QueryTaskFuture} representing the evaluation of {@code expr} as a mutable,
+   * thread safe {@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> QueryTaskFuture<Set<T>> evalAll(
+  public static <T> QueryTaskFuture<ThreadSafeMutableSet<T>> evalAll(
       QueryEnvironment<T> env, VariableContext<T> context, QueryExpression expr) {
-    final AggregateAllCallback<T> callback = newAggregateAllCallback();
+    final AggregateAllCallback<T, ThreadSafeMutableSet<T>> callback = newAggregateAllCallback(env);
     return env.whenSucceedsCall(
         env.eval(expr, context, callback),
-        new QueryTaskCallable<Set<T>>() {
+        new QueryTaskCallable<ThreadSafeMutableSet<T>>() {
           @Override
-          public Set<T> call() {
+          public ThreadSafeMutableSet<T> call() {
             return callback.getResult();
           }
         });
   }
 
+  /**
+   * A mutable thread safe {@link Set} that uses a {@link KeyExtractor} for determining equality of
+   * its elements. This is useful e.g. when {@code T} isn't guaranteed to have a useful
+   * {@link Object#equals} and {@link Object#hashCode} but {@code K} is.
+   */
+  public static class ThreadSafeMutableKeyExtractorBackedSetImpl<T, K>
+      extends AbstractSet<T> implements ThreadSafeMutableSet<T> {
+    private final KeyExtractor<T, K> extractor;
+    private final Class<T> elementClass;
+    private final ConcurrentMap<K, T> map;
+
+    public ThreadSafeMutableKeyExtractorBackedSetImpl(
+        KeyExtractor<T, K> extractor, Class<T> elementClass) {
+      this(extractor, elementClass, /*concurrencyLevel=*/ 1);
+    }
+
+    public ThreadSafeMutableKeyExtractorBackedSetImpl(
+        KeyExtractor<T, K> extractor,
+        Class<T> elementClass,
+        int concurrencyLevel) {
+      this.extractor = extractor;
+      this.elementClass = elementClass;
+      this.map = new MapMaker().concurrencyLevel(concurrencyLevel).makeMap();
+    }
+
+    @Override
+    public Iterator<T> iterator() {
+      return map.values().iterator();
+    }
+
+    @Override
+    public int size() {
+      return map.size();
+    }
+
+    @Override
+    public boolean add(T element) {
+      return map.putIfAbsent(extractor.extractKey(element), element) == null;
+    }
+
+    @Override
+    public boolean contains(Object obj) {
+      if (!elementClass.isInstance(obj)) {
+        return false;
+      }
+      T element = elementClass.cast(obj);
+      return map.containsKey(extractor.extractKey(element));
+    }
+
+    @Override
+    public boolean remove(Object obj) {
+      if (!elementClass.isInstance(obj)) {
+        return false;
+      }
+      T element = elementClass.cast(obj);
+      return map.remove(extractor.extractKey(element)) != null;
+    }
+  }
+
+  /**
+   * A {@link MutableMap} implementation that uses a {@link KeyExtractor} for determining equality
+   * of its keys.
+   */
+  public static class MutableKeyExtractorBackedMapImpl<T, K, V> implements MutableMap<T, V> {
+    private final KeyExtractor<T, K> extractor;
+    private final HashMap<K, V> map;
+
+    public MutableKeyExtractorBackedMapImpl(KeyExtractor<T, K> extractor) {
+      this.extractor = extractor;
+      this.map = new HashMap<>();
+    }
+
+    @Override
+    @Nullable
+    public V get(T key) {
+      return map.get(extractor.extractKey(key));
+    }
+
+    @Override
+    public V put(T key, V value) {
+      return map.put(extractor.extractKey(key), value);
+    }
+  }
+
   /** A trivial {@link Uniquifier} implementation. */
   public static class UniquifierImpl<T, K> implements Uniquifier<T> {
     private final KeyExtractor<T, K> extractor;