Use BTreeMap to store attributes.

10 files changed, 117 insertions, 47 deletions

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 18ac884c7cb..43b93d03ff6 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -258,12 +258,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         ex.span = e.span;
                     }
                     // Merge attributes into the inner expression.
-                    self.attrs[ex.hir_id] = &*self.arena.alloc_from_iter(
-                        e.attrs
-                            .iter()
-                            .map(|a| self.lower_attr(a))
-                            .chain(self.attrs[ex.hir_id].iter().cloned()),
-                    );
+                    if !e.attrs.is_empty() {
+                        let old_attrs = self.attrs.get(&ex.hir_id).map(|la| *la).unwrap_or(&[]);
+                        self.attrs.insert(
+                            ex.hir_id,
+                            &*self.arena.alloc_from_iter(
+                                e.attrs
+                                    .iter()
+                                    .map(|a| self.lower_attr(a))
+                                    .chain(old_attrs.iter().cloned()),
+                            ),
+                        );
+                    }
                     return ex;
                 }
 
@@ -1016,7 +1022,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
         // Introduce a `let` for destructuring: `let (lhs1, lhs2) = t`.
         let destructure_let = self.stmt_let_pat(
-            &[],
+            None,
             whole_span,
             Some(rhs),
             pat,
@@ -1775,7 +1781,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
         // `let mut __next`
         let next_let = self.stmt_let_pat(
-            &[],
+            None,
             desugared_span,
             None,
             next_pat,
@@ -1785,7 +1791,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         // `let <pat> = __next`
         let pat = self.lower_pat(pat);
         let pat_let = self.stmt_let_pat(
-            &[],
+            None,
             desugared_span,
             Some(next_expr),
             pat,
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index 570aaab0d91..777d3a5060b 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -251,7 +251,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         id: NodeId,
         hir_id: hir::HirId,
         ident: &mut Ident,
-        attrs: &'hir [Attribute],
+        attrs: Option<&'hir [Attribute]>,
         vis: &mut hir::Visibility<'hir>,
         i: &ItemKind,
     ) -> hir::ItemKind<'hir> {
@@ -502,7 +502,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         id: NodeId,
         vis: &mut hir::Visibility<'hir>,
         ident: &mut Ident,
-        attrs: &'hir [Attribute],
+        attrs: Option<&'hir [Attribute]>,
     ) -> hir::ItemKind<'hir> {
         debug!("lower_use_tree(tree={:?})", tree);
         debug!("lower_use_tree: vis = {:?}", vis);
@@ -551,7 +551,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         let path = this.lower_path_extra(res, &path, ParamMode::Explicit, None);
                         let kind = hir::ItemKind::Use(path, hir::UseKind::Single);
                         let vis = this.rebuild_vis(&vis);
-                        this.attrs.push_sparse(new_id, attrs);
+                        if let Some(attrs) = attrs {
+                            this.attrs.insert(new_id, attrs);
+                        }
 
                         this.insert_item(hir::Item {
                             def_id: new_id.expect_owner(),
@@ -623,7 +625,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
                         let kind =
                             this.lower_use_tree(use_tree, &prefix, id, &mut vis, &mut ident, attrs);
-                        this.attrs.push_sparse(new_hir_id, attrs);
+                        if let Some(attrs) = attrs {
+                            this.attrs.insert(new_hir_id, attrs);
+                        }
 
                         this.insert_item(hir::Item {
                             def_id: new_hir_id.expect_owner(),
@@ -770,7 +774,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             ),
             VariantData::Tuple(ref fields, id) => {
                 let ctor_id = self.lower_node_id(id);
-                self.attrs.push_sparse(ctor_id, self.attrs[parent_id]);
+                self.alias_attrs(ctor_id, parent_id);
                 hir::VariantData::Tuple(
                     self.arena.alloc_from_iter(
                         fields.iter().enumerate().map(|f| self.lower_struct_field(f)),
@@ -780,7 +784,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
             VariantData::Unit(id) => {
                 let ctor_id = self.lower_node_id(id);
-                self.attrs.push_sparse(ctor_id, self.attrs[parent_id]);
+                self.alias_attrs(ctor_id, parent_id);
                 hir::VariantData::Unit(ctor_id)
             }
         }
@@ -1168,7 +1172,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 //
                 // If this is the simple case, this parameter will end up being the same as the
                 // original parameter, but with a different pattern id.
-                let stmt_attrs = this.attrs[parameter.hir_id];
+                let stmt_attrs = this.attrs.get(&parameter.hir_id).copied();
                 let (new_parameter_pat, new_parameter_id) = this.pat_ident(desugared_span, ident);
                 let new_parameter = hir::Param {
                     hir_id: parameter.hir_id,
@@ -1213,7 +1217,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     );
                     let move_expr = this.expr_ident(desugared_span, ident, new_parameter_id);
                     let move_stmt = this.stmt_let_pat(
-                        &[],
+                        None,
                         desugared_span,
                         Some(move_expr),
                         move_pat,
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 52646557fc3..057e4d81f80 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -114,7 +114,7 @@ struct LoweringContext<'a, 'hir: 'a> {
 
     generator_kind: Option<hir::GeneratorKind>,
 
-    attrs: hir::HirIdVec<&'hir [Attribute]>,
+    attrs: BTreeMap<hir::HirId, &'hir [Attribute]>,
 
     /// When inside an `async` context, this is the `HirId` of the
     /// `task_context` local bound to the resume argument of the generator.
@@ -311,7 +311,7 @@ pub fn lower_crate<'a, 'hir>(
         bodies: BTreeMap::new(),
         trait_impls: BTreeMap::new(),
         modules: BTreeMap::new(),
-        attrs: hir::HirIdVec::default(),
+        attrs: BTreeMap::default(),
         exported_macros: Vec::new(),
         non_exported_macro_attrs: Vec::new(),
         catch_scopes: Vec::new(),
@@ -595,8 +595,13 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         self.resolver.definitions().init_def_id_to_hir_id_mapping(def_id_to_hir_id);
 
-        // Not all HIR owners have declared attrs. Complete with empty IndexVecs.
-        self.attrs.push_owner(Idx::new(self.resolver.definitions().def_index_count() - 1));
+        #[cfg(debug_assertions)]
+        for (&id, attrs) in self.attrs.iter() {
+            // Verify that we do not store empty slices in the map.
+            if attrs.is_empty() {
+                panic!("Stored empty attributes for {:?}", id);
+            }
+        }
 
         hir::Crate {
             item: hir::CrateItem { module, span: c.span },
@@ -973,10 +978,15 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         ret
     }
 
-    fn lower_attrs(&mut self, id: hir::HirId, attrs: &[Attribute]) -> &'hir [Attribute] {
-        let ret = self.arena.alloc_from_iter(attrs.iter().map(|a| self.lower_attr(a)));
-        self.attrs.push_sparse(id, ret);
-        ret
+    fn lower_attrs(&mut self, id: hir::HirId, attrs: &[Attribute]) -> Option<&'hir [Attribute]> {
+        if attrs.is_empty() {
+            None
+        } else {
+            let ret = self.arena.alloc_from_iter(attrs.iter().map(|a| self.lower_attr(a)));
+            debug_assert!(!ret.is_empty());
+            self.attrs.insert(id, ret);
+            Some(ret)
+        }
     }
 
     fn lower_attr(&self, attr: &Attribute) -> Attribute {
@@ -999,6 +1009,13 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         Attribute { kind, id: attr.id, style: attr.style, span: attr.span }
     }
 
+    fn alias_attrs(&mut self, id: hir::HirId, target_id: hir::HirId) {
+        if let Some(&a) = self.attrs.get(&target_id) {
+            debug_assert!(!a.is_empty());
+            self.attrs.insert(id, a);
+        }
+    }
+
     fn lower_mac_args(&self, args: &MacArgs) -> MacArgs {
         match *args {
             MacArgs::Empty => MacArgs::Empty,
@@ -2447,7 +2464,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                     })
                     .collect();
                 let hir_id = self.lower_node_id(s.id);
-                self.attrs.push_sparse(hir_id, self.attrs[l.hir_id]);
+                self.alias_attrs(hir_id, l.hir_id);
                 ids.push({
                     hir::Stmt {
                         hir_id,
@@ -2476,13 +2493,13 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             StmtKind::Expr(ref e) => {
                 let e = self.lower_expr(e);
                 let hir_id = self.lower_node_id(s.id);
-                self.attrs.push_sparse(hir_id, self.attrs[e.hir_id]);
+                self.alias_attrs(hir_id, e.hir_id);
                 (hir_id, hir::StmtKind::Expr(e))
             }
             StmtKind::Semi(ref e) => {
                 let e = self.lower_expr(e);
                 let hir_id = self.lower_node_id(s.id);
-                self.attrs.push_sparse(hir_id, self.attrs[e.hir_id]);
+                self.alias_attrs(hir_id, e.hir_id);
                 (hir_id, hir::StmtKind::Semi(e))
             }
             StmtKind::Empty => return smallvec![],
@@ -2532,14 +2549,17 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
     fn stmt_let_pat(
         &mut self,
-        attrs: &'hir [Attribute],
+        attrs: Option<&'hir [Attribute]>,
         span: Span,
         init: Option<&'hir hir::Expr<'hir>>,
         pat: &'hir hir::Pat<'hir>,
         source: hir::LocalSource,
     ) -> hir::Stmt<'hir> {
         let hir_id = self.next_id();
-        self.attrs.push_sparse(hir_id, attrs);
+        if let Some(a) = attrs {
+            debug_assert!(!a.is_empty());
+            self.attrs.insert(hir_id, a);
+        }
         let local = hir::Local { hir_id, init, pat, source, span, ty: None };
         self.stmt(span, hir::StmtKind::Local(self.arena.alloc(local)))
     }
diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
index ff1cee1fc20..8f61adcd8e2 100644
--- a/compiler/rustc_hir/src/hir.rs
+++ b/compiler/rustc_hir/src/hir.rs
@@ -2,7 +2,7 @@
 use crate::def::{CtorKind, DefKind, Namespace, Res};
 use crate::def_id::DefId;
 crate use crate::hir_id::HirId;
-use crate::{itemlikevisit, HirIdVec, LangItem};
+use crate::{itemlikevisit, LangItem};
 
 use rustc_ast::util::parser::ExprPrecedence;
 use rustc_ast::{self as ast, CrateSugar, LlvmAsmDialect};
@@ -675,7 +675,7 @@ pub struct Crate<'hir> {
     pub trait_map: BTreeMap<HirId, Vec<TraitCandidate>>,
 
     /// Collected attributes from HIR nodes.
-    pub attrs: HirIdVec<&'hir [Attribute]>,
+    pub attrs: BTreeMap<HirId, &'hir [Attribute]>,
 }
 
 impl Crate<'hir> {
diff --git a/compiler/rustc_hir/src/intravisit.rs b/compiler/rustc_hir/src/intravisit.rs
index b83c008d717..df63f0d48c3 100644
--- a/compiler/rustc_hir/src/intravisit.rs
+++ b/compiler/rustc_hir/src/intravisit.rs
@@ -477,7 +477,7 @@ pub trait Visitor<'v>: Sized {
 pub fn walk_crate<'v, V: Visitor<'v>>(visitor: &mut V, krate: &'v Crate<'v>) {
     visitor.visit_mod(&krate.item.module, krate.item.span, CRATE_HIR_ID);
     walk_list!(visitor, visit_macro_def, krate.exported_macros);
-    for (id, attrs) in krate.attrs.iter_enumerated() {
+    for (&id, attrs) in krate.attrs.iter() {
         for a in *attrs {
             visitor.visit_attribute(id, a)
         }
diff --git a/compiler/rustc_hir_pretty/src/lib.rs b/compiler/rustc_hir_pretty/src/lib.rs
index 4da714aff31..9c2766a1631 100644
--- a/compiler/rustc_hir_pretty/src/lib.rs
+++ b/compiler/rustc_hir_pretty/src/lib.rs
@@ -16,6 +16,7 @@ use rustc_target::spec::abi::Abi;
 
 use std::borrow::Cow;
 use std::cell::Cell;
+use std::collections::BTreeMap;
 use std::vec;
 
 pub fn id_to_string(map: &dyn rustc_hir::intravisit::Map<'_>, hir_id: hir::HirId) -> String {
@@ -82,7 +83,7 @@ impl PpAnn for &dyn rustc_hir::intravisit::Map<'_> {
 pub struct State<'a> {
     pub s: pp::Printer,
     comments: Option<Comments<'a>>,
-    attrs: &'a hir::HirIdVec<&'a [ast::Attribute]>,
+    attrs: &'a BTreeMap<hir::HirId, &'a [ast::Attribute]>,
     ann: &'a (dyn PpAnn + 'a),
 }
 
@@ -169,7 +170,7 @@ pub fn print_crate<'a>(
     // When printing the AST, we sometimes need to inject `#[no_std]` here.
     // Since you can't compile the HIR, it's not necessary.
 
-    s.print_mod(&krate.item.module, krate.attrs[hir::CRATE_HIR_ID]);
+    s.print_mod(&krate.item.module, s.attrs(hir::CRATE_HIR_ID));
     s.print_remaining_comments();
     s.s.eof()
 }
@@ -179,7 +180,7 @@ impl<'a> State<'a> {
         sm: &'a SourceMap,
         filename: FileName,
         input: String,
-        attrs: &'a hir::HirIdVec<&[ast::Attribute]>,
+        attrs: &'a BTreeMap<hir::HirId, &[ast::Attribute]>,
         ann: &'a dyn PpAnn,
     ) -> State<'a> {
         State {
@@ -191,7 +192,7 @@ impl<'a> State<'a> {
     }
 
     fn attrs(&self, id: hir::HirId) -> &'a [ast::Attribute] {
-        self.attrs.get(id).map_or(&[], |la| *la)
+        self.attrs.get(&id).map_or(&[], |la| *la)
     }
 }
 
@@ -200,7 +201,7 @@ where
     F: FnOnce(&mut State<'_>),
 {
     let mut printer =
-        State { s: pp::mk_printer(), comments: None, attrs: &hir::HirIdVec::default(), ann };
+        State { s: pp::mk_printer(), comments: None, attrs: &BTreeMap::default(), ann };
     f(&mut printer);
     printer.s.eof()
 }
diff --git a/compiler/rustc_middle/src/hir/map/collector.rs b/compiler/rustc_middle/src/hir/map/collector.rs
index 0d2db4af7a8..b1dd405a6be 100644
--- a/compiler/rustc_middle/src/hir/map/collector.rs
+++ b/compiler/rustc_middle/src/hir/map/collector.rs
@@ -132,11 +132,7 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
             hcx,
             hir_body_nodes,
             map: (0..definitions.def_index_count())
-                .map(|id| HirOwnerData {
-                    attrs: krate.attrs.get_owner(Idx::new(id)),
-                    signature: None,
-                    with_bodies: None,
-                })
+                .map(|_| HirOwnerData { signature: None, with_bodies: None })
                 .collect(),
         };
         collector.insert_entry(
diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index f0deb8bb668..41ecffb9c56 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/compiler/rustc_middle/src/hir/map/mod.rs
@@ -88,7 +88,6 @@ fn is_body_owner<'hir>(node: Node<'hir>, hir_id: HirId) -> bool {
 
 #[derive(Debug)]
 pub(super) struct HirOwnerData<'hir> {
-    pub(super) attrs: &'hir IndexVec<ItemLocalId, &'hir [ast::Attribute]>,
     pub(super) signature: Option<&'hir Owner<'hir>>,
     pub(super) with_bodies: Option<&'hir mut OwnerNodes<'hir>>,
 }
@@ -851,7 +850,7 @@ impl<'hir> Map<'hir> {
     /// Given a node ID, gets a list of attributes associated with the AST
     /// corresponding to the node-ID.
     pub fn attrs(&self, id: HirId) -> &'hir [ast::Attribute] {
-        self.tcx.hir_attrs(id.owner).get(id.local_id).copied().unwrap_or(&[])
+        self.tcx.hir_attrs(id.owner).get(id.local_id)
     }
 
     /// Gets the span of the definition of the specified HIR node.
diff --git a/compiler/rustc_middle/src/hir/mod.rs b/compiler/rustc_middle/src/hir/mod.rs
index 491934b845e..cf4e473d8ac 100644
--- a/compiler/rustc_middle/src/hir/mod.rs
+++ b/compiler/rustc_middle/src/hir/mod.rs
@@ -9,6 +9,7 @@ pub mod place;
 use crate::ich::StableHashingContext;
 use crate::ty::query::Providers;
 use crate::ty::TyCtxt;
+use rustc_ast::Attribute;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
@@ -16,6 +17,7 @@ use rustc_hir::def_id::{LocalDefId, LOCAL_CRATE};
 use rustc_hir::*;
 use rustc_index::vec::IndexVec;
 use rustc_span::DUMMY_SP;
+use std::collections::BTreeMap;
 
 #[derive(Debug)]
 pub struct Owner<'tcx> {
@@ -55,6 +57,48 @@ impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for OwnerNodes<'tcx> {
     }
 }
 
+#[derive(Copy, Clone)]
+pub struct AttributeMap<'tcx> {
+    map: &'tcx BTreeMap<HirId, &'tcx [Attribute]>,
+    prefix: LocalDefId,
+}
+
+impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for AttributeMap<'tcx> {
+    fn hash_stable(&self, hcx: &mut StableHashingContext<'a>, hasher: &mut StableHasher) {
+        let range = self.range();
+
+        range.clone().count().hash_stable(hcx, hasher);
+        for (key, value) in range {
+            key.hash_stable(hcx, hasher);
+            value.hash_stable(hcx, hasher);
+        }
+    }
+}
+
+impl<'tcx> std::fmt::Debug for AttributeMap<'tcx> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AttributeMap")
+            .field("prefix", &self.prefix)
+            .field("range", &&self.range().collect::<Vec<_>>()[..])
+            .finish()
+    }
+}
+
+impl<'tcx> AttributeMap<'tcx> {
+    fn get(&self, id: ItemLocalId) -> &'tcx [Attribute] {
+        self.map.get(&HirId { owner: self.prefix, local_id: id }).copied().unwrap_or(&[])
+    }
+
+    fn range(&self) -> std::collections::btree_map::Range<'_, rustc_hir::HirId, &[Attribute]> {
+        let local_zero = ItemLocalId::from_u32(0);
+        let range = HirId { owner: self.prefix, local_id: local_zero }..HirId {
+            owner: LocalDefId { local_def_index: self.prefix.local_def_index + 1 },
+            local_id: local_zero,
+        };
+        self.map.range(range)
+    }
+}
+
 impl<'tcx> TyCtxt<'tcx> {
     #[inline(always)]
     pub fn hir(self) -> map::Map<'tcx> {
@@ -76,7 +120,7 @@ pub fn provide(providers: &mut Providers) {
     providers.hir_module_items = |tcx, id| &tcx.untracked_crate.modules[&id];
     providers.hir_owner = |tcx, id| tcx.index_hir(LOCAL_CRATE).map[id].signature;
     providers.hir_owner_nodes = |tcx, id| tcx.index_hir(LOCAL_CRATE).map[id].with_bodies.as_deref();
-    providers.hir_attrs = |tcx, id| &tcx.index_hir(LOCAL_CRATE).map[id].attrs;
+    providers.hir_attrs = |tcx, id| AttributeMap { map: &tcx.untracked_crate.attrs, prefix: id };
     providers.def_span = |tcx, def_id| tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP);
     providers.fn_arg_names = |tcx, id| {
         let hir = tcx.hir();
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 8edcefe72da..b03b26d6460 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -65,7 +65,7 @@ rustc_queries! {
     ///
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
-    query hir_attrs(key: LocalDefId) -> &'tcx IndexVec<ItemLocalId, &'tcx [ast::Attribute]> {
+    query hir_attrs(key: LocalDefId) -> rustc_middle::hir::AttributeMap<'tcx> {
         eval_always
         desc { |tcx| "HIR owner attributes in `{}`", tcx.def_path_str(key.to_def_id()) }
     }