1 files changed, 52 insertions, 16 deletions

diff --git a/crates/ra_hir_def/src/nameres/raw.rs b/crates/ra_hir_def/src/nameres/raw.rs
index b10e458a2..ecb4d7c03 100644
--- a/crates/ra_hir_def/src/nameres/raw.rs
+++ b/crates/ra_hir_def/src/nameres/raw.rs
@@ -7,24 +7,24 @@
 
 use std::{ops::Index, sync::Arc};
 
+use either::Either;
 use hir_expand::{
     ast_id_map::AstIdMap,
     db::AstDatabase,
     hygiene::Hygiene,
     name::{AsName, Name},
 };
-use ra_arena::{impl_arena_id, Arena, RawId};
+use ra_arena::{impl_arena_id, map::ArenaMap, Arena, RawId};
 use ra_syntax::{
     ast::{self, AttrsOwner, NameOwner},
-    AstNode,
+    AstNode, AstPtr,
 };
 use test_utils::tested_by;
 
-use crate::{attr::Attrs, db::DefDatabase, path::ModPath, FileAstId, HirFileId, InFile};
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub(super) struct LocalImportId(RawId);
-impl_arena_id!(LocalImportId);
+use crate::{
+    attr::Attrs, db::DefDatabase, path::ModPath, trace::Trace, FileAstId, HirFileId, InFile,
+    LocalImportId,
+};
 
 /// `RawItems` is a set of top-level items in a file (except for impls).
 ///
@@ -41,14 +41,35 @@ pub struct RawItems {
     items: Vec<RawItem>,
 }
 
+#[derive(Debug, Default, PartialEq, Eq)]
+pub struct ImportSourceMap {
+    map: ArenaMap<LocalImportId, ImportSourcePtr>,
+}
+
+type ImportSourcePtr = Either<AstPtr<ast::UseTree>, AstPtr<ast::ExternCrateItem>>;
+
+impl ImportSourceMap {
+    pub fn get(&self, import: LocalImportId) -> ImportSourcePtr {
+        self.map[import].clone()
+    }
+}
+
 impl RawItems {
     pub(crate) fn raw_items_query(
         db: &(impl DefDatabase + AstDatabase),
         file_id: HirFileId,
     ) -> Arc<RawItems> {
+        db.raw_items_with_source_map(file_id).0
+    }
+
+    pub(crate) fn raw_items_with_source_map_query(
+        db: &(impl DefDatabase + AstDatabase),
+        file_id: HirFileId,
+    ) -> (Arc<RawItems>, Arc<ImportSourceMap>) {
         let mut collector = RawItemsCollector {
             raw_items: RawItems::default(),
             source_ast_id_map: db.ast_id_map(file_id),
+            imports: Trace::new(),
             file_id,
             hygiene: Hygiene::new(db, file_id),
         };
@@ -59,8 +80,11 @@ impl RawItems {
                 collector.process_module(None, item_list);
             }
         }
-        let raw_items = collector.raw_items;
-        Arc::new(raw_items)
+        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] {
@@ -199,6 +223,7 @@ pub(super) struct ImplData {
 
 struct RawItemsCollector {
     raw_items: RawItems,
+    imports: Trace<LocalImportId, ImportData, ImportSourcePtr>,
     source_ast_id_map: Arc<AstIdMap>,
     file_id: HirFileId,
     hygiene: Hygiene,
@@ -305,7 +330,7 @@ impl RawItemsCollector {
         ModPath::expand_use_item(
             InFile { value: use_item, file_id: self.file_id },
             &self.hygiene,
-            |path, _use_tree, is_glob, alias| {
+            |path, use_tree, is_glob, alias| {
                 let import_data = ImportData {
                     path,
                     alias,
@@ -314,11 +339,11 @@ impl RawItemsCollector {
                     is_extern_crate: false,
                     is_macro_use: false,
                 };
-                buf.push(import_data);
+                buf.push((import_data, Either::Left(AstPtr::new(use_tree))));
             },
         );
-        for import_data in buf {
-            self.push_import(current_module, attrs.clone(), import_data);
+        for (import_data, ptr) in buf {
+            self.push_import(current_module, attrs.clone(), import_data, ptr);
         }
     }
 
@@ -341,7 +366,12 @@ impl RawItemsCollector {
                 is_extern_crate: true,
                 is_macro_use,
             };
-            self.push_import(current_module, attrs, import_data);
+            self.push_import(
+                current_module,
+                attrs,
+                import_data,
+                Either::Right(AstPtr::new(&extern_crate)),
+            );
         }
     }
 
@@ -372,8 +402,14 @@ impl RawItemsCollector {
         self.push_item(current_module, attrs, RawItemKind::Impl(imp))
     }
 
-    fn push_import(&mut self, current_module: Option<Module>, attrs: Attrs, data: ImportData) {
-        let import = self.raw_items.imports.alloc(data);
+    fn push_import(
+        &mut self,
+        current_module: Option<Module>,
+        attrs: Attrs,
+        data: ImportData,
+        source: ImportSourcePtr,
+    ) {
+        let import = self.imports.alloc(|| source, || data);
         self.push_item(current_module, attrs, RawItemKind::Import(import))
     }