1 files changed, 68 insertions, 0 deletions
diff --git a/compiler/rustc_query_system/src/dep_graph/query.rs b/compiler/rustc_query_system/src/dep_graph/query.rs
new file mode 100644
index 000000000..27b3b5e13
--- /dev/null
+++ b/compiler/rustc_query_system/src/dep_graph/query.rs
@@ -0,0 +1,68 @@
+use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::graph::implementation::{Direction, Graph, NodeIndex, INCOMING};
+use rustc_index::vec::IndexVec;
+
+use super::{DepKind, DepNode, DepNodeIndex};
+
+pub struct DepGraphQuery<K> {
+    pub graph: Graph<DepNode<K>, ()>,
+    pub indices: FxHashMap<DepNode<K>, NodeIndex>,
+    pub dep_index_to_index: IndexVec<DepNodeIndex, Option<NodeIndex>>,
+}
+
+impl<K: DepKind> DepGraphQuery<K> {
+    pub fn new(prev_node_count: usize) -> DepGraphQuery<K> {
+        let node_count = prev_node_count + prev_node_count / 4;
+        let edge_count = 6 * node_count;
+
+        let graph = Graph::with_capacity(node_count, edge_count);
+        let indices = FxHashMap::default();
+        let dep_index_to_index = IndexVec::new();
+
+        DepGraphQuery { graph, indices, dep_index_to_index }
+    }
+
+    pub fn push(&mut self, index: DepNodeIndex, node: DepNode<K>, edges: &[DepNodeIndex]) {
+        let source = self.graph.add_node(node);
+        if index.index() >= self.dep_index_to_index.len() {
+            self.dep_index_to_index.resize(index.index() + 1, None);
+        }
+        self.dep_index_to_index[index] = Some(source);
+        self.indices.insert(node, source);
+
+        for &target in edges.iter() {
+            let target = self.dep_index_to_index[target];
+            // We may miss the edges that are pushed while the `DepGraphQuery` is being accessed.
+            // Skip them to issues.
+            if let Some(target) = target {
+                self.graph.add_edge(source, target, ());
+            }
+        }
+    }
+
+    pub fn nodes(&self) -> Vec<&DepNode<K>> {
+        self.graph.all_nodes().iter().map(|n| &n.data).collect()
+    }
+
+    pub fn edges(&self) -> Vec<(&DepNode<K>, &DepNode<K>)> {
+        self.graph
+            .all_edges()
+            .iter()
+            .map(|edge| (edge.source(), edge.target()))
+            .map(|(s, t)| (self.graph.node_data(s), self.graph.node_data(t)))
+            .collect()
+    }
+
+    fn reachable_nodes(&self, node: &DepNode<K>, direction: Direction) -> Vec<&DepNode<K>> {
+        if let Some(&index) = self.indices.get(node) {
+            self.graph.depth_traverse(index, direction).map(|s| self.graph.node_data(s)).collect()
+        } else {
+            vec![]
+        }
+    }
+
+    /// All nodes that can reach `node`.
+    pub fn transitive_predecessors(&self, node: &DepNode<K>) -> Vec<&DepNode<K>> {
+        self.reachable_nodes(node, INCOMING)
+    }
+}