3 files changed, 39 insertions, 29 deletions

diff --git a/crates/ra_hir_def/src/adt.rs b/crates/ra_hir_def/src/adt.rs
index 20e9a1eb5..a43f553aa 100644
--- a/crates/ra_hir_def/src/adt.rs
+++ b/crates/ra_hir_def/src/adt.rs
@@ -95,7 +95,7 @@ fn lower_enum(
                 name: var.name().map(|it| it.as_name()),
                 variant_data: Arc::new(VariantData::new(var.kind())),
             },
-        )
+        );
     }
 }
 
@@ -160,7 +160,7 @@ fn lower_struct(
                         name: Name::new_tuple_field(i),
                         type_ref: TypeRef::from_ast_opt(fd.type_ref()),
                     },
-                )
+                );
             }
             StructKind::Tuple
         }
@@ -172,7 +172,7 @@ fn lower_struct(
                         name: fd.name().map(|n| n.as_name()).unwrap_or_else(Name::missing),
                         type_ref: TypeRef::from_ast_opt(fd.ascribed_type()),
                     },
-                )
+                );
             }
             StructKind::Record
         }
diff --git a/crates/ra_hir_def/src/nameres/raw.rs b/crates/ra_hir_def/src/nameres/raw.rs
index 2ec84f2cc..401af031c 100644
--- a/crates/ra_hir_def/src/nameres/raw.rs
+++ b/crates/ra_hir_def/src/nameres/raw.rs
@@ -22,7 +22,8 @@ use ra_syntax::{
 use test_utils::tested_by;
 
 use crate::{
-    attr::Attrs, db::DefDatabase, path::Path, FileAstId, HirFileId, LocalImportId, Source,
+    attr::Attrs, db::DefDatabase, path::Path, trace::Trace, FileAstId, HirFileId, LocalImportId,
+    Source,
 };
 
 /// `RawItems` is a set of top-level items in a file (except for impls).
@@ -48,10 +49,6 @@ pub struct ImportSourceMap {
 type ImportSourcePtr = Either<AstPtr<ast::UseTree>, AstPtr<ast::ExternCrateItem>>;
 
 impl ImportSourceMap {
-    fn insert(&mut self, import: LocalImportId, ptr: ImportSourcePtr) {
-        self.map.insert(import, ptr)
-    }
-
     pub fn get(&self, import: LocalImportId) -> ImportSourcePtr {
         self.map[import].clone()
     }
@@ -72,7 +69,7 @@ impl RawItems {
         let mut collector = RawItemsCollector {
             raw_items: RawItems::default(),
             source_ast_id_map: db.ast_id_map(file_id),
-            source_map: ImportSourceMap::default(),
+            imports: Trace::new(),
             file_id,
             hygiene: Hygiene::new(db, file_id),
         };
@@ -83,7 +80,11 @@ impl RawItems {
                 collector.process_module(None, item_list);
             }
         }
-        (Arc::new(collector.raw_items), Arc::new(collector.source_map))
+        let mut raw_items = collector.raw_items;
+        let (arena, map) = collector.imports.into_arena_and_map();
+        raw_items.imports = arena;
+        let source_map = ImportSourceMap { map };
+        (Arc::new(raw_items), Arc::new(source_map))
     }
 
     pub(super) fn items(&self) -> &[RawItem] {
@@ -207,8 +208,8 @@ pub(super) struct ImplData {
 
 struct RawItemsCollector {
     raw_items: RawItems,
+    imports: Trace<LocalImportId, ImportData, ImportSourcePtr>,
     source_ast_id_map: Arc<AstIdMap>,
-    source_map: ImportSourceMap,
     file_id: HirFileId,
     hygiene: Hygiene,
 }
@@ -392,8 +393,7 @@ impl RawItemsCollector {
         data: ImportData,
         source: ImportSourcePtr,
     ) {
-        let import = self.raw_items.imports.alloc(data);
-        self.source_map.insert(import, source);
+        let import = self.imports.alloc(|| source, || data);
         self.push_item(current_module, attrs, RawItemKind::Import(import))
     }
 
diff --git a/crates/ra_hir_def/src/trace.rs b/crates/ra_hir_def/src/trace.rs
index fc26f5a48..2bcd707bc 100644
--- a/crates/ra_hir_def/src/trace.rs
+++ b/crates/ra_hir_def/src/trace.rs
@@ -12,38 +12,48 @@
 use ra_arena::{map::ArenaMap, Arena, ArenaId, RawId};
 
 pub(crate) struct Trace<ID: ArenaId, T, V> {
-    for_arena: bool,
-    arena: Arena<ID, T>,
-    map: ArenaMap<ID, V>,
+    arena: Option<Arena<ID, T>>,
+    map: Option<ArenaMap<ID, V>>,
     len: u32,
 }
 
-impl<ID: ra_arena::ArenaId, T, V> Trace<ID, T, V> {
+impl<ID: ra_arena::ArenaId + Copy, T, V> Trace<ID, T, V> {
+    pub(crate) fn new() -> Trace<ID, T, V> {
+        Trace { arena: Some(Arena::default()), map: Some(ArenaMap::default()), len: 0 }
+    }
+
     pub(crate) fn new_for_arena() -> Trace<ID, T, V> {
-        Trace { for_arena: true, arena: Arena::default(), map: ArenaMap::default(), len: 0 }
+        Trace { arena: Some(Arena::default()), map: None, len: 0 }
     }
 
     pub(crate) fn new_for_map() -> Trace<ID, T, V> {
-        Trace { for_arena: false, arena: Arena::default(), map: ArenaMap::default(), len: 0 }
+        Trace { arena: None, map: Some(ArenaMap::default()), len: 0 }
     }
 
-    pub(crate) fn alloc(&mut self, value: impl Fn() -> V, data: impl Fn() -> T) {
-        if self.for_arena {
-            self.arena.alloc(data());
+    pub(crate) fn alloc(&mut self, value: impl FnOnce() -> V, data: impl FnOnce() -> T) -> ID {
+        let id = if let Some(arena) = &mut self.arena {
+            arena.alloc(data())
         } else {
             let id = ID::from_raw(RawId::from(self.len));
             self.len += 1;
-            self.map.insert(id, value());
+            id
+        };
+
+        if let Some(map) = &mut self.map {
+            map.insert(id, value());
         }
+        id
+    }
+
+    pub(crate) fn into_arena(mut self) -> Arena<ID, T> {
+        self.arena.take().unwrap()
     }
 
-    pub(crate) fn into_arena(self) -> Arena<ID, T> {
-        assert!(self.for_arena);
-        self.arena
+    pub(crate) fn into_map(mut self) -> ArenaMap<ID, V> {
+        self.map.take().unwrap()
     }
 
-    pub(crate) fn into_map(self) -> ArenaMap<ID, V> {
-        assert!(!self.for_arena);
-        self.map
+    pub(crate) fn into_arena_and_map(mut self) -> (Arena<ID, T>, ArenaMap<ID, V>) {
+        (self.arena.take().unwrap(), self.map.take().unwrap())
     }
 }