perf: reuse green-marking's edge walk when promoting a node

compiler/rustc_middle/src/dep_graph/graph.rs

@@ -1,4 +1,5 @@
 use std::assert_matches;
+use std::cell::Cell;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::sync::Arc;
@@ -9,7 +10,7 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::profiling::QueryInvocationId;
 use rustc_data_structures::sharded::{self, ShardedHashMap};
 use rustc_data_structures::stable_hash::{StableHash, StableHasher};
-use rustc_data_structures::sync::{AtomicU64, Lock};
+use rustc_data_structures::sync::{AtomicU64, Lock, WorkerLocal};
 use rustc_data_structures::unord::UnordMap;
 use rustc_errors::DiagInner;
 use rustc_index::IndexVec;
@@ -89,6 +90,38 @@ pub(crate) struct MarkFrame<'a> {
     parent: Option<&'a MarkFrame<'a>>,
 }
 
+/// The edge list of one node being marked green: it occupies `buf[start..]` of the shared
+/// scratch buffer and is popped again on drop, restoring the buffer for the enclosing call.
+struct EdgeFrame<'a> {
+    buf: &'a mut Vec<DepNodeIndex>,
+    start: usize,
+}
+
+impl<'a> EdgeFrame<'a> {
+    #[inline]
+    fn new(buf: &'a mut Vec<DepNodeIndex>) -> Self {
+        EdgeFrame { start: buf.len(), buf }
+    }
+
+    #[inline]
+    fn push(&mut self, edge: DepNodeIndex) {
+        self.buf.push(edge);
+    }
+
+    /// The edges pushed onto this frame so far.
+    #[inline]
+    fn get(&self) -> &[DepNodeIndex] {
+        &self.buf[self.start..]
+    }
+}
+
+impl Drop for EdgeFrame<'_> {
+    #[inline]
+    fn drop(&mut self) {
+        self.buf.truncate(self.start);
+    }
+}
+
 #[derive(Debug)]
 pub(super) enum DepNodeColor {
     Green(DepNodeIndex),
@@ -119,6 +152,9 @@ pub struct DepGraphData {
     /// a particular query result was decoded from disk
     /// (not just marked green)
     debug_loaded_from_disk: Lock<FxHashSet<DepNode>>,
+
+    /// Per-worker edge buffer amortized across `try_mark_green` calls.
+    green_edge_buf: WorkerLocal<Cell<Vec<DepNodeIndex>>>,
 }
 
 pub fn hash_result<R>(hcx: &mut StableHashState<'_>, result: &R) -> Fingerprint
@@ -175,6 +211,7 @@ impl DepGraph {
                 previous: prev_graph,
                 colors,
                 debug_loaded_from_disk: Default::default(),
+                green_edge_buf: WorkerLocal::default(),
             })),
             virtual_dep_node_index: Arc::new(AtomicU32::new(0)),
         }
@@ -790,8 +827,9 @@ impl DepGraphData {
     fn promote_node_and_deps_to_current(
         &self,
         prev_index: SerializedDepNodeIndex,
+        edges: &[DepNodeIndex],
     ) -> Option<DepNodeIndex> {
-        let dep_node_index = self.current.encoder.send_promoted(prev_index, &self.colors);
+        let dep_node_index = self.current.encoder.send_promoted(prev_index, &self.colors, edges);
 
         #[cfg(debug_assertions)]
         if let Some(dep_node_index) = dep_node_index {
@@ -878,20 +916,29 @@ impl DepGraphData {
                 // in the previous compilation session too, so we can try to
                 // mark it as green by recursively marking all of its
                 // dependencies green.
-                self.try_mark_previous_green(tcx, prev_index, None)
-                    .map(|dep_node_index| (prev_index, dep_node_index))
+
+                // Reuse a per-worker buffer for the edges instead of allocating one per call.
+                // The recursion gives it back empty: each `EdgeFrame` pops its edges on drop.
+                let mut edge_buf = self.green_edge_buf.take();
+                let result = self.try_mark_previous_green(tcx, prev_index, None, &mut edge_buf);
+                debug_assert!(edge_buf.is_empty());
+                self.green_edge_buf.set(edge_buf);
+                result.map(|dep_node_index| (prev_index, dep_node_index))
             }
         }
     }
 
     /// Try to mark a dep-node which existed in the previous compilation session as green.
-    #[instrument(skip(self, tcx, prev_dep_node_index, frame), level = "debug")]
+    #[instrument(skip(self, tcx, prev_dep_node_index, frame, edge_buf), level = "debug")]
     fn try_mark_previous_green<'tcx>(
         &self,
         tcx: TyCtxt<'tcx>,
         prev_dep_node_index: SerializedDepNodeIndex,
         frame: Option<&MarkFrame<'_>>,
+        // Amortized buffer to store edges in.
+        edge_buf: &mut Vec<DepNodeIndex>,
     ) -> Option<DepNodeIndex> {
+        let mut edges = EdgeFrame::new(edge_buf);
         let frame = MarkFrame { index: prev_dep_node_index, parent: frame };
 
         // We never try to mark eval_always nodes as green
@@ -901,7 +948,10 @@ impl DepGraphData {
             match self.colors.get(parent_dep_node_index) {
                 // This dependency has been marked as green before, we are still ok and can
                 // continue checking the remaining dependencies.
-                DepNodeColor::Green(_) => continue,
+                DepNodeColor::Green(parent_index) => {
+                    edges.push(parent_index);
+                    continue;
+                }
 
                 // This dependency's result is different to the previous compilation session. We
                 // cannot mark this dep_node as green, so stop checking.
@@ -915,8 +965,16 @@ impl DepGraphData {
 
             // If this dependency isn't eval_always, try to mark it green recursively.
             if !tcx.is_eval_always(parent_dep_node.kind)
-                && self.try_mark_previous_green(tcx, parent_dep_node_index, Some(&frame)).is_some()
+                && let Some(parent_index) = self.try_mark_previous_green(
+                    tcx,
+                    parent_dep_node_index,
+                    Some(&frame),
+                    // Pass the edge buffer to the recursive call.
+                    // It will use an `EdgeFrame` to give it back unchanged.
+                    edges.buf,
+                )
             {
+                edges.push(parent_index);
                 continue;
             }
 
@@ -926,7 +984,10 @@ impl DepGraphData {
             }
 
             match self.colors.get(parent_dep_node_index) {
-                DepNodeColor::Green(_) => continue,
+                DepNodeColor::Green(parent_index) => {
+                    edges.push(parent_index);
+                    continue;
+                }
                 DepNodeColor::Red => return None,
                 DepNodeColor::Unknown => {}
             }
@@ -954,7 +1015,8 @@ impl DepGraphData {
         // adding all the appropriate edges imported from the previous graph.
         //
         // `no_hash` nodes may fail this promotion due to already being conservatively colored red.
-        let dep_node_index = self.promote_node_and_deps_to_current(prev_dep_node_index)?;
+        let dep_node_index =
+            self.promote_node_and_deps_to_current(prev_dep_node_index, edges.get())?;
 
         // ... and finally storing a "Green" entry in the color map.
         // Multiple threads can all write the same color here.

compiler/rustc_middle/src/dep_graph/serialized.rs

@@ -603,16 +603,12 @@ impl NodeInfo {
         node: &DepNode,
         index: DepNodeIndex,
         value_fingerprint: Fingerprint,
-        prev_index: SerializedDepNodeIndex,
-        colors: &DepNodeColorMap,
-        previous: &SerializedDepGraph,
+        edges: &[DepNodeIndex],
     ) -> usize {
-        let edges = previous.edge_targets_from(prev_index);
-        let edge_count = edges.size_hint().0;
+        let edge_count = edges.len();
 
         // Find the highest edge in the new dep node indices
-        let edge_max =
-            edges.clone().map(|i| colors.current(i).unwrap().as_u32()).max().unwrap_or(0);
+        let edge_max = edges.iter().map(|x| x.as_u32()).max().unwrap_or(0);
 
         let header =
             SerializedNodeHeader::new(node, index, value_fingerprint, edge_max, edge_count);
@@ -624,10 +620,10 @@ impl NodeInfo {
         }
 
         let bytes_per_index = header.bytes_per_index();
-        for node_index in edges {
-            let node_index = colors.current(node_index).unwrap();
+        for edge in edges {
+            let edge = edge.as_u32();
             e.write_with(|dest| {
-                *dest = node_index.as_u32().to_le_bytes();
+                *dest = edge.to_le_bytes();
                 bytes_per_index
             });
         }
@@ -728,25 +724,22 @@ impl EncoderState {
         node: &DepNode,
         index: DepNodeIndex,
         edge_count: usize,
-        edges: impl FnOnce(&Self) -> Vec<DepNodeIndex>,
+        edges: &[DepNodeIndex],
         retained_graph: &Option<Lock<RetainedDepGraph>>,
         local: &mut LocalEncoderState,
     ) {
         local.kind_stats[node.kind.as_usize()] += 1;
         local.edge_count += edge_count;
 
         if let Some(retained_graph) = &retained_graph {
-            // Call `edges` before the outlined code to allow the closure to be optimized out.
-            let edges = edges(self);
-
             // Outline the build of the full dep graph as it's typically disabled and cold.
             outline(move || {
                 // Block on the lock rather than using `try_lock`: under the parallel frontend
                 // several threads record nodes concurrently, and dropping a node on lock
                 // contention would make the retained graph nondeterministic. Readers take a
                 // clone of the graph (`retained_dep_graph`) rather than holding the lock, so
                 // this never deadlocks against a reentrant `record`.
-                retained_graph.lock().push(index, *node, &edges);
+                retained_graph.lock().push(index, *node, edges);
             });
         }
 
@@ -783,14 +776,7 @@ impl EncoderState {
     ) {
         node.encode(&mut local.encoder, index);
         self.flush_mem_encoder(&mut *local);
-        self.record(
-            &node.node,
-            index,
-            node.edges.len(),
-            |_| node.edges[..].to_vec(),
-            retained_graph,
-            &mut *local,
-        );
+        self.record(&node.node, index, node.edges.len(), &node.edges, retained_graph, &mut *local);
     }
 
     /// Encodes a node that was promoted from the previous graph. It reads the information directly from
@@ -805,34 +791,15 @@ impl EncoderState {
         index: DepNodeIndex,
         prev_index: SerializedDepNodeIndex,
         retained_graph: &Option<Lock<RetainedDepGraph>>,
-        colors: &DepNodeColorMap,
         local: &mut LocalEncoderState,
+        edges: &[DepNodeIndex],
     ) {
         let node = self.previous.index_to_node(prev_index);
         let value_fingerprint = self.previous.value_fingerprint_for_index(prev_index);
-        let edge_count = NodeInfo::encode_promoted(
-            &mut local.encoder,
-            node,
-            index,
-            value_fingerprint,
-            prev_index,
-            colors,
-            &self.previous,
-        );
+        let edge_count =
+            NodeInfo::encode_promoted(&mut local.encoder, node, index, value_fingerprint, edges);
         self.flush_mem_encoder(&mut *local);
-        self.record(
-            node,
-            index,
-            edge_count,
-            |this| {
-                this.previous
-                    .edge_targets_from(prev_index)
-                    .map(|i| colors.current(i).unwrap())
-                    .collect()
-            },
-            retained_graph,
-            &mut *local,
-        );
+        self.record(node, index, edge_count, edges, retained_graph, &mut *local);
     }
 
     fn finish(&self, profiler: &SelfProfilerRef, current: &CurrentDepGraph) -> FileEncodeResult {
@@ -1045,6 +1012,7 @@ impl GraphEncoder {
         &self,
         prev_index: SerializedDepNodeIndex,
         colors: &DepNodeColorMap,
+        edges: &[DepNodeIndex],
     ) -> Option<DepNodeIndex> {
         let _prof_timer = self.profiler.generic_activity("incr_comp_encode_dep_graph");
 
@@ -1060,8 +1028,8 @@ impl GraphEncoder {
                     index,
                     prev_index,
                     &self.retained_graph,
-                    colors,
                     &mut *local,
+                    edges,
                 );
                 Some(index)
             }