1 files changed, 19 insertions, 12 deletions

diff --git a/crates/ra_hir_def/src/src.rs b/crates/ra_hir_def/src/src.rs
index 46e90da70..043b93fad 100644
--- a/crates/ra_hir_def/src/src.rs
+++ b/crates/ra_hir_def/src/src.rs
@@ -2,30 +2,37 @@
 
 use hir_expand::InFile;
 use ra_arena::map::ArenaMap;
-use ra_syntax::AstNode;
 
-use crate::{db::DefDatabase, AssocItemLoc, ItemLoc};
+use crate::{db::DefDatabase, item_tree::ItemTreeNode, AssocItemLoc, ItemLoc};
 
 pub trait HasSource {
     type Value;
     fn source(&self, db: &dyn DefDatabase) -> InFile<Self::Value>;
 }
 
-impl<N: AstNode> HasSource for AssocItemLoc<N> {
-    type Value = N;
+impl<N: ItemTreeNode> HasSource for AssocItemLoc<N> {
+    type Value = N::Source;
 
-    fn source(&self, db: &dyn DefDatabase) -> InFile<N> {
-        let node = self.ast_id.to_node(db.upcast());
-        InFile::new(self.ast_id.file_id, node)
+    fn source(&self, db: &dyn DefDatabase) -> InFile<N::Source> {
+        let tree = db.item_tree(self.id.file_id);
+        let ast_id_map = db.ast_id_map(self.id.file_id);
+        let root = db.parse_or_expand(self.id.file_id).unwrap();
+        let node = &tree[self.id.value];
+
+        InFile::new(self.id.file_id, ast_id_map.get(node.ast_id()).to_node(&root))
     }
 }
 
-impl<N: AstNode> HasSource for ItemLoc<N> {
-    type Value = N;
+impl<N: ItemTreeNode> HasSource for ItemLoc<N> {
+    type Value = N::Source;
+
+    fn source(&self, db: &dyn DefDatabase) -> InFile<N::Source> {
+        let tree = db.item_tree(self.id.file_id);
+        let ast_id_map = db.ast_id_map(self.id.file_id);
+        let root = db.parse_or_expand(self.id.file_id).unwrap();
+        let node = &tree[self.id.value];
 
-    fn source(&self, db: &dyn DefDatabase) -> InFile<N> {
-        let node = self.ast_id.to_node(db.upcast());
-        InFile::new(self.ast_id.file_id, node)
+        InFile::new(self.id.file_id, ast_id_map.get(node.ast_id()).to_node(&root))
     }
 }