type-safer source-map for bindings

7 files changed, 45 insertions, 30 deletions

diff --git a/crates/ra_hir/src/either.rs b/crates/ra_hir/src/either.rs
index 6714529d9..4073cc82e 100644
--- a/crates/ra_hir/src/either.rs
+++ b/crates/ra_hir/src/either.rs
@@ -5,6 +5,16 @@ pub enum Either<A, 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,
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/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));
diff --git a/crates/ra_ide_api/src/completion/complete_path.rs b/crates/ra_ide_api/src/completion/complete_path.rs
index e54fe7b7e..7e47fa6bd 100644
--- a/crates/ra_ide_api/src/completion/complete_path.rs
+++ b/crates/ra_ide_api/src/completion/complete_path.rs
@@ -1,4 +1,4 @@
-use hir::Resolution;
+use hir::{Resolution, Either};
 use ra_syntax::AstNode;
 use test_utils::tested_by;
 
@@ -19,10 +19,8 @@ pub(super) fn complete_path(acc: &mut Completions, ctx: &CompletionContext) {
             for (name, res) in module_scope.entries() {
                 if Some(module) == ctx.module {
                     if let Some(import) = res.import {
-                        if let hir::ImportSource::UseTree(tree) =
-                            module.import_source(ctx.db, import)
-                        {
-                            if tree.syntax().range().contains_inclusive(ctx.offset) {
+                        if let Either::A(use_tree) = module.import_source(ctx.db, import) {
+                            if use_tree.syntax().range().contains_inclusive(ctx.offset) {
                                 // for `use self::foo<|>`, don't suggest `foo` as a completion
                                 tested_by!(dont_complete_current_use);
                                 continue;
diff --git a/crates/ra_ide_api/src/goto_definition.rs b/crates/ra_ide_api/src/goto_definition.rs
index 660b43cfa..60c1f5085 100644
--- a/crates/ra_ide_api/src/goto_definition.rs
+++ b/crates/ra_ide_api/src/goto_definition.rs
@@ -113,6 +113,7 @@ pub(crate) fn reference_definition(
                 let ptr = source_map.pat_syntax(pat).expect("pattern not found in syntax mapping");
                 let name =
                     path.as_ident().cloned().expect("local binding from a multi-segment path");
+                let ptr = ptr.either(|it| it.into(), |it| it.into());
                 let nav = NavigationTarget::from_scope_entry(file_id, name, ptr);
                 return Exact(nav);
             }
diff --git a/crates/ra_ide_api/src/references.rs b/crates/ra_ide_api/src/references.rs
index 20bbf11a3..3e30e047c 100644
--- a/crates/ra_ide_api/src/references.rs
+++ b/crates/ra_ide_api/src/references.rs
@@ -1,5 +1,5 @@
 use relative_path::{RelativePath, RelativePathBuf};
-use hir::{ModuleSource, source_binder};
+use hir::{ModuleSource, source_binder, Either};
 use ra_db::{SourceDatabase};
 use ra_syntax::{
     AstNode, SyntaxNode, SourceFile,
@@ -89,9 +89,12 @@ pub(crate) fn find_all_refs(
             source_binder::function_from_child_node(db, position.file_id, name_ref.syntax())?;
         let scope = descr.scopes(db);
         let resolved = scope.resolve_local_name(name_ref)?;
-        let resolved = resolved.ptr().to_node(source_file);
-        let binding = find_node_at_offset::<ast::BindPat>(syntax, resolved.range().end())?;
-        Some((binding, descr))
+        if let Either::A(ptr) = resolved.ptr() {
+            if let ast::PatKind::BindPat(binding) = ptr.to_node(source_file).kind() {
+                return Some((binding, descr));
+            }
+        }
+        None
     }
 }