3 files changed, 40 insertions, 9 deletions

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};