Use AtomicU32 for the index counter

Author: Zoxc

compiler/rustc_query_system/src/dep_graph/serialized.rs

@@ -43,7 +43,7 @@ use std::cell::RefCell;
 use std::cmp::max;
 use std::marker::PhantomData;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::{iter, mem, u64};
 
 use rustc_data_structures::fingerprint::{Fingerprint, PackedFingerprint};
@@ -544,7 +544,7 @@ struct LocalEncoderResult {
 }
 
 struct EncoderState<D: Deps> {
-    next_node_index: AtomicU64,
+    next_node_index: AtomicU32,
     previous: Arc<SerializedDepGraph>,
     file: Lock<Option<FileEncoder>>,
     local: WorkerLocal<RefCell<LocalEncoderState>>,
@@ -556,7 +556,7 @@ impl<D: Deps> EncoderState<D> {
     fn new(encoder: FileEncoder, record_stats: bool, previous: Arc<SerializedDepGraph>) -> Self {
         Self {
             previous,
-            next_node_index: AtomicU64::new(0),
+            next_node_index: AtomicU32::new(0),
             stats: record_stats.then(|| Lock::new(FxHashMap::default())),
             file: Lock::new(Some(encoder)),
             local: WorkerLocal::new(|_| {
@@ -578,11 +578,10 @@ impl<D: Deps> EncoderState<D> {
         if local.remaining_node_index == 0 {
             const COUNT: u32 = 256;
 
-            // We assume that there won't be enough active threads to overflow `u64` from `u32::MAX` here.
-            // This can exceed u32::MAX by at most `N` * `COUNT` where `N` is the thread pool count since
-            // `try_into().unwrap()` will make threads panic when `self.next_node_index` exceeds u32::MAX.
-            local.next_node_index =
-                self.next_node_index.fetch_add(COUNT as u64, Ordering::Relaxed).try_into().unwrap();
+            local.next_node_index = self
+                .next_node_index
+                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| v.checked_add(COUNT))
+                .unwrap();
 
             // Check that we'll stay within `u32`
             local.next_node_index.checked_add(COUNT).unwrap();
@@ -714,7 +713,7 @@ impl<D: Deps> EncoderState<D> {
 
     fn finish(&self, profiler: &SelfProfilerRef, current: &CurrentDepGraph<D>) -> FileEncodeResult {
         // Prevent more indices from being allocated.
-        self.next_node_index.store(u32::MAX as u64 + 1, Ordering::SeqCst);
+        self.next_node_index.store(u32::MAX, Ordering::SeqCst);
 
         let results = broadcast(|_| {
             let mut local = self.local.borrow_mut();