Bucket dep-graph nodes by kind during decode instead of a second pass

The reverse index groups nodes by DepKind so each kind's fingerprint map can
be built from a contiguous range. That grouping was a counting sort that ran
as a separate pass over the decoded nodes, re-reading every node after the
main decode loop had already visited it.

Move the per-kind node counts into the fixed-size trailer so the decoder can
read them up front, lay out the per-kind ranges before the main loop, and
place each node into its range as it is decoded. This drops the separate pass.

Author: xmakro

compiler/rustc_middle/src/dep_graph/serialized.rs

@@ -315,19 +315,28 @@ impl SerializedDepGraph {
 
         // `node_max` is the number of indices including empty nodes while `node_count`
         // is the number of actually encoded nodes.
-        let (node_max, node_count, edge_count) =
-            d.with_position(d.len() - 3 * IntEncodedWithFixedSize::ENCODED_SIZE, |d| {
+        //
+        // The trailer holds the fixed-size node/edge counts followed by the per-kind
+        // node counts, so all of them can be read by seeking from the end of the file.
+        // Reading the per-kind counts up front lets us bucket nodes by kind during the
+        // single decode pass below, rather than in a separate counting-sort pass.
+        let kind_count = DepKind::MAX as usize + 1;
+        let trailer_len = (3 + kind_count) * IntEncodedWithFixedSize::ENCODED_SIZE;
+        let (node_max, node_count, edge_count, kind_counts) =
+            d.with_position(d.len() - trailer_len, |d| {
                 debug!("position: {:?}", d.position());
                 let node_max = IntEncodedWithFixedSize::decode(d).0 as usize;
                 let node_count = IntEncodedWithFixedSize::decode(d).0 as usize;
                 let edge_count = IntEncodedWithFixedSize::decode(d).0 as usize;
-                (node_max, node_count, edge_count)
+                let kind_counts: Vec<u32> =
+                    (0..kind_count).map(|_| IntEncodedWithFixedSize::decode(d).0 as u32).collect();
+                (node_max, node_count, edge_count, kind_counts)
             });
         debug!("position: {:?}", d.position());
 
         debug!(?node_count, ?edge_count);
 
-        let graph_bytes = d.len() - (3 * IntEncodedWithFixedSize::ENCODED_SIZE) - d.position();
+        let graph_bytes = d.len() - trailer_len - d.position();
 
         let mut nodes = IndexVec::from_elem_n(
             DepNode {
@@ -352,6 +361,19 @@ impl SerializedDepGraph {
         let mut edge_list_data =
             Vec::with_capacity(graph_bytes - node_count * size_of::<SerializedNodeHeader>());
 
+        // Lay out the contiguous per-`DepKind` ranges of `nodes_by_kind` from the
+        // counts read above, so each node can be placed into its kind's range as it
+        // is decoded below. `fill[k]` is the next free slot in kind `k`'s range.
+        let mut kinds = Vec::with_capacity(kind_count);
+        let mut offset = 0u32;
+        for &len in &kind_counts {
+            kinds.push(LazyKindIndex { start: offset, len, map: OnceLock::new() });
+            offset += len;
+        }
+        debug_assert_eq!(offset as usize, node_count);
+        let mut nodes_by_kind = vec![None; node_count];
+        let mut fill: Vec<u32> = kinds.iter().map(|k| k.start).collect();
+
         for _ in 0..node_count {
             // Decode the header for this edge; the header packs together as many of the fixed-size
             // fields as possible to limit the number of times we update decoder state.
@@ -364,6 +386,12 @@ impl SerializedDepGraph {
             assert!(node_header.node().kind != DepKind::Null && node.kind == DepKind::Null);
             *node = node_header.node();
 
+            // Place this node into its kind's range for the lazy reverse index. Only
+            // encoded (non-`Null`) nodes reach this loop, so each fills one slot once.
+            let k = node.kind.as_usize();
+            nodes_by_kind[fill[k] as usize] = Some(index);
+            fill[k] += 1;
+
             value_fingerprints[index] = node_header.value_fingerprint();
 
             // If the length of this node's edge list is small, the length is stored in the header.
@@ -388,36 +416,9 @@ impl SerializedDepGraph {
         // end of the array. This padding ensure it doesn't.
         edge_list_data.extend(&[0u8; DEP_NODE_PAD]);
 
-        // Read the number of nodes of each dep kind, and perform
-        // counting sort for `LazyNodeIndex`.
-        let mut kinds = Vec::with_capacity(DepKind::MAX as usize + 1);
-        let mut offset = 0u32;
-        for _ in 0..(DepKind::MAX + 1) {
-            let len = d.read_u32();
-            kinds.push(LazyKindIndex { start: offset, len, map: OnceLock::new() });
-            offset += len;
-        }
-        debug_assert_eq!(offset as usize, node_count);
-
         let session_count = d.read_u64();
 
-        // Counting sort: place each node index into its kind's range. `fill[k]`
-        // points at the next free slot in kind `k`'s range, so a kind's nodes end
-        // up contiguous. Slots start as `None` and are each filled exactly once
-        // (the counts sum to the number of non-`Null` nodes).
-        let mut nodes_by_kind = vec![None; node_count];
-        let mut fill: Vec<u32> = kinds.iter().map(|k| k.start).collect();
-        for (idx, node) in nodes.iter_enumerated() {
-            // Unused indices from batch allocation stay `Null`; they carry no
-            // encoded node and are never looked up by fingerprint, so skip them.
-            if node.kind == DepKind::Null {
-                continue;
-            }
-            let k = node.kind.as_usize();
-            nodes_by_kind[fill[k] as usize] = Some(idx);
-            fill[k] += 1;
-        }
-        // Each kind's range was filled exactly to its end.
+        // Each kind's range was filled exactly to its end by the decode loop.
         debug_assert!(kinds.iter().zip(&fill).all(|(k, &f)| f == k.start + k.len));
         let reverse_index = LazyNodeIndex { nodes_by_kind, kinds };
 
@@ -885,20 +886,22 @@ impl EncoderState {
             }
         }
 
-        // Encode the number of each dep kind encountered
-        for count in kind_stats.iter() {
-            count.encode(&mut encoder);
-        }
-
+        // The session count is read sequentially right after the node data.
         self.previous.session_count.checked_add(1).unwrap().encode(&mut encoder);
 
         debug!(?node_max, ?node_count, ?edge_count);
         debug!("position: {:?}", encoder.position());
+        // Fixed-size trailer: the node/edge counts followed by the number of nodes of
+        // each dep kind, all readable by seeking from the end of the file. The decoder
+        // reads the per-kind counts up front to bucket nodes by kind in a single pass.
         IntEncodedWithFixedSize(node_max.try_into().unwrap()).encode(&mut encoder);
         IntEncodedWithFixedSize(node_count.try_into().unwrap()).encode(&mut encoder);
         IntEncodedWithFixedSize(edge_count.try_into().unwrap()).encode(&mut encoder);
+        for count in kind_stats.iter() {
+            IntEncodedWithFixedSize(*count as u64).encode(&mut encoder);
+        }
         debug!("position: {:?}", encoder.position());
-        // Drop the encoder so that nothing is written after the counts.
+        // Drop the encoder so that nothing is written after the trailer.
         let result = encoder.finish();
         if let Ok(position) = result {
             // FIXME(rylev): we hardcode the dep graph file name so we