Add batch methods to WalkableGraph and convert SkyQueryEnvironment to use them.

--
MOS_MIGRATED_REVID=96214911

2 files changed, 77 insertions, 2 deletions

diff --git a/src/main/java/com/google/devtools/build/skyframe/DelegatingWalkableGraph.java b/src/main/java/com/google/devtools/build/skyframe/DelegatingWalkableGraph.java
index 5328a487b5..fee4f9c579 100644
--- a/src/main/java/com/google/devtools/build/skyframe/DelegatingWalkableGraph.java
+++ b/src/main/java/com/google/devtools/build/skyframe/DelegatingWalkableGraph.java
@@ -13,7 +13,13 @@
 // limitations under the License.
 package com.google.devtools.build.skyframe;
 
+import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Predicates;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Maps;
+
+import java.util.Map;
 
 import javax.annotation.Nullable;
 
@@ -33,6 +39,15 @@ public class DelegatingWalkableGraph implements WalkableGraph {
     return entry;
   }
 
+  private Map<SkyKey, NodeEntry> getEntries(Iterable<SkyKey> keys) {
+    Map<SkyKey, NodeEntry> result = graph.getBatch(keys);
+    Preconditions.checkState(result.size() == Iterables.size(keys), "%s %s", keys, result);
+    for (Map.Entry<SkyKey, NodeEntry> entry : result.entrySet()) {
+      Preconditions.checkState(entry.getValue().isDone(), entry);
+    }
+    return result;
+  }
+
   @Override
   public boolean exists(SkyKey key) {
     NodeEntry entry = graph.get(key);
@@ -45,6 +60,21 @@ public class DelegatingWalkableGraph implements WalkableGraph {
     return getEntry(key).getValue();
   }
 
+  private static final Function<NodeEntry, SkyValue> GET_SKY_VALUE_FUNCTION =
+      new Function<NodeEntry, SkyValue>() {
+        @Nullable
+        @Override
+        public SkyValue apply(NodeEntry entry) {
+          return entry.isDone() ? entry.getValue() : null;
+        }
+      };
+
+  @Override
+  public Map<SkyKey, SkyValue> getValuesMaybe(Iterable<SkyKey> keys) {
+    return Maps.filterValues(Maps.transformValues(graph.getBatch(keys), GET_SKY_VALUE_FUNCTION),
+        Predicates.notNull());
+  }
+
   @Nullable
   @Override
   public Exception getException(SkyKey key) {
@@ -57,8 +87,34 @@ public class DelegatingWalkableGraph implements WalkableGraph {
     return getEntry(key).getDirectDeps();
   }
 
+  private static final Function<NodeEntry, Iterable<SkyKey>> GET_DIRECT_DEPS_FUNCTION =
+      new Function<NodeEntry, Iterable<SkyKey>>() {
+        @Override
+        public Iterable<SkyKey> apply(NodeEntry entry) {
+          return entry.getDirectDeps();
+        }
+      };
+
+  @Override
+  public Map<SkyKey, Iterable<SkyKey>> getDirectDeps(Iterable<SkyKey> keys) {
+    return Maps.transformValues(getEntries(keys), GET_DIRECT_DEPS_FUNCTION);
+  }
+
   @Override
   public Iterable<SkyKey> getReverseDeps(SkyKey key) {
     return getEntry(key).getReverseDeps();
   }
+
+  private static final Function<NodeEntry, Iterable<SkyKey>> GET_REVERSE_DEPS_FUNCTION =
+      new Function<NodeEntry, Iterable<SkyKey>>() {
+        @Override
+        public Iterable<SkyKey> apply(NodeEntry entry) {
+          return entry.getReverseDeps();
+        }
+      };
+
+  @Override
+  public Map<SkyKey, Iterable<SkyKey>> getReverseDeps(Iterable<SkyKey> keys) {
+    return Maps.transformValues(getEntries(keys), GET_REVERSE_DEPS_FUNCTION);
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java b/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
index 68c39c9a2a..e020a26ba7 100644
--- a/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
+++ b/src/main/java/com/google/devtools/build/skyframe/WalkableGraph.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.skyframe;
 import com.google.devtools.build.lib.events.EventHandler;
 
 import java.util.Collection;
+import java.util.Map;
 
 import javax.annotation.Nullable;
 
@@ -40,6 +41,12 @@ public interface WalkableGraph {
   SkyValue getValue(SkyKey key);
 
   /**
+   * Returns a map giving the values of the given keys for done keys. Keys not present in the graph
+   * or whose nodes are not done will not be present in the returned map.
+   */
+  Map<SkyKey, SkyValue> getValuesMaybe(Iterable<SkyKey> keys);
+
+  /**
    * Returns the exception thrown when computing the node with the given key, if any. If the node
    * was computed successfully, returns null. A node with this key must exist in the graph.
    */
@@ -52,11 +59,23 @@ public interface WalkableGraph {
   Iterable<SkyKey> getDirectDeps(SkyKey key);
 
   /**
-   * Returns the reverse dependencies of the node with the given key. A node with this key must
-   * exist in the graph.
+   * Returns a map giving the direct dependencies of the nodes with the given keys. Same semantics
+   * as {@link #getDirectDeps(SkyKey)}.
    */
+  Map<SkyKey, Iterable<SkyKey>> getDirectDeps(Iterable<SkyKey> keys);
+
+    /**
+     * Returns the reverse dependencies of the node with the given key. A node with this key must
+     * exist in the graph.
+     */
   Iterable<SkyKey> getReverseDeps(SkyKey key);
 
+  /**
+   * Returns a map giving the reverse dependencies of the nodes with the given keys. Same semantics
+   * as {@link #getReverseDeps(SkyKey)}.
+   */
+  Map<SkyKey, Iterable<SkyKey>> getReverseDeps(Iterable<SkyKey> keys);
+
   /** Provides a WalkableGraph on demand after preparing it. */
   interface WalkableGraphFactory {
     WalkableGraph prepareAndGet(Collection<String> roots, int numThreads,