Collect spans directly into an IndexVec.

Author: cjgillot

src/librustc_ast_lowering/expr.rs

@@ -203,7 +203,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     // Include parens in span, but only if it is a super-span.
                     if e.span.contains(ex.span) {
                         ex.span = e.span;
-                        self.spans.insert(ex.hir_id, e.span);
+                        self.spans[ex.hir_id] = e.span;
                     }
                     // Merge attributes into the inner expression.
                     let mut attrs = e.attrs.clone();
@@ -1365,7 +1365,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let mut head = self.lower_expr_mut(head);
         let desugared_span = self.mark_span_with_reason(DesugaringKind::ForLoop, head.span, None);
         head.span = desugared_span;
-        self.spans.insert(head.hir_id, desugared_span);
+        self.spans[head.hir_id] = desugared_span;
 
         let iter = Ident::with_dummy_span(sym::iter);
 

src/librustc_ast_lowering/lib.rs

@@ -53,8 +53,8 @@ use rustc_hir::def::{DefKind, Namespace, PartialRes, PerNS, Res};
 use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId, CRATE_DEF_INDEX};
 use rustc_hir::definitions::{DefKey, DefPathData, Definitions};
 use rustc_hir::intravisit;
-use rustc_hir::{ConstArg, GenericArg, ParamName};
-use rustc_index::vec::IndexVec;
+use rustc_hir::{ConstArg, GenericArg, HirIdVec, ParamName};
+use rustc_index::vec::{Idx, IndexVec};
 use rustc_session::config::nightly_options;
 use rustc_session::lint::{builtin::BARE_TRAIT_OBJECTS, BuiltinLintDiagnostics, LintBuffer};
 use rustc_session::parse::ParseSess;
@@ -115,7 +115,7 @@ struct LoweringContext<'a, 'hir: 'a> {
     modules: BTreeMap<hir::HirId, hir::ModuleItems>,
 
     /// Collected spans from the AST.
-    spans: BTreeMap<hir::HirId, Span>,
+    spans: HirIdVec<Span>,
 
     generator_kind: Option<hir::GeneratorKind>,
 
@@ -307,7 +307,7 @@ pub fn lower_crate<'a, 'hir>(
         bodies: BTreeMap::new(),
         trait_impls: BTreeMap::new(),
         modules: BTreeMap::new(),
-        spans: BTreeMap::new(),
+        spans: Default::default(),
         exported_macros: Vec::new(),
         non_exported_macro_attrs: Vec::new(),
         catch_scopes: Vec::new(),
@@ -570,6 +570,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         self.resolver.definitions().init_node_id_to_hir_id_mapping(self.node_id_to_hir_id);
 
+        //FIXME(cjgillot) Ideally, each LocalDefId would be a HIR owner.
+        // In the mean time, allocate the missing empty vectors.
+        self.spans.push_owner(Idx::new(self.resolver.definitions().def_index_count() - 1));
+
         hir::Crate {
             item: hir::CrateItem { module, attrs, span: c.span },
             exported_macros: self.arena.alloc_from_iter(self.exported_macros),
@@ -619,22 +623,22 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         }
 
         let hir_id = if let Some(existing_hir_id) = self.node_id_to_hir_id[ast_node_id] {
+            if span != DUMMY_SP {
+                debug_assert!(
+                    self.spans[existing_hir_id] == DUMMY_SP || self.spans[existing_hir_id] == span
+                );
+                self.spans[existing_hir_id] = span;
+            }
             existing_hir_id
         } else {
             // Generate a new `HirId`.
             let hir_id = alloc_hir_id(self);
             self.node_id_to_hir_id[ast_node_id] = Some(hir_id);
+            self.spans.push(hir_id, span);
 
             hir_id
         };
 
-        let stored_span = self.spans.entry(hir_id).or_insert(span);
-        if *stored_span == DUMMY_SP {
-            *stored_span = span;
-        } else if span != DUMMY_SP {
-            assert_eq!(*stored_span, span);
-        }
-
         hir_id
     }
 

src/librustc_hir/hir.rs

@@ -1,6 +1,7 @@
 use crate::def::{DefKind, Namespace, Res};
 use crate::def_id::DefId;
 crate use crate::hir_id::HirId;
+use crate::hir_id::HirIdVec;
 use crate::itemlikevisit;
 
 use rustc_ast::ast::{self, CrateSugar, LlvmAsmDialect};
@@ -643,7 +644,7 @@ pub struct Crate<'hir> {
     pub trait_map: BTreeMap<HirId, Vec<TraitCandidate>>,
 
     /// Collected spans from the AST.
-    pub spans: BTreeMap<HirId, Span>,
+    pub spans: HirIdVec<Span>,
 }
 
 impl Crate<'hir> {