type-safer source-map for bindings

4 files changed, 34 insertions, 21 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));