In the Blaze Query implementation, use Set and Map implementations backed by the same KeyExtractor used that the Uniquifier implementation uses. This fixes a hypothetical issue where we were previously relying on Target#equals/hashCode.

RELNOTES: None
PiperOrigin-RevId: 159741545

1 files changed, 7 insertions, 6 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java
index 8b1fc37d62..de4cd341dd 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java
@@ -14,14 +14,13 @@
 package com.google.devtools.build.lib.query2.engine;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Sets;
+import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ArgumentType;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskFuture;
-import java.util.Collection;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.ThreadSafeMutableSet;
 import java.util.List;
-import java.util.Set;
 
 /**
  * A "deps" query expression, which computes the dependencies of the argument. An optional
@@ -64,8 +63,9 @@ final class DepsFunction implements QueryFunction {
     return env.eval(args.get(0).getExpression(), context, new Callback<T>() {
       @Override
       public void process(Iterable<T> partialResult) throws QueryException, InterruptedException {
-        Collection<T> current = Sets.newHashSet(partialResult);
-        env.buildTransitiveClosure(expression, (Set<T>) current, depthBound);
+        ThreadSafeMutableSet<T> current = env.createThreadSafeMutableSet();
+        Iterables.addAll(current, partialResult);
+        env.buildTransitiveClosure(expression, current, depthBound);
 
         // We need to iterate depthBound + 1 times.
         for (int i = 0; i <= depthBound; i++) {
@@ -75,7 +75,8 @@ final class DepsFunction implements QueryFunction {
           ImmutableList<T> toProcess =
               minDepthUniquifier.uniqueAtDepthLessThanOrEqualTo(current, i);
           callback.process(toProcess);
-          current = ImmutableList.copyOf(env.getFwdDeps(toProcess));
+          current = env.createThreadSafeMutableSet();
+          Iterables.addAll(current, env.getFwdDeps(toProcess));
           if (current.isEmpty()) {
             // Exit when there are no more nodes to visit.
             break;