9 files changed, 74 insertions, 65 deletions

diff --git a/crates/ra_hir/src/code_model_api.rs b/crates/ra_hir/src/code_model_api.rs
index 9e6170440..f13a6b37a 100644
--- a/crates/ra_hir/src/code_model_api.rs
+++ b/crates/ra_hir/src/code_model_api.rs
@@ -4,7 +4,7 @@ use ra_db::{CrateId, SourceRootId, Edition};
 use ra_syntax::{ast::self, TreeArc};
 
 use crate::{
-    Name, ScopesWithSourceMap, Ty, HirFileId, ImportSource,
+    Name, ScopesWithSourceMap, Ty, HirFileId, Either,
     HirDatabase, DefDatabase,
     type_ref::TypeRef,
     nameres::{ModuleScope, Namespace, ImportId, CrateModuleId},
@@ -117,8 +117,14 @@ impl Module {
     }
 
     /// Returns the syntax of the last path segment corresponding to this import
-    pub fn import_source(&self, db: &impl HirDatabase, import: ImportId) -> ImportSource {
-        self.import_source_impl(db, import)
+    pub fn import_source(
+        &self,
+        db: &impl HirDatabase,
+        import: ImportId,
+    ) -> Either<TreeArc<ast::UseTree>, TreeArc<ast::ExternCrateItem>> {
+        let (file_id, source) = self.definition_source(db);
+        let (_, source_map) = db.raw_items_with_source_map(file_id);
+        source_map.get(&source, import)
     }
 
     /// Returns the crate this module is part of.
diff --git a/crates/ra_hir/src/code_model_impl/module.rs b/crates/ra_hir/src/code_model_impl/module.rs
index 88dee3a69..5c2ea73ce 100644
--- a/crates/ra_hir/src/code_model_impl/module.rs
+++ b/crates/ra_hir/src/code_model_impl/module.rs
@@ -3,9 +3,9 @@ use ra_syntax::{ast, TreeArc};
 
 use crate::{
     Module, ModuleSource, Name, AstId,
-    nameres::{CrateModuleId, ImportId},
+    nameres::CrateModuleId,
     HirDatabase, DefDatabase,
-    HirFileId, ImportSource,
+    HirFileId,
 };
 
 impl ModuleSource {
@@ -68,16 +68,6 @@ impl Module {
         Some((decl.file_id(), ast))
     }
 
-    pub(crate) fn import_source_impl(
-        &self,
-        db: &impl HirDatabase,
-        import: ImportId,
-    ) -> ImportSource {
-        let (file_id, source) = self.definition_source(db);
-        let (_, source_map) = db.raw_items_with_source_map(file_id);
-        source_map.get(&source, import)
-    }
-
     pub(crate) fn crate_root_impl(&self, db: &impl DefDatabase) -> Module {
         let def_map = db.crate_def_map(self.krate);
         self.with_module_id(def_map.root())
diff --git a/crates/ra_hir/src/either.rs b/crates/ra_hir/src/either.rs
new file mode 100644
index 000000000..4073cc82e
--- /dev/null
+++ b/crates/ra_hir/src/either.rs
@@ -0,0 +1,28 @@
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum Either<A, B> {
+    A(A),
+    B(B),
+}
+
+impl<A, B> Either<A, B> {
+    pub fn either<R, F1, F2>(self, f1: F1, f2: F2) -> R
+    where
+        F1: FnOnce(A) -> R,
+        F2: FnOnce(B) -> R,
+    {
+        match self {
+            Either::A(a) => f1(a),
+            Either::B(b) => f2(b),
+        }
+    }
+    pub fn map<U, V, F1, F2>(self, f1: F1, f2: F2) -> Either<U, V>
+    where
+        F1: FnOnce(A) -> U,
+        F2: FnOnce(B) -> V,
+    {
+        match self {
+            Either::A(a) => Either::A(f1(a)),
+            Either::B(b) => Either::B(f2(b)),
+        }
+    }
+}
diff --git a/crates/ra_hir/src/expr.rs b/crates/ra_hir/src/expr.rs
index 589a9b2db..a2840c15d 100644
--- a/crates/ra_hir/src/expr.rs
+++ b/crates/ra_hir/src/expr.rs
@@ -10,7 +10,7 @@ use ra_syntax::{
 };
 
 use crate::{
-    Path, Name, HirDatabase, Resolver,DefWithBody,
+    Path, Name, HirDatabase, Resolver,DefWithBody, Either,
     name::AsName,
     type_ref::{Mutability, TypeRef},
 };
@@ -51,11 +51,13 @@ pub struct Body {
 pub struct BodySourceMap {
     expr_map: FxHashMap<SyntaxNodePtr, ExprId>,
     expr_map_back: ArenaMap<ExprId, SyntaxNodePtr>,
-    pat_map: FxHashMap<SyntaxNodePtr, PatId>,
-    pat_map_back: ArenaMap<PatId, SyntaxNodePtr>,
+    pat_map: FxHashMap<PatPrr, PatId>,
+    pat_map_back: ArenaMap<PatId, PatPrr>,
     field_map: FxHashMap<(ExprId, usize), AstPtr<ast::NamedField>>,
 }
 
+type PatPrr = Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>;
+
 impl Body {
     pub fn params(&self) -> &[PatId] {
         &self.params
@@ -127,16 +129,16 @@ impl BodySourceMap {
         self.expr_map.get(&SyntaxNodePtr::new(node.syntax())).cloned()
     }
 
-    pub fn pat_syntax(&self, pat: PatId) -> Option<SyntaxNodePtr> {
+    pub fn pat_syntax(&self, pat: PatId) -> Option<PatPrr> {
         self.pat_map_back.get(pat).cloned()
     }
 
-    pub fn syntax_pat(&self, ptr: SyntaxNodePtr) -> Option<PatId> {
+    pub fn syntax_pat(&self, ptr: PatPrr) -> Option<PatId> {
         self.pat_map.get(&ptr).cloned()
     }
 
     pub fn node_pat(&self, node: &ast::Pat) -> Option<PatId> {
-        self.pat_map.get(&SyntaxNodePtr::new(node.syntax())).cloned()
+        self.pat_map.get(&Either::A(AstPtr::new(node))).cloned()
     }
 
     pub fn field_syntax(&self, expr: ExprId, field: usize) -> AstPtr<ast::NamedField> {
@@ -504,10 +506,10 @@ impl ExprCollector {
         id
     }
 
-    fn alloc_pat(&mut self, pat: Pat, syntax_ptr: SyntaxNodePtr) -> PatId {
+    fn alloc_pat(&mut self, pat: Pat, ptr: PatPrr) -> PatId {
         let id = self.pats.alloc(pat);
-        self.source_map.pat_map.insert(syntax_ptr, id);
-        self.source_map.pat_map_back.insert(id, syntax_ptr);
+        self.source_map.pat_map.insert(ptr, id);
+        self.source_map.pat_map_back.insert(id, ptr);
         id
     }
 
@@ -886,8 +888,8 @@ impl ExprCollector {
             ast::PatKind::LiteralPat(_) => Pat::Missing,
             ast::PatKind::SlicePat(_) | ast::PatKind::RangePat(_) => Pat::Missing,
         };
-        let syntax_ptr = SyntaxNodePtr::new(pat.syntax());
-        self.alloc_pat(pattern, syntax_ptr)
+        let ptr = AstPtr::new(pat);
+        self.alloc_pat(pattern, Either::A(ptr))
     }
 
     fn collect_pat_opt(&mut self, pat: Option<&ast::Pat>) -> PatId {
@@ -911,14 +913,14 @@ impl ExprCollector {
     fn collect_fn_body(&mut self, node: &ast::FnDef) {
         if let Some(param_list) = node.param_list() {
             if let Some(self_param) = param_list.self_param() {
-                let self_param = SyntaxNodePtr::new(self_param.syntax());
+                let ptr = AstPtr::new(self_param);
                 let param_pat = self.alloc_pat(
                     Pat::Bind {
                         name: Name::self_param(),
                         mode: BindingAnnotation::Unannotated,
                         subpat: None,
                     },
-                    self_param,
+                    Either::B(ptr),
                 );
                 self.params.push(param_pat);
             }
diff --git a/crates/ra_hir/src/expr/scope.rs b/crates/ra_hir/src/expr/scope.rs
index f1e6e0f02..725b6c00e 100644
--- a/crates/ra_hir/src/expr/scope.rs
+++ b/crates/ra_hir/src/expr/scope.rs
@@ -3,14 +3,14 @@ use std::sync::Arc;
 use rustc_hash::{FxHashMap, FxHashSet};
 
 use ra_syntax::{
-    AstNode, SyntaxNode, TextUnit, TextRange, SyntaxNodePtr,
+    AstNode, SyntaxNode, TextUnit, TextRange, SyntaxNodePtr, AstPtr,
     algo::generate,
     ast,
 };
 use ra_arena::{Arena, RawId, impl_arena_id};
 
 use crate::{
-    Name, AsName,DefWithBody,
+    Name, AsName,DefWithBody, Either,
     expr::{PatId, ExprId, Pat, Expr, Body, Statement, BodySourceMap},
     HirDatabase,
 };
@@ -116,7 +116,7 @@ pub struct ScopesWithSourceMap {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ScopeEntryWithSyntax {
     name: Name,
-    ptr: SyntaxNodePtr,
+    ptr: Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>,
 }
 
 impl ScopeEntryWithSyntax {
@@ -124,7 +124,7 @@ impl ScopeEntryWithSyntax {
         &self.name
     }
 
-    pub fn ptr(&self) -> SyntaxNodePtr {
+    pub fn ptr(&self) -> Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>> {
         self.ptr
     }
 }
@@ -192,14 +192,14 @@ impl ScopesWithSourceMap {
 
     pub fn find_all_refs(&self, pat: &ast::BindPat) -> Vec<ReferenceDescriptor> {
         let fn_def = pat.syntax().ancestors().find_map(ast::FnDef::cast).unwrap();
-        let name_ptr = SyntaxNodePtr::new(pat.syntax());
+        let ptr = Either::A(AstPtr::new(pat.into()));
         fn_def
             .syntax()
             .descendants()
             .filter_map(ast::NameRef::cast)
             .filter(|name_ref| match self.resolve_local_name(*name_ref) {
                 None => false,
-                Some(entry) => entry.ptr() == name_ptr,
+                Some(entry) => entry.ptr() == ptr,
             })
             .map(|name_ref| ReferenceDescriptor {
                 name: name_ref.syntax().text().to_string(),
@@ -429,7 +429,8 @@ mod tests {
         let scopes =
             ScopesWithSourceMap { scopes: Arc::new(scopes), source_map: Arc::new(source_map) };
         let local_name_entry = scopes.resolve_local_name(name_ref).unwrap();
-        let local_name = local_name_entry.ptr();
+        let local_name =
+            local_name_entry.ptr().either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr());
         assert_eq!(local_name.range(), expected_name.syntax().range());
     }
 
diff --git a/crates/ra_hir/src/lib.rs b/crates/ra_hir/src/lib.rs
index 4d337d2e3..0881939a2 100644
--- a/crates/ra_hir/src/lib.rs
+++ b/crates/ra_hir/src/lib.rs
@@ -17,6 +17,8 @@ macro_rules! impl_froms {
     }
 }
 
+mod either;
+
 pub mod db;
 #[macro_use]
 pub mod mock;
@@ -52,11 +54,12 @@ use crate::{
 };
 
 pub use self::{
+    either::Either,
     path::{Path, PathKind},
     name::Name,
     source_id::{AstIdMap, ErasedFileAstId},
     ids::{HirFileId, MacroDefId, MacroCallId, MacroCallLoc},
-    nameres::{PerNs, Namespace, ImportId, ImportSource},
+    nameres::{PerNs, Namespace, ImportId},
     ty::{Ty, ApplicationTy, TypeCtor, Substs, display::HirDisplay},
     impl_block::{ImplBlock, ImplItem},
     docs::{Docs, Documentation},
diff --git a/crates/ra_hir/src/nameres.rs b/crates/ra_hir/src/nameres.rs
index 4ae04514a..0eddfab12 100644
--- a/crates/ra_hir/src/nameres.rs
+++ b/crates/ra_hir/src/nameres.rs
@@ -75,7 +75,7 @@ pub(crate) use self::raw::{RawItems, ImportSourceMap};
 
 pub use self::{
     per_ns::{PerNs, Namespace},
-    raw::{ImportId, ImportSource},
+    raw::ImportId,
 };
 
 /// Contans all top-level defs from a macro-expanded crate
diff --git a/crates/ra_hir/src/nameres/raw.rs b/crates/ra_hir/src/nameres/raw.rs
index b7416ede6..43c97a0bf 100644
--- a/crates/ra_hir/src/nameres/raw.rs
+++ b/crates/ra_hir/src/nameres/raw.rs
@@ -12,7 +12,7 @@ use ra_syntax::{
 
 use crate::{
     DefDatabase, Name, AsName, Path, HirFileId, ModuleSource,
-    AstIdMap, FileAstId,
+    AstIdMap, FileAstId, Either,
 };
 
 /// `RawItems` is a set of top-level items in a file (except for impls).
@@ -34,28 +34,15 @@ pub struct ImportSourceMap {
     map: ArenaMap<ImportId, ImportSourcePtr>,
 }
 
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
-enum ImportSourcePtr {
-    UseTree(AstPtr<ast::UseTree>),
-    ExternCrate(AstPtr<ast::ExternCrateItem>),
-}
+type ImportSourcePtr = Either<AstPtr<ast::UseTree>, AstPtr<ast::ExternCrateItem>>;
+type ImportSource = Either<TreeArc<ast::UseTree>, TreeArc<ast::ExternCrateItem>>;
 
 impl ImportSourcePtr {
     fn to_node(self, file: &SourceFile) -> ImportSource {
-        match self {
-            ImportSourcePtr::UseTree(ptr) => ImportSource::UseTree(ptr.to_node(file).to_owned()),
-            ImportSourcePtr::ExternCrate(ptr) => {
-                ImportSource::ExternCrate(ptr.to_node(file).to_owned())
-            }
-        }
+        self.map(|ptr| ptr.to_node(file).to_owned(), |ptr| ptr.to_node(file).to_owned())
     }
 }
 
-pub enum ImportSource {
-    UseTree(TreeArc<ast::UseTree>),
-    ExternCrate(TreeArc<ast::ExternCrateItem>),
-}
-
 impl ImportSourceMap {
     fn insert(&mut self, import: ImportId, ptr: ImportSourcePtr) {
         self.map.insert(import, ptr)
@@ -281,11 +268,7 @@ impl RawItemsCollector {
         Path::expand_use_item(use_item, |path, use_tree, is_glob, alias| {
             let import_data =
                 ImportData { path, alias, is_glob, is_prelude, is_extern_crate: false };
-            self.push_import(
-                current_module,
-                import_data,
-                ImportSourcePtr::UseTree(AstPtr::new(use_tree)),
-            );
+            self.push_import(current_module, import_data, Either::A(AstPtr::new(use_tree)));
         })
     }
 
@@ -304,11 +287,7 @@ impl RawItemsCollector {
                 is_prelude: false,
                 is_extern_crate: true,
             };
-            self.push_import(
-                current_module,
-                import_data,
-                ImportSourcePtr::ExternCrate(AstPtr::new(extern_crate)),
-            );
+            self.push_import(current_module, import_data, Either::B(AstPtr::new(extern_crate)));
         }
     }
 
diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index f6a325033..ecc63f376 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -2324,7 +2324,7 @@ fn infer(content: &str) -> String {
 
         for (pat, ty) in inference_result.type_of_pat.iter() {
             let syntax_ptr = match body_source_map.pat_syntax(pat) {
-                Some(sp) => sp,
+                Some(sp) => sp.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr()),
                 None => continue,
             };
             types.push((syntax_ptr, ty));