4 files changed, 152 insertions, 143 deletions

diff --git a/crates/ra_hir/src/module_tree.rs b/crates/ra_hir/src/module_tree.rs
index 47c14af35..b201bf69b 100644
--- a/crates/ra_hir/src/module_tree.rs
+++ b/crates/ra_hir/src/module_tree.rs
@@ -62,14 +62,26 @@ impl Submodule {
             file_items: &SourceFileItems,
             root: &impl ast::ModuleItemOwner,
         ) -> Vec<Submodule> {
-            modules(root)
-                .map(|(name, m)| Submodule {
-                    name,
-                    is_declaration: m.has_semi(),
-                    source: SourceItemId {
-                        file_id,
-                        item_id: Some(file_items.id_of(file_id, m.syntax())),
-                    },
+            root.items()
+                .filter_map(|item| match item.kind() {
+                    ast::ModuleItemKind::Module(m) => Some(m),
+                    _ => None,
+                })
+                .filter_map(|module| {
+                    let name = module.name()?.as_name();
+                    if !module.has_semi() && module.item_list().is_none() {
+                        tested_by!(name_res_works_for_broken_modules);
+                        return None;
+                    }
+                    let sub = Submodule {
+                        name,
+                        is_declaration: module.has_semi(),
+                        source: SourceItemId {
+                            file_id,
+                            item_id: Some(file_items.id_of(file_id, module.syntax())),
+                        },
+                    };
+                    Some(sub)
                 })
                 .collect()
         }
@@ -119,7 +131,8 @@ impl ModuleTree {
         source_root: SourceRootId,
     ) -> Arc<ModuleTree> {
         db.check_canceled();
-        let res = create_module_tree(db, source_root);
+        let mut res = ModuleTree::default();
+        res.init(db, source_root);
         Arc::new(res)
     }
 
@@ -131,6 +144,96 @@ impl ModuleTree {
         let (res, _) = self.mods.iter().find(|(_, m)| m.source == source)?;
         Some(res)
     }
+
+    fn init(&mut self, db: &impl HirDatabase, source_root: SourceRootId) {
+        let mut roots = FxHashMap::default();
+        let mut visited = FxHashSet::default();
+
+        let source_root = db.source_root(source_root);
+        for &file_id in source_root.files.values() {
+            let source = SourceItemId {
+                file_id: file_id.into(),
+                item_id: None,
+            };
+            if visited.contains(&source) {
+                continue; // TODO: use explicit crate_roots here
+            }
+            assert!(!roots.contains_key(&file_id));
+            let module_id =
+                self.init_subtree(db, &source_root, &mut visited, &mut roots, None, source);
+            roots.insert(file_id, module_id);
+        }
+    }
+
+    fn init_subtree(
+        &mut self,
+        db: &impl HirDatabase,
+        source_root: &SourceRoot,
+        visited: &mut FxHashSet<SourceItemId>,
+        roots: &mut FxHashMap<FileId, ModuleId>,
+        parent: Option<LinkId>,
+        source: SourceItemId,
+    ) -> ModuleId {
+        visited.insert(source);
+        let id = self.alloc_mod(ModuleData {
+            source,
+            parent,
+            children: Vec::new(),
+        });
+        for sub in db.submodules(source).iter() {
+            let link = self.alloc_link(LinkData {
+                source: sub.source,
+                name: sub.name.clone(),
+                owner: id,
+                points_to: Vec::new(),
+                problem: None,
+            });
+
+            let (points_to, problem) = if sub.is_declaration {
+                let (points_to, problem) = resolve_submodule(db, source.file_id, &sub.name);
+                let points_to = points_to
+                    .into_iter()
+                    .map(|file_id| match roots.remove(&file_id) {
+                        Some(module_id) => {
+                            self.mods[module_id].parent = Some(link);
+                            module_id
+                        }
+                        None => self.init_subtree(
+                            db,
+                            source_root,
+                            visited,
+                            roots,
+                            Some(link),
+                            SourceItemId {
+                                file_id: file_id.into(),
+                                item_id: None,
+                            },
+                        ),
+                    })
+                    .collect::<Vec<_>>();
+                (points_to, problem)
+            } else {
+                let points_to =
+                    self.init_subtree(db, source_root, visited, roots, Some(link), sub.source);
+                (vec![points_to], None)
+            };
+
+            self.links[link].points_to = points_to;
+            self.links[link].problem = problem;
+        }
+        id
+    }
+
+    fn alloc_mod(&mut self, data: ModuleData) -> ModuleId {
+        self.mods.alloc(data)
+    }
+
+    fn alloc_link(&mut self, data: LinkData) -> LinkId {
+        let owner = data.owner;
+        let id = self.links.alloc(data);
+        self.mods[owner].children.push(id);
+        id
+    }
 }
 
 impl ModuleId {
@@ -198,131 +301,6 @@ impl LinkId {
     }
 }
 
-impl ModuleTree {
-    fn push_mod(&mut self, data: ModuleData) -> ModuleId {
-        self.mods.alloc(data)
-    }
-    fn push_link(&mut self, data: LinkData) -> LinkId {
-        let owner = data.owner;
-        let id = self.links.alloc(data);
-        self.mods[owner].children.push(id);
-        id
-    }
-}
-
-fn modules(root: &impl ast::ModuleItemOwner) -> impl Iterator<Item = (Name, &ast::Module)> {
-    root.items()
-        .filter_map(|item| match item.kind() {
-            ast::ModuleItemKind::Module(m) => Some(m),
-            _ => None,
-        })
-        .filter_map(|module| {
-            let name = module.name()?.as_name();
-            if !module.has_semi() && module.item_list().is_none() {
-                tested_by!(name_res_works_for_broken_modules);
-                return None;
-            }
-            Some((name, module))
-        })
-}
-
-fn create_module_tree<'a>(db: &impl HirDatabase, source_root: SourceRootId) -> ModuleTree {
-    let mut tree = ModuleTree::default();
-
-    let mut roots = FxHashMap::default();
-    let mut visited = FxHashSet::default();
-
-    let source_root = db.source_root(source_root);
-    for &file_id in source_root.files.values() {
-        let source = SourceItemId {
-            file_id: file_id.into(),
-            item_id: None,
-        };
-        if visited.contains(&source) {
-            continue; // TODO: use explicit crate_roots here
-        }
-        assert!(!roots.contains_key(&file_id));
-        let module_id = build_subtree(
-            db,
-            &source_root,
-            &mut tree,
-            &mut visited,
-            &mut roots,
-            None,
-            source,
-        );
-        roots.insert(file_id, module_id);
-    }
-    tree
-}
-
-fn build_subtree(
-    db: &impl HirDatabase,
-    source_root: &SourceRoot,
-    tree: &mut ModuleTree,
-    visited: &mut FxHashSet<SourceItemId>,
-    roots: &mut FxHashMap<FileId, ModuleId>,
-    parent: Option<LinkId>,
-    source: SourceItemId,
-) -> ModuleId {
-    visited.insert(source);
-    let id = tree.push_mod(ModuleData {
-        source,
-        parent,
-        children: Vec::new(),
-    });
-    for sub in db.submodules(source).iter() {
-        let link = tree.push_link(LinkData {
-            source: sub.source,
-            name: sub.name.clone(),
-            owner: id,
-            points_to: Vec::new(),
-            problem: None,
-        });
-
-        let (points_to, problem) = if sub.is_declaration {
-            let (points_to, problem) = resolve_submodule(db, source.file_id, &sub.name);
-            let points_to = points_to
-                .into_iter()
-                .map(|file_id| match roots.remove(&file_id) {
-                    Some(module_id) => {
-                        tree.mods[module_id].parent = Some(link);
-                        module_id
-                    }
-                    None => build_subtree(
-                        db,
-                        source_root,
-                        tree,
-                        visited,
-                        roots,
-                        Some(link),
-                        SourceItemId {
-                            file_id: file_id.into(),
-                            item_id: None,
-                        },
-                    ),
-                })
-                .collect::<Vec<_>>();
-            (points_to, problem)
-        } else {
-            let points_to = build_subtree(
-                db,
-                source_root,
-                tree,
-                visited,
-                roots,
-                Some(link),
-                sub.source,
-            );
-            (vec![points_to], None)
-        };
-
-        tree.links[link].points_to = points_to;
-        tree.links[link].problem = problem;
-    }
-    id
-}
-
 fn resolve_submodule(
     db: &impl HirDatabase,
     file_id: HirFileId,
diff --git a/crates/ra_hir/src/nameres/lower.rs b/crates/ra_hir/src/nameres/lower.rs
index ab6f3a9bc..921ba3c98 100644
--- a/crates/ra_hir/src/nameres/lower.rs
+++ b/crates/ra_hir/src/nameres/lower.rs
@@ -1,10 +1,10 @@
 use std::sync::Arc;
 
 use ra_syntax::{
-    SyntaxKind, AstNode, SourceFile, TreeArc, SyntaxNodePtr,
+    SyntaxKind, AstNode, SourceFile, TreeArc, AstPtr,
     ast::{self, ModuleItemOwner},
 };
-use ra_db::{SourceRootId};
+use ra_db::SourceRootId;
 use ra_arena::{Arena, RawId, impl_arena_id, map::ArenaMap};
 
 use crate::{
@@ -72,13 +72,12 @@ pub struct LoweredModule {
 
 #[derive(Debug, Default, PartialEq, Eq)]
 pub struct ImportSourceMap {
-    map: ArenaMap<ImportId, SyntaxNodePtr>,
+    map: ArenaMap<ImportId, AstPtr<ast::PathSegment>>,
 }
 
 impl ImportSourceMap {
     fn insert(&mut self, import: ImportId, segment: &ast::PathSegment) {
-        self.map
-            .insert(import, SyntaxNodePtr::new(segment.syntax()))
+        self.map.insert(import, AstPtr::new(segment))
     }
 
     pub fn get(&self, source: &ModuleSource, import: ImportId) -> TreeArc<ast::PathSegment> {
@@ -87,9 +86,7 @@ impl ImportSourceMap {
             ModuleSource::Module(m) => m.syntax().ancestors().find_map(SourceFile::cast).unwrap(),
         };
 
-        ast::PathSegment::cast(self.map[import].to_node(file))
-            .unwrap()
-            .to_owned()
+        self.map[import].to_node(file).to_owned()
     }
 }
 
diff --git a/crates/ra_syntax/src/lib.rs b/crates/ra_syntax/src/lib.rs
index 97b196118..104f32851 100644
--- a/crates/ra_syntax/src/lib.rs
+++ b/crates/ra_syntax/src/lib.rs
@@ -43,7 +43,7 @@ pub use crate::{
     lexer::{tokenize, Token},
     syntax_kinds::SyntaxKind,
     yellow::{Direction, SyntaxError, SyntaxNode, WalkEvent, Location, TreeArc},
-    ptr::SyntaxNodePtr,
+    ptr::{SyntaxNodePtr, AstPtr},
 };
 
 use ra_text_edit::AtomTextEdit;
diff --git a/crates/ra_syntax/src/ptr.rs b/crates/ra_syntax/src/ptr.rs
index e8c40e5d3..b50cd8a52 100644
--- a/crates/ra_syntax/src/ptr.rs
+++ b/crates/ra_syntax/src/ptr.rs
@@ -1,3 +1,5 @@
+use std::marker::PhantomData;
+
 use crate::{
     AstNode, SourceFile, SyntaxKind, SyntaxNode, TextRange,
     algo::generate,
@@ -37,6 +39,38 @@ impl SyntaxNodePtr {
     }
 }
 
+/// Like `SyntaxNodePtr`, but remembers the type of node
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct AstPtr<N: AstNode> {
+    ptr: SyntaxNodePtr,
+    _ty: PhantomData<N>,
+}
+
+impl<N: AstNode> Copy for AstPtr<N> {}
+impl<N: AstNode> Clone for AstPtr<N> {
+    fn clone(&self) -> AstPtr<N> {
+        *self
+    }
+}
+
+impl<N: AstNode> AstPtr<N> {
+    pub fn new(node: &N) -> AstPtr<N> {
+        AstPtr {
+            ptr: SyntaxNodePtr::new(node.syntax()),
+            _ty: PhantomData,
+        }
+    }
+
+    pub fn to_node(self, source_file: &SourceFile) -> &N {
+        let syntax_node = self.ptr.to_node(source_file);
+        N::cast(syntax_node).unwrap()
+    }
+
+    pub fn syntax_node_ptr(self) -> SyntaxNodePtr {
+        self.ptr
+    }
+}
+
 #[test]
 fn test_local_syntax_ptr() {
     use crate::{ast, AstNode};