Reuse the index from promoted nodes when coloring executed tasks

2 files changed, 88 insertions, 36 deletions

diff --git a/compiler/rustc_query_system/src/dep_graph/graph.rs b/compiler/rustc_query_system/src/dep_graph/graph.rs
index 9f34417973e..56fd3a660d0 100644
--- a/compiler/rustc_query_system/src/dep_graph/graph.rs
+++ b/compiler/rustc_query_system/src/dep_graph/graph.rs
@@ -141,7 +141,7 @@ impl<D: Deps> DepGraph<D> {
         let colors = DepNodeColorMap::new(prev_graph_node_count);
 
         // Instantiate a node with zero dependencies only once for anonymous queries.
-        let _green_node_index = current.alloc_node(
+        let _green_node_index = current.alloc_new_node(
             DepNode { kind: D::DEP_KIND_ANON_ZERO_DEPS, hash: current.anon_id_seed.into() },
             EdgesVec::new(),
             Fingerprint::ZERO,
@@ -149,7 +149,7 @@ impl<D: Deps> DepGraph<D> {
         assert_eq!(_green_node_index, DepNodeIndex::SINGLETON_ZERO_DEPS_ANON_NODE);
 
         // Instantiate a dependy-less red node only once for anonymous queries.
-        let red_node_index = current.alloc_node(
+        let red_node_index = current.alloc_new_node(
             DepNode { kind: D::DEP_KIND_RED, hash: Fingerprint::ZERO.into() },
             EdgesVec::new(),
             Fingerprint::ZERO,
@@ -438,7 +438,7 @@ impl<D: Deps> DepGraphData<D> {
                 // memory impact of this `anon_node_to_index` map remains tolerable, and helps
                 // us avoid useless growth of the graph with almost-equivalent nodes.
                 self.current.anon_node_to_index.get_or_insert_with(target_dep_node, || {
-                    self.current.alloc_node(target_dep_node, task_deps, Fingerprint::ZERO)
+                    self.current.alloc_new_node(target_dep_node, task_deps, Fingerprint::ZERO)
                 })
             }
         };
@@ -680,8 +680,8 @@ impl<D: Deps> DepGraphData<D> {
         qcx: Qcx,
         diagnostic: &DiagInner,
     ) -> DepNodeIndex {
-        // Use `send` so we get an unique index, even though the dep node is not.
-        let dep_node_index = self.current.encoder.send(
+        // Use `send_new` so we get an unique index, even though the dep node is not.
+        let dep_node_index = self.current.encoder.send_new(
             DepNode {
                 kind: D::DEP_KIND_SIDE_EFFECT,
                 hash: PackedFingerprint::from(Fingerprint::ZERO),
@@ -713,20 +713,22 @@ impl<D: Deps> DepGraphData<D> {
                 }
             }
 
-            // Manually recreate the node as `promote_node_and_deps_to_current` expects all
-            // green dependencies.
-            let dep_node_index = self.current.encoder.send(
+            // Use `send_and_color` as `promote_node_and_deps_to_current` expects all
+            // green dependencies. `send_and_color` will also prevent multiple nodes
+            // being encoded for concurrent calls.
+            let dep_node_index = self.current.encoder.send_and_color(
+                prev_index,
+                &self.colors,
                 DepNode {
                     kind: D::DEP_KIND_SIDE_EFFECT,
                     hash: PackedFingerprint::from(Fingerprint::ZERO),
                 },
                 Fingerprint::ZERO,
                 std::iter::once(DepNodeIndex::FOREVER_RED_NODE).collect(),
+                true,
             );
+            // This will just overwrite the same value for concurrent calls.
             qcx.store_side_effect(dep_node_index, side_effect);
-
-            // Mark the node as green.
-            self.colors.insert(prev_index, DepNodeColor::Green(dep_node_index));
         })
     }
 
@@ -736,38 +738,43 @@ impl<D: Deps> DepGraphData<D> {
         edges: EdgesVec,
         fingerprint: Option<Fingerprint>,
     ) -> DepNodeIndex {
-        let dep_node_index =
-            self.current.alloc_node(key, edges, fingerprint.unwrap_or(Fingerprint::ZERO));
-
         if let Some(prev_index) = self.previous.node_to_index_opt(&key) {
             // Determine the color and index of the new `DepNode`.
-            let color = if let Some(fingerprint) = fingerprint {
+            let is_green = if let Some(fingerprint) = fingerprint {
                 if fingerprint == self.previous.fingerprint_by_index(prev_index) {
                     // This is a green node: it existed in the previous compilation,
                     // its query was re-executed, and it has the same result as before.
-                    DepNodeColor::Green(dep_node_index)
+                    true
                 } else {
                     // This is a red node: it existed in the previous compilation, its query
                     // was re-executed, but it has a different result from before.
-                    DepNodeColor::Red
+                    false
                 }
             } else {
                 // This is a red node, effectively: it existed in the previous compilation
                 // session, its query was re-executed, but it doesn't compute a result hash
                 // (i.e. it represents a `no_hash` query), so we have no way of determining
                 // whether or not the result was the same as before.
-                DepNodeColor::Red
+                false
             };
 
-            debug_assert!(
-                self.colors.get(prev_index).is_none(),
-                "DepGraph::with_task() - Duplicate DepNodeColor insertion for {key:?}",
+            let fingerprint = fingerprint.unwrap_or(Fingerprint::ZERO);
+
+            let dep_node_index = self.current.encoder.send_and_color(
+                prev_index,
+                &self.colors,
+                key,
+                fingerprint,
+                edges,
+                is_green,
             );
 
-            self.colors.insert(prev_index, color);
-        }
+            self.current.record_node(dep_node_index, key, fingerprint);
 
-        dep_node_index
+            dep_node_index
+        } else {
+            self.current.alloc_new_node(key, edges, fingerprint.unwrap_or(Fingerprint::ZERO))
+        }
     }
 
     fn promote_node_and_deps_to_current(&self, prev_index: SerializedDepNodeIndex) -> DepNodeIndex {
@@ -1246,19 +1253,15 @@ impl<D: Deps> CurrentDepGraph<D> {
         assert_eq!(previous, fingerprint, "Unstable fingerprints for {:?}", key);
     }
 
-    /// Writes the node to the current dep-graph and allocates a `DepNodeIndex` for it.
-    /// Assumes that this is a node that has no equivalent in the previous dep-graph.
     #[inline(always)]
-    fn alloc_node(
+    fn record_node(
         &self,
+        dep_node_index: DepNodeIndex,
         key: DepNode,
-        edges: EdgesVec,
-        current_fingerprint: Fingerprint,
-    ) -> DepNodeIndex {
-        let dep_node_index = self.encoder.send(key, current_fingerprint, edges);
-
+        _current_fingerprint: Fingerprint,
+    ) {
         #[cfg(debug_assertions)]
-        self.record_edge(dep_node_index, key, current_fingerprint);
+        self.record_edge(dep_node_index, key, _current_fingerprint);
 
         if let Some(ref nodes_in_current_session) = self.nodes_in_current_session {
             outline(|| {
@@ -1267,6 +1270,20 @@ impl<D: Deps> CurrentDepGraph<D> {
                 }
             });
         }
+    }
+
+    /// Writes the node to the current dep-graph and allocates a `DepNodeIndex` for it.
+    /// Assumes that this is a node that has no equivalent in the previous dep-graph.
+    #[inline(always)]
+    fn alloc_new_node(
+        &self,
+        key: DepNode,
+        edges: EdgesVec,
+        current_fingerprint: Fingerprint,
+    ) -> DepNodeIndex {
+        let dep_node_index = self.encoder.send_new(key, current_fingerprint, edges);
+
+        self.record_node(dep_node_index, key, current_fingerprint);
 
         dep_node_index
     }
diff --git a/compiler/rustc_query_system/src/dep_graph/serialized.rs b/compiler/rustc_query_system/src/dep_graph/serialized.rs
index 7750d6d1fef..7556a2456aa 100644
--- a/compiler/rustc_query_system/src/dep_graph/serialized.rs
+++ b/compiler/rustc_query_system/src/dep_graph/serialized.rs
@@ -707,7 +707,8 @@ impl<D: Deps> GraphEncoder<D> {
         }
     }
 
-    pub(crate) fn send(
+    /// Encodes a node that does not exists in the previous graph.
+    pub(crate) fn send_new(
         &self,
         node: DepNode,
         fingerprint: Fingerprint,
@@ -718,6 +719,40 @@ impl<D: Deps> GraphEncoder<D> {
         self.status.lock().as_mut().unwrap().encode_node(&node, &self.record_graph)
     }
 
+    /// Encodes a node that exists in the previous graph, but was re-executed.
+    ///
+    /// This will also ensure the dep node is colored either red or green.
+    pub(crate) fn send_and_color(
+        &self,
+        prev_index: SerializedDepNodeIndex,
+        colors: &DepNodeColorMap,
+        node: DepNode,
+        fingerprint: Fingerprint,
+        edges: EdgesVec,
+        is_green: bool,
+    ) -> DepNodeIndex {
+        let _prof_timer = self.profiler.generic_activity("incr_comp_encode_dep_graph");
+        let node = NodeInfo { node, fingerprint, edges };
+
+        let mut status = self.status.lock();
+        let status = status.as_mut().unwrap();
+
+        // Check colors inside the lock to avoid racing when `send_promoted` is called concurrently
+        // on the same index.
+        match colors.get(prev_index) {
+            None => {
+                let dep_node_index = status.encode_node(&node, &self.record_graph);
+                colors.insert(
+                    prev_index,
+                    if is_green { DepNodeColor::Green(dep_node_index) } else { DepNodeColor::Red },
+                );
+                dep_node_index
+            }
+            Some(DepNodeColor::Green(dep_node_index)) => dep_node_index,
+            Some(DepNodeColor::Red) => panic!(),
+        }
+    }
+
     /// Encodes a node that was promoted from the previous graph. It reads the information directly from
     /// the previous dep graph and expects all edges to already have a new dep node index assigned.
     ///
@@ -733,8 +768,8 @@ impl<D: Deps> GraphEncoder<D> {
         let mut status = self.status.lock();
         let status = status.as_mut().unwrap();
 
-        // Check colors inside the lock to avoid racing when `send_promoted` is called concurrently
-        // on the same index.
+        // Check colors inside the lock to avoid racing when `send_promoted` or `send_and_color`
+        // is called concurrently on the same index.
         match colors.get(prev_index) {
             None => {
                 let dep_node_index =