Auto merge of #101313 - SparrowLii:mk_attr_id, r=cjgillot

make `mk_attr_id` part of `ParseSess`

Updates #48685

The current `mk_attr_id` uses the `AtomicU32` type, which is not very efficient and adds a lot of lock contention in a parallel environment.

This PR refers to the task list in #48685, uses `mk_attr_id` as a method of the `AttrIdGenerator` struct, and adds a new field `attr_id_generator` to `ParseSess`.

`AttrIdGenerator` uses the `WorkerLocal`, which has two advantages: 1. `Cell` is more efficient than `AtomicU32`, and does not increase any lock contention. 2. We put the index of the work thread in the first few bits of the generated `AttrId`, so that the `AttrId` generated in different threads can be easily guaranteed to be unique.

cc `@cjgillot`

2 files changed, 57 insertions, 17 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index b4bb89af038..09f376a4800 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -2524,8 +2524,8 @@ impl<S: Encoder> Encodable<S> for AttrId {
 }
 
 impl<D: Decoder> Decodable<D> for AttrId {
-    fn decode(_: &mut D) -> AttrId {
-        crate::attr::mk_attr_id()
+    default fn decode(_: &mut D) -> AttrId {
+        panic!("cannot decode `AttrId` with `{}`", std::any::type_name::<D>());
     }
 }
 
diff --git a/compiler/rustc_ast/src/attr/mod.rs b/compiler/rustc_ast/src/attr/mod.rs
index 28198e69bff..990f4f8f132 100644
--- a/compiler/rustc_ast/src/attr/mod.rs
+++ b/compiler/rustc_ast/src/attr/mod.rs
@@ -11,12 +11,18 @@ use crate::tokenstream::{DelimSpan, Spacing, TokenTree};
 use crate::tokenstream::{LazyAttrTokenStream, TokenStream};
 use crate::util::comments;
 
+use rustc_data_structures::sync::WorkerLocal;
 use rustc_index::bit_set::GrowableBitSet;
 use rustc_span::source_map::BytePos;
 use rustc_span::symbol::{sym, Ident, Symbol};
 use rustc_span::Span;
 
+use std::cell::Cell;
 use std::iter;
+#[cfg(debug_assertions)]
+use std::ops::BitXor;
+#[cfg(debug_assertions)]
+use std::sync::atomic::{AtomicU32, Ordering};
 
 pub struct MarkedAttrs(GrowableBitSet<AttrId>);
 
@@ -346,22 +352,55 @@ pub fn mk_nested_word_item(ident: Ident) -> NestedMetaItem {
     NestedMetaItem::MetaItem(mk_word_item(ident))
 }
 
-pub(crate) fn mk_attr_id() -> AttrId {
-    use std::sync::atomic::AtomicU32;
-    use std::sync::atomic::Ordering;
+pub struct AttrIdGenerator(WorkerLocal<Cell<u32>>);
 
-    static NEXT_ATTR_ID: AtomicU32 = AtomicU32::new(0);
+#[cfg(debug_assertions)]
+static MAX_ATTR_ID: AtomicU32 = AtomicU32::new(u32::MAX);
 
-    let id = NEXT_ATTR_ID.fetch_add(1, Ordering::SeqCst);
-    assert!(id != u32::MAX);
-    AttrId::from_u32(id)
+impl AttrIdGenerator {
+    pub fn new() -> Self {
+        // We use `(index as u32).reverse_bits()` to initialize the
+        // starting value of AttrId in each worker thread.
+        // The `index` is the index of the worker thread.
+        // This ensures that the AttrId generated in each thread is unique.
+        AttrIdGenerator(WorkerLocal::new(|index| {
+            let index: u32 = index.try_into().unwrap();
+
+            #[cfg(debug_assertions)]
+            {
+                let max_id = ((index + 1).next_power_of_two() - 1).bitxor(u32::MAX).reverse_bits();
+                MAX_ATTR_ID.fetch_min(max_id, Ordering::Release);
+            }
+
+            Cell::new(index.reverse_bits())
+        }))
+    }
+
+    pub fn mk_attr_id(&self) -> AttrId {
+        let id = self.0.get();
+
+        // Ensure the assigned attr_id does not overlap the bits
+        // representing the number of threads.
+        #[cfg(debug_assertions)]
+        assert!(id <= MAX_ATTR_ID.load(Ordering::Acquire));
+
+        self.0.set(id + 1);
+        AttrId::from_u32(id)
+    }
 }
 
-pub fn mk_attr(style: AttrStyle, path: Path, args: MacArgs, span: Span) -> Attribute {
-    mk_attr_from_item(AttrItem { path, args, tokens: None }, None, style, span)
+pub fn mk_attr(
+    g: &AttrIdGenerator,
+    style: AttrStyle,
+    path: Path,
+    args: MacArgs,
+    span: Span,
+) -> Attribute {
+    mk_attr_from_item(g, AttrItem { path, args, tokens: None }, None, style, span)
 }
 
 pub fn mk_attr_from_item(
+    g: &AttrIdGenerator,
     item: AttrItem,
     tokens: Option<LazyAttrTokenStream>,
     style: AttrStyle,
@@ -369,29 +408,30 @@ pub fn mk_attr_from_item(
 ) -> Attribute {
     Attribute {
         kind: AttrKind::Normal(P(ast::NormalAttr { item, tokens })),
-        id: mk_attr_id(),
+        id: g.mk_attr_id(),
         style,
         span,
     }
 }
 
 /// Returns an inner attribute with the given value and span.
-pub fn mk_attr_inner(item: MetaItem) -> Attribute {
-    mk_attr(AttrStyle::Inner, item.path, item.kind.mac_args(item.span), item.span)
+pub fn mk_attr_inner(g: &AttrIdGenerator, item: MetaItem) -> Attribute {
+    mk_attr(g, AttrStyle::Inner, item.path, item.kind.mac_args(item.span), item.span)
 }
 
 /// Returns an outer attribute with the given value and span.
-pub fn mk_attr_outer(item: MetaItem) -> Attribute {
-    mk_attr(AttrStyle::Outer, item.path, item.kind.mac_args(item.span), item.span)
+pub fn mk_attr_outer(g: &AttrIdGenerator, item: MetaItem) -> Attribute {
+    mk_attr(g, AttrStyle::Outer, item.path, item.kind.mac_args(item.span), item.span)
 }
 
 pub fn mk_doc_comment(
+    g: &AttrIdGenerator,
     comment_kind: CommentKind,
     style: AttrStyle,
     data: Symbol,
     span: Span,
 ) -> Attribute {
-    Attribute { kind: AttrKind::DocComment(comment_kind, data), id: mk_attr_id(), style, span }
+    Attribute { kind: AttrKind::DocComment(comment_kind, data), id: g.mk_attr_id(), style, span }
 }
 
 pub fn list_contains_name(items: &[NestedMetaItem], name: Symbol) -> bool {